You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
182 lines
6.0 KiB
Python
182 lines
6.0 KiB
Python
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_det_dict = [perdict for perdict in det_dict if list(perdict.keys())[0] == 'person']
|
|
|
|
# print('person_det_dict:',det_dict)
|
|
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):
|
|
|
|
# print('11111')
|
|
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'])
|
|
|
|
# print('hands_list:',hands_list)
|
|
imgsize2 = person_cut_frame.shape
|
|
|
|
# 手部坐标修正
|
|
hands_bbox_list = tools_function.para_list_correction(images_size=imgsize2,bbox_list=hands_list,dertpara=5)
|
|
|
|
# print('hands_bbox_list:',hands_bbox_list)
|
|
|
|
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)
|
|
|
|
# print('hands_result_list:',hands_result_list)
|
|
|
|
re_dict = {'hands':hands_result_list}
|
|
|
|
hand_bbox_list.append(re_dict)
|
|
|
|
# hands_result_original_dict = {'results':hand_bbox_list}
|
|
|
|
# print(hands_result_original_dict)
|
|
|
|
return hand_bbox_list
|
|
|
|
def save_annotations_xml(self,xml_save_file,save_infors,images):
|
|
|
|
# images = 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):
|
|
|
|
|
|
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()]
|
|
|
|
# print('per_frame_list:',per_frame_list)
|
|
|
|
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):
|
|
|
|
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)
|
|
|
|
|
|
# person_det = self.detect_yolo(images_path=images,labels_name_list=labels_name_list)
|
|
|
|
def det_process(self,images,xml_save_file):
|
|
|
|
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)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|