|
|
import cv2
|
|
|
import os
|
|
|
import mediapipe as mp
|
|
|
from ultralytics import YOLO
|
|
|
from personDet import analysis_yolov8
|
|
|
import tools_function
|
|
|
from holisticDet import MediapipeProcess
|
|
|
from add_xml import add_xml
|
|
|
from create_xml import create_xml
|
|
|
import queue
|
|
|
|
|
|
|
|
|
class DetProcess():
|
|
|
|
|
|
def __init__(self, person_det_model, hand_det_model):
|
|
|
|
|
|
self.person_det_model = person_det_model
|
|
|
self.hand_det_model = hand_det_model
|
|
|
|
|
|
def get_person_cut(self, frame, det_dict, imgsize):
|
|
|
|
|
|
person_list = tools_function.get_dict_values(det_dict)
|
|
|
|
|
|
# 坐标参数修正
|
|
|
person_bbox_list = tools_function.para_list_correction(
|
|
|
images_size=imgsize, bbox_list=person_list, dertpara=5)
|
|
|
|
|
|
frame_list = []
|
|
|
for per_bbox in person_bbox_list:
|
|
|
|
|
|
# 裁剪后人的图片
|
|
|
person_cut_frame = tools_function.img_cut(
|
|
|
frame=frame, bbox=per_bbox)
|
|
|
|
|
|
frame_cut_dict = {tuple(per_bbox): person_cut_frame}
|
|
|
|
|
|
frame_list.append(frame_cut_dict)
|
|
|
|
|
|
return frame_list
|
|
|
|
|
|
def hand_det(self, person_cut_frame, per_bbox):
|
|
|
|
|
|
hands_result = MediapipeProcess.mediapipe_det(
|
|
|
image=person_cut_frame, holistic=self.hand_det_model)
|
|
|
hands_result_dict = MediapipeProcess.get_analysis_result(
|
|
|
image=person_cut_frame, results=hands_result)
|
|
|
hands_list = tools_function.select_list(hands_result_dict['hand_bbox'])
|
|
|
|
|
|
imgsize2 = person_cut_frame.shape
|
|
|
|
|
|
# 手部坐标修正
|
|
|
hands_bbox_list = tools_function.para_list_correction(
|
|
|
images_size=imgsize2, bbox_list=hands_list, dertpara=5)
|
|
|
|
|
|
hand_bbox_list = []
|
|
|
for hand in hands_bbox_list:
|
|
|
|
|
|
hands_result_list = tools_function.change_bbox(
|
|
|
bbox_person=[per_bbox[0], per_bbox[1]], bbox_hand=hand)
|
|
|
|
|
|
re_dict = {'hands': hands_result_list}
|
|
|
|
|
|
hand_bbox_list.append(re_dict)
|
|
|
|
|
|
return hand_bbox_list
|
|
|
|
|
|
def save_annotations_xml(self, xml_save_file, save_infors, images):
|
|
|
|
|
|
results = save_infors
|
|
|
img = os.path.basename(images)
|
|
|
img_frame = cv2.imread(images)
|
|
|
xml_save_path = os.path.join(xml_save_file, img.split('.')[0] + '.xml')
|
|
|
w, h, d = img_frame.shape
|
|
|
img_shape = (w, h, d, img)
|
|
|
|
|
|
if os.path.isfile(xml_save_path):
|
|
|
|
|
|
add_labels = add_xml(inforsDict=results, xmlFilePath=xml_save_path)
|
|
|
else:
|
|
|
create_new = create_xml(
|
|
|
boxs=results, img_shape=img_shape, xml_path=xml_save_path)
|
|
|
|
|
|
def person_cut_process(self, images, img_save_files):
|
|
|
'''
|
|
|
使用模型将检测到的人从大图中裁剪出来
|
|
|
images：图片全路径
|
|
|
img_save_files：裁剪后保存小图的文件夹
|
|
|
通过设置labels_name_list列表中的标签名，调整需要裁剪出来的目标
|
|
|
'''
|
|
|
|
|
|
frame = cv2.imread(images)
|
|
|
imgsize = frame.shape
|
|
|
labels_name_list = ['person']
|
|
|
|
|
|
per_det_dict = analysis_yolov8(frame=frame,
|
|
|
model_coco=self.person_det_model,
|
|
|
labels_names=labels_name_list,
|
|
|
confidence_set=0.2)
|
|
|
|
|
|
per_frame_cut = self.get_person_cut(
|
|
|
frame=frame, det_dict=per_det_dict, imgsize=imgsize)
|
|
|
per_frame_list = [
|
|
|
value for dictionary in per_frame_cut for value in dictionary.values()]
|
|
|
|
|
|
for id_num, cut_frame in enumerate(per_frame_list):
|
|
|
|
|
|
cut_frame_save = tools_function.img_write(
|
|
|
frame=cut_frame, img_file=images, id_num=id_num, save_file=img_save_files)
|
|
|
|
|
|
def hands_det_process(self, images, xml_save_file):
|
|
|
'''
|
|
|
使用目标检测模型检测到行人，然后串联使用mediapipe模型检测到人的手部，后将检测到的手部坐标框保存成xml标注文件
|
|
|
images：检测图片的全路径
|
|
|
xml_save_file: 保存标注文件路径
|
|
|
'''
|
|
|
|
|
|
frame = cv2.imread(images)
|
|
|
imgsize = frame.shape
|
|
|
labels_name_list = ['person']
|
|
|
|
|
|
per_det_dict = analysis_yolov8(frame=frame,
|
|
|
model_coco=self.person_det_model,
|
|
|
labels_names=labels_name_list,
|
|
|
confidence_set=0.2)
|
|
|
|
|
|
per_frame_cut = self.get_person_cut(
|
|
|
frame=frame, det_dict=per_det_dict, imgsize=imgsize)
|
|
|
|
|
|
for frame_dict in per_frame_cut:
|
|
|
|
|
|
per_bbox = list(frame_dict.keys())[0]
|
|
|
person_cut_frame = list(frame_dict.values())[0]
|
|
|
|
|
|
hands_dict = self.hand_det(
|
|
|
person_cut_frame=person_cut_frame, per_bbox=per_bbox)
|
|
|
|
|
|
self.save_annotations_xml(
|
|
|
xml_save_file=xml_save_file, save_infors=hands_dict, images=images)
|
|
|
|
|
|
def det_process(self, images, xml_save_file):
|
|
|
'''
|
|
|
检测指定目标，并保存检测结果到xml文件中
|
|
|
通过设置labels_name_list列表中的标签名，调整需要裁剪出来的目标
|
|
|
'''
|
|
|
|
|
|
frame = cv2.imread(images)
|
|
|
imgsize = frame.shape
|
|
|
labels_name_list = ['cell phone', 'mouse', 'keyboard']
|
|
|
|
|
|
per_det_dict = analysis_yolov8(frame=frame,
|
|
|
model_coco=self.person_det_model,
|
|
|
labels_names=labels_name_list,
|
|
|
confidence_set=0.2)
|
|
|
|
|
|
self.save_annotations_xml(
|
|
|
xml_save_file=xml_save_file, save_infors=per_det_dict, images=images)
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
|
|
images_files = 'images'
|
|
|
images_list = tools_function.get_path_list(images_files)
|
|
|
img_save_files = 'images_cut'
|
|
|
xml_save_file = 'annotations'
|
|
|
|
|
|
# 初始化目标检测
|
|
|
person_model = YOLO("model_files/yolov8x.pt")
|
|
|
|
|
|
# 初始化mediapipe
|
|
|
mp_holistic = mp.solutions.holistic
|
|
|
holistic = mp_holistic.Holistic(
|
|
|
min_detection_confidence=0.1, min_tracking_confidence=0.1)
|
|
|
|
|
|
Det = DetProcess(person_det_model=person_model, hand_det_model=holistic)
|
|
|
|
|
|
for images in images_list:
|
|
|
|
|
|
Det.person_cut_process(images=images, img_save_files=img_save_files)
|
|
|
Det.hands_det_process(images=images, xml_save_file=xml_save_file)
|
|
|
Det.det_process(images=images, xml_save_file=xml_save_file)
|