头肩检测数据清洗小工具

import os
import numpy as np
import cv2


class WiderFaceDetection():
    def __init__(self, txt_path):
        self.imgs_path = []
        self.words = []
        f = open(txt_path,'r')
        lines = f.readlines()
        isFirst = True
        labels = []
        txt_path = txt_path[0:txt_path.rindex("/")+1]
        for line in lines:
            line = line.rstrip()
            if line.startswith('#'):
                if isFirst is True:
                    isFirst = False
                else:
                    labels_copy = labels.copy()
                    self.words.append(labels_copy)
                    labels.clear()
                path = line[2:]
                path = txt_path + path
                self.imgs_path.append(path)
            else:
                line = line.split(' ')
                label = [float(x) for x in line]
                labels.append(label)

        self.words.append(labels)

face_index_select = []
labels = []
labels_add = []
new_label = []
annotations = np.zeros((0, 5))
training_dataset = "/data/ir_face_dataset/u2outdoor_face_shoulder_clear.txt"
dataset = WiderFaceDetection(training_dataset)
dataset_new_words = dataset.words.copy()
point1 = [0,0]
point2 = [0,0]

labels_merger = []


add_flag = 0
add_finish_flag = 0

def getSelectFace(x,y):
    face_count = len(labels_merger)
    for i in range(face_count):
        anno = labels_merger[i]
        if x>=anno[0] and x < (anno[0] + anno[2]) and y >=anno[1] and y<(anno[1]+anno[3]):
            return i
    return -1

def onMouse(event, x, y, flags, param):
    global point2,point1,labels_merger,add_flag,add_finish_flag
    if x<0 or y <0:
        add_flag =0
        return
    face_select = getSelectFace(x,y)
    if event == cv2.EVENT_MOUSEMOVE and add_flag == 1:
        param_show = param.copy()
        # print("x y point1",x,y,point1)
        cv2.rectangle(param_show, (point1[0],point1[1]), (x, y), (0, 255, 255), thickness=2)
        cv2.imshow('img', param_show)
        return

    if event == cv2.EVENT_LBUTTONDOWN:
        print("1-EVENT_LBUTTONDOWN")
        # if add_flag == 1:
        #     add_flag = 0
        #     point2 = [x, y]
        #     w = abs(x-point1[0])
        #     h = abs(y-point1[1])
        #     if w < 10 or h < 10:
        #         print("shoulder size too small ")
        #         return
        #     label_x = min(x,point1[0])
        #     label_y = min(y,point1[1])
        #     label_add_one = [label_x,label_y,w,h,-1,-1,-1,-1,2]
        #     labels_merger.append(label_add_one)
        #     cv2.rectangle(param, (point1[0],point1[1]), (point2[0],point2[1]), (0, 255, 255), thickness=2)
        #     cv2.imshow('img', param)
        #     return 
        # else:
        point1 = [x,y]
        add_flag = 1
        return  
    if event == cv2.EVENT_LBUTTONUP:
        print("2-EVENT_LBUTTONUP")
        if add_flag == 1:
            add_flag = 0
            point2 = [x, y]
            w = abs(x-point1[0])
            h = abs(y-point1[1])
            if w < 10 or h < 10:
                print("shoulder size too small ")
                return
            label_x = min(x,point1[0])
            label_y = min(y,point1[1])
            label_add_one = [label_x,label_y,w,h,-1,-1,-1,-1,2]
            labels_merger.append(label_add_one)
            cv2.rectangle(param, (point1[0],point1[1]), (point2[0],point2[1]), (0, 255, 255), thickness=2)
            cv2.imshow('img', param)
            return 
    if event == cv2.EVENT_MBUTTONDOWN and face_select != -1:
        add_flag = 0
        if face_select not in face_index_select:
            face_index_select.append(face_select)
            print("EVENT_MBUTTONDOWN", face_select)
            x = int(labels_merger[face_select][0])
            x2 = int(labels_merger[face_select][0]+labels_merger[face_select][2])
            y = int(labels_merger[face_select][1])
            y2 = int(labels_merger[face_select][1] + labels_merger[face_select][3])
            param[y:y2,x:x2,2] = 255
            cv2.imshow("img", param)
            return 
        else:
            face_index_select.remove(face_select)
            x = int(labels_merger[face_select][0])
            x2 = int(labels_merger[face_select][0]+labels_merger[face_select][2])
            y = int(labels_merger[face_select][1])
            y2 = int(labels_merger[face_select][1] + labels_merger[face_select][3])
            x = max(0,x)
            param[y:y2,x:x2,2] = param[y:y2,x:x2,0]
            cv2.imshow("img", param)
            print("EVENT_MBUTTONDOWN", face_select)
            return 



