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)