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

import os
import cv2


def get_path_list(file):

    retutn_list = []


    if os.path.isdir(file):

        video_ext = [".mp4", ".avi",".MP4",'.jpg']
        for maindir, subdir, file_name_list in os.walk(file):
            for filename in file_name_list:
                apath = os.path.join(maindir, filename)
                ext = os.path.splitext(apath)[1]
                if ext in video_ext:
                    retutn_list.append(apath)
    else:
        retutn_list.append(file)
    
    return retutn_list

# 获得字典中所有values值（这个值是列表)
def get_dict_values(lst):
    """
    获取列表中所有字典的 values 值（如果值是列表）

    参数:
    lst: 包含字典的列表

    返回值:
    values: 包含所有字典的 values 值的列表（如果值是列表）
    """
    return [value for dictionary in lst for value in dictionary.values() if isinstance(value, list)]

    # 解析检测后的结果，为检测后的结果排序

def img_cut(frame,bbox):

    frame_target = frame[int(bbox[1]):int(bbox[3]), int(bbox[0]):int(bbox[2])]

    return frame_target

def img_write(frame,img_file,id_num,save_file):

    filename = os.path.basename(img_file).split('.')[0] +'_'+ f"person_{id_num}.jpg"
    person_image_path = os.path.join(save_file, filename)
    cv2.imwrite(person_image_path, frame)

    return person_image_path

def change_bbox(bbox_person,bbox_hand):

    '''
    将小图上手的坐标转换到原图大图的坐标
    '''

    bbox_xmin = bbox_hand[0] + bbox_person[0]
    bbox_xmax = bbox_hand[2] + bbox_person[0]
    bbox_ymin = bbox_hand[1] + bbox_person[1]
    bbox_ymax = bbox_hand[3] + bbox_person[1]

    bbox = [bbox_xmin,bbox_ymin,bbox_xmax,bbox_ymax]

    return bbox


def select_list(result_list):
    '''
    筛选列表中的空列表
    '''
    if result_list:
        result_only = []

        for result in result_list:

            if result == None:
                pass
            else:
                result_only.append(result)

        return result_only
    

def para_list_correction(images_size, bbox_list, dertpara):


    updata_result_list = []

    for bbox in bbox_list:

        updata_bbox = para_correction(images_size, bbox, dertpara)

        size = [(updata_bbox[2] - updata_bbox[0]),
                (updata_bbox[3] - updata_bbox[1])]

        if determine_zero(size):

            if check_bbox(size):

                updata_result_list.append(updata_bbox)

    return updata_result_list

def para_correction(images_size, bbox, dertpara):
    '''
    修正检测后标注框过小的情况,如果有修正参数则使用修正参数，如果没有就按照坐标值扩大两倍

    '''

    # if dertpara == 0:
    #     pass
    # else:
    w = (bbox[2] - bbox[0]) / int(dertpara)
    h = (bbox[3] - bbox[1]) / int(dertpara)

    bbox_extand_list_x = [bbox[0] - w, bbox[2] + w]
    bbox_extand_list_y = [bbox[1] - h, bbox[3] + h]

    bbox_list_x = contrast(
        size=images_size[1], bbox_extand_list=bbox_extand_list_x)
    bbox_list_y = contrast(
        size=images_size[0], bbox_extand_list=bbox_extand_list_y)

    bbox_list = [bbox_list_x[0], bbox_list_y[0],
                 bbox_list_x[1], bbox_list_y[1]]

    return bbox_list


def contrast(size, bbox_extand_list):
    '''
    对比数值是否在这个范围内
    '''

    # print('bbox_extand_list:',bbox_extand_list)
    # print('size:',size)
    bbox_list = []

    for x in bbox_extand_list:

        # print('size:',size)

        if 0 <= int(x) <= int(size):
            # print('in:',x,size)
            bbox_list.append(x)
        if int(x) > int(size):
            # print('>:',x,size)
            bbox_list.append(size)
        if int(x) < 0:
            # print('<:',x,size)
            bbox_list.append(0)

    # print('bbox_list:',bbox_list)

    return bbox_list


def determine_zero(num_list):

    for num in num_list:

        if num == 0:
            return False
    return True


def check_bbox(size):

    ratio = size[0]/size[1]

    if ratio >= 5:

        return False

    if ratio <= 1/5:

        return False

    else:
        return True