if __name__ == '__main__':
    
    cv2.namedWindow("img",0)

    label_file = training_dataset[0:-4] + "_clear.txt"
    while 1 :
        if os.path.exists(label_file):
            label_file = label_file[0:-4] + "_clear.txt"
        else :
            break

    f = open(label_file, 'w')
    img_root = training_dataset[0:training_dataset.rindex("/")+1]
    img_list_len = len(dataset.imgs_path)
    i = -1
    # print(img_list_len)
    while i < (img_list_len-1):
        i = i + 1
        i = max(0,i)
        img_path = dataset.imgs_path[i]
       
        print(i,"/", img_list_len, img_path)

        img_raw = cv2.imread(img_path)
        if img_raw is None:
           print("img not exist ", img_path)
           continue

        height, width, _ = img_raw.shape
        labels = dataset.words[i]

        face_index_select = []
        labels_merger = []

        face_count = len(labels)
        for face_i in range(face_count):
            min_face_Slid = min(labels[face_i][2], labels[face_i][3])
            if min_face_Slid<16:
                if face_i not in face_index_select:
                    face_index_select.append(face_i)
                continue
            if labels[face_i][8] == 1:
                cv2.rectangle(img_raw,(int(labels[face_i][0]), int(labels[face_i][1])),(int(labels[face_i][0]+labels[face_i][2]), int(labels[face_i][1] + labels[face_i][3])), (0, 0, 255), 2)
            else:
                cv2.rectangle(img_raw,(int(labels[face_i][0]), int(labels[face_i][1])),(int(labels[face_i][0]+labels[face_i][2]), int(labels[face_i][1] + labels[face_i][3])), (0, 255, 255), 2)
        new_label = labels.copy()
        labels_merger = labels.copy()
        print("labels ", labels)
        add_flag = 0
        add_finish_flag =0

        cv2.imshow("img", img_raw)
        cv2.waitKey(1)
        cv2.setMouseCallback("img", onMouse, img_raw)
        
        key = cv2.waitKey(0) 
        key = int(key)
        if key == 113 or key == 27: #113 q esc 27
            print("enter q")
            break
        elif key == 97: #a back
            i = i - 2
            print("i ", i)
            print("back")
            continue
        elif key == 32 or key == 13 or add_finish_flag == 1: #32 sapce 13 enter
            print("continue") 
        
    
        face_select_num = len(face_index_select)
        if face_select_num > 0:
            new_label = [labels_merger[i] for i, e in enumerate(labels_merger) if i not in face_index_select]
        else:
            new_label = labels_merger
         
        
        print("new_label ", new_label)
        dataset_new_words[i] = new_label
    
    for i in range(img_list_len):
        label_size = len(dataset_new_words[i])
        if label_size == 0:
            continue
        img_path = dataset.imgs_path[i]
        img_path = img_path.replace(img_root,'')
        line = "# {}\n".format(img_path)
        f.writelines(line)
        f.flush()

        for f_i in range(label_size):
            l = dataset_new_words[i][f_i]
            line = "{} {} {} {} {} {} {} {} {}\n".format(l[0],l[1],l[2],l[3],l[4],l[5],
            l[6],l[7],l[8])
            f.writelines(line)
            f.flush()
        
    f.close()               

    # print(label_file_data)