[딥러닝][yolov5]custom dataset tutorial

비디오를 png파일로 포맷하여 images에 사진들을 넣고

json을 읽어서 같은 프레임에 대한 정보를 labels에 넣는다.

 

%주의사항

(1)

cap.read()를 하는순간 

frame 은 첫번째순서로 넘어간다.

따라서 cap.get(1) = 1 이 되기때문에 

json[0]부터읽어야하기때문에 frame을 한프레임 조정해야하는 점

 

(2)

opencv가 (x,y,w,z)로 주어진 데이터의 정확한 좌표가 어디일까?

첫번째가 cv2.rectangle에서의 (x,y,w,z)

object의 왼쪽모서리이다.

 

두번째가 yolov5에서의 bounding box (x,y,w,z)

yolov5에서의 xy는 object의 중심이다.

 

 

따라서 cv2.rectangle에서의 (x,y,w,z)는 yolo format으로 변환하기위해서

x_center = (2*x + w) / width

y_center = (2*y + h) / heigh

 

# coding=utf-8
import os
import cv2
import json
#import matplotlib.pyplot as plt


width = 1920
height = 1080

def zero(idx):
    if idx ==1:
        return 0
    else:
        None

def convert(idx):
    if idx ==[]:
        return None
    dw = 1./width
    dh = 1./height
    
    x = (2*idx[0]+idx[2])/2
    y = (2*idx[1]+idx[3])/2

    x = x * dw
    w = idx[2] * dw
    y = y * dh
    h = idx[3] * dh
    return (x,y,w,h)


def video2png(video_path,json_path):
    cap = cv2.VideoCapture(video_path)
    frame_count = 0
    success = True
    print('read video in', video_path)
    videoclass = os.path.basename((os.path.dirname(video_path)))
    #videoclass=os.path.basename(video_path).split('.')[0]
    save_png_path = 'path/to/images/data'
    save_json = '/path/to/labels/data'
    with open(json_path,"r") as json_file:
        json_read = json.load(json_file)
    #json
    jsonclass = os.path.basename(os.path.dirname(video_path))  
    while success:
        success, frame = cap.read()
        if not success:
            break
        if (int(cap.get(1)-1)%40 == 0):
            #jsoncut
            os.chdir(save_json)

            f = open(jsonclass+"I{}.txt".format(str(frame_count).zfill(4)),'w')
            normalize = convert(json_read['gt_rect'][int(cap.get(1)-1)])
            if zero(json_read['exist'][int(cap.get(1)-1)] is None):
                f.write("\t")
            else:
                f.write("{}\t".format(zero(json_read['exist'][int(cap.get(1)-1)])))
            if normalize is None:
                f.write("\n")
            else:
                f.write("{}\t{}\t{}\t{}\n".format(normalize[0],normalize[1],normalize[2],normalize[3]))
            f.close()
            #framecut
            
            params = [cv2.IMWRITE_PNG_COMPRESSION, 3]
            print("Frame %d" %frame_count)
            cv2.imwrite(save_png_path + "/{}I{}.png".format(videoclass,str(frame_count).zfill(4)), frame, params)
            frame_count = frame_count + 1
           
            # frame = cv2.rectangle(frame,json_read['gt_rect'][int(cap.get(1)-1)],(0,0,255),3)
            # plt.imshow(frame)

    cap.release()


def get_all_dir(root_dir):
    video_pathes = []
    for dir in os.listdir(root_dir):
        dir = os.path.join(root_dir, dir)
        if os.path.isdir(dir):
            video_pathes.extend(get_all_dir(dir))
        if dir[-5:] == 'e.mp4':
            video_pathes.append(dir)

    return video_pathes

def get_all_dir_json(root_dir):
    json_pathes = []
    for dir in os.listdir(root_dir):
        dir = os.path.join(root_dir, dir)
        if os.path.isdir(dir):
            json_pathes.extend(get_all_dir_json(dir))
        if dir[-6:] == 'e.json':
            json_pathes.append(dir)
    return json_pathes


root_dir = '/path/to/directory'

json_pathes = get_all_dir_json(root_dir)
video_pathes = get_all_dir(root_dir)

zips= zip(video_pathes,json_pathes)

for video_path,json_path in zips:
    video2png(video_path,json_path)

 

마지막으로...

코드가 파이썬으로 짠 것 같지 않다는 선임님의 말을 들었다.

 

파이썬같아 보이기 위해서는 

class도 써주고

if __name__ == '__main__' 도 써줘야 깔끔해보이는데

너무 def만 남발했다..

 

개발자로 사는기간이 길어지면 깔끔하게 코드를 작성하기를 기대하고 

그때의 모습을 매우 기다리고 있다. ㅎㅎ