create_dataset.py

import cv2
from collections import OrderedDict
import os
from import_data import get_list

# 0: 정지, 1: 어린이 보호구역, 2: 버스전용, 3: 직진, 4: 좌회전, 5: 직좌, 6: 빨간불

root_dir = 'data/'
image_dir = os.path.join(root_dir, 'red_light')
save_dir = os.path.join(root_dir, 'TL_new_text')
down = 0

label = 0
objects = ["Stop", "Children", "Bus", "TL_Go", "TL_Left", "TL_GoLeft", "TL_Stop"]

BoundingBox = [[] for row in range(len(objects))]
class_names_to_id={'Stop':0, 'Children':1, 'Bus':2, 'TL_Go':3, 'TL_Left':4, 'TL_GoLeft':5, 'TL_Stop':6}

def mouse_callback(event, x, y, flags, params):
    global x_min, x_max, y_min, y_max, BoundingBox, down, img_result, img_copy

    if event == cv2.EVENT_LBUTTONDOWN:
        x_min = x
        y_min = y
        down = 1

    if event == cv2.EVENT_MOUSEMOVE:
        img_copy = img_color.copy()
        cv2.line(img_copy, (0, y), (width, y), (122, 80, 0))
        cv2.line(img_copy, (x, 0), (x, height), (122, 80, 0))
        if down == 1:
            cv2.rectangle(img_copy, (x_min, y_min), (x, y), (10, 10, 122), -1)

        img_result = cv2.addWeighted(img_color, 0.6, img_copy, 0.4, 0)

    if event == cv2.EVENT_LBUTTONUP:
        x_max = x
        y_max = y
        down = 0
        BoundingBox[label].append([x_min, y_min, x_max, y_max])



cv2.namedWindow('img_color')
cv2.setMouseCallback('img_color', mouse_callback)

images_ = get_list(os.listdir(image_dir), ".jpg")
images_.sort()
images = []
real_name = []
for idx, _ in enumerate(images_):
    images.append(images_[idx].split('/')[-1])
    real_name.append(images_[idx].split('\\')[-1].split('.')[0])

print(real_name)

stay = 1

for idx, data in enumerate(images):
    if stay == 0:
        label = 0

    file_data = OrderedDict()
    directory=image_dir + '/' + data.split('\\')[-1]
    img_color = cv2.imread(directory)
    height, width, ch = img_color.shape
    img_copy = img_color.copy()
    img_result = img_color.copy()

    while(True):
        cv2.imshow('img_color', img_result)
        waitkey = cv2.waitKey(2)
        cv2.putText(img_result, "label: {}".format(objects[label]), (0, 15), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255))

        if waitkey & 0xFF == 97:
            with open(os.path.join(save_dir, '{}.txt'.format(real_name[idx])), 'w', encoding = "utf-8") as make_file:
                for i in range(len(objects)):
                        for a in BoundingBox[i]:
                            b_center_x =(a[0]+a[2])/2/width
                            b_center_y =(a[1]+a[3])/2/height
                            b_width =abs((a[2]-a[0])/width)
                            b_height =abs((a[3]-a[1])/height)
                            line="{} {:.3f} {:.3f} {:.3f} {:.3f}\n".format(class_names_to_id[objects[i]], b_center_x, b_center_y, b_width, b_height)
                            make_file.writelines(line)
            BoundingBox = [[] for row in range(len(objects))]
            break

        if waitkey & 0xFF == 100:
            BoundingBox = [[] for row in range(len(objects))]
            print("Deleted Bounding Box!")
        if waitkey & 0xFF == 0x31: # class 1
            img_result = cv2.addWeighted(img_color, 0.8, img_copy, 0.2, 0)
            label = 0
        if waitkey & 0xFF == 0x32: # class 1
            img_result = cv2.addWeighted(img_color, 0.8, img_copy, 0.2, 0)
            label = 1
        if waitkey & 0xFF == 0x33: # class 1
            img_result = cv2.addWeighted(img_color, 0.8, img_copy, 0.2, 0)
            label = 2
        if waitkey & 0xFF == 0x34: # class 1
            img_result = cv2.addWeighted(img_color, 0.8, img_copy, 0.2, 0)
            label = 3
        if waitkey & 0xFF == 0x35: # class 1
            img_result = cv2.addWeighted(img_color, 0.8, img_copy, 0.2, 0)
            label = 4
        if waitkey & 0xFF == 0x36:
            img_result = cv2.addWeighted(img_color, 0.8, img_copy, 0.2, 0)
            label = 5
        if waitkey & 0xFF == 0x37:
            img_result = cv2.addWeighted(img_color, 0.8, img_copy, 0.2, 0)
            label = 6
    print(data, "done // ", idx+1, "/", len(images), "done")