XZNSH-Code-AI/Bank_second_part/model_annotation/model_annotations_process.py

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)