红叉yolo自动标注

import numpy as np
import cv2
import sys
import glob
import os
import shutil
from os import listdir, getcwd

sets=[‘5’,‘6’]

def convert(size, box):
dw = 1./size[0]
dh = 1./size[1]
x = (box[0] + box[1])/2.0
y = (box[2] + box[3])/2.0
w = box[1] - box[0]
h = box[3] - box[2]
x = xdw
w = wdw
y = ydh
h = hdh
return (x,y,w,h)

def find_X(img):
x, y = img.shape
point_x = []
point_y = []
for i in range(200,(x-200),4):
if(np.sum(img[(i-195):(i+195),:]) == 0):
continue
for j in range(200,(y-200),4):
if (np.sum(img[i,(j-195):(j+195)]) == 0):
continue
top = np.sum(img[(i-195):(i-95),j])/255.0
bottom = np.sum(img[(i+95):(i+195),j])/255.0
left = np.sum(img[i,(j-195):(j-95)])/255.0
right = np.sum(img[i,(j+95):(j+195)])/255.0

        th = top + bottom + left + right
        if(th > 330):
            point_x.append(i)
            point_y.append(j)

for i in range(1,200,4):
    if(np.sum(img[i:(i+145),:]) == 0):
        continue
    for j in range(200,(y-200),4):
        if (np.sum(img[i,(j-195):(j+195)]) == 0):
            continue
        bottom = np.sum(img[(i+95):(i+195),j])/255.0
        left = np.sum(img[i,(j-195):(j-95)])/255.0
        right = np.sum(img[i,(j+95):(j+195)])/255.0
        th =  bottom + left + right
        if(th > 230):
            point_x.append(i)
            point_y.append(j)


for i in range((x-200),x,4):
    if(np.sum(img[(i-195):i,:]) == 0):
        continue
    for j in range(200,(y-200),4):
        if (np.sum(img[i,(j-195):(j+195)]) == 0):
            continue
        top = np.sum(img[(i-195):(i-95),j])/255.0

        left = np.sum(img[i,(j-195):(j-95)])/255.0
        right = np.sum(img[i,(j+95):(j+195)])/255.0
        th = top + left + right
        if(th > 230):
            point_x.append(i)
            point_y.append(j)

return point_x,point_y

def find_NG(i,image_set):
out_file = open(r’./data/’ + image_set + ‘/’ + i.split(’/’)[2].split(’\’)[1].split(’.’)[0] + ‘.txt’, “w”)

img_c = cv2.imread(i)
img0 = cv2.imread(i,0)
size = (3941,1935) 
img0 = cv2.resize(img0, size, interpolation=cv2.INTER_AREA)
img_c = cv2.resize(img_c, size, interpolation=cv2.INTER_AREA)

hsv = cv2.cvtColor(img_c,cv2.COLOR_BGR2HSV)
H, S, V = cv2.split(hsv)
S1 = S.copy()
S1[S1<100] = 0
S1[S1>0] = 255
S2 = S.copy()
S2[S2<100] = 0
S2[S2>0] = 255    


X,Y = find_X(S2)

dst = cv2.inpaint(img0,S1,1,cv2.INPAINT_TELEA)

if (len(X) > 0):
    cv2.rectangle(dst, (Y[0] - 145, X[0] - 145), (Y[0] + 145, X[0] + 145), (0, 255, 0), 2)
    b=(float(Y[0] - 145),float(Y[0] + 145),float(X[0] - 145),float(X[0] + 145))
    bb=convert((size),b)
    out_file.write(str(image_set) + " " + " ".join([str(a) for a in bb]) + '\n')

    # cv2.namedWindow("i", cv2.WINDOW_NORMAL)
    # cv2.imshow('i', dst)
    # cv2.imwrite('./i.jpg',dst)
    # cv2.waitKey(0)

return dst

for image_set in sets:

filelist = sorted(glob.glob(r'./data/' + image_set + '/*.jpg'))
num = 0
for i in filelist:

    num = num + 1
    if(num == 1):
        th = find_NG(i,image_set)

    if(num > 1):
        th = find_NG(i,image_set)