[Yolo5] Yolo5의 detect.py 사용해서 Yolo 학습용 데이터셋 만들기

 

 

우리집엔 봉식이뿐이지만, 고양이가 고양이를 불러서인지 다묘가정이 더 많은 것 같다.

 

그래서 다묘가정을 위한... 두,세 마리 고양이를 각각 구분하는 모델을 만들어보고자 하는데

저번처럼 또 roboflow에서 하나하나 바운딩 박스 그릴 생각하니 머리가 아파졌다...

 

그리고 봉식이 말고 다른 특정 고양이 사진을 어디서 가져와야 할 지도 고민이였는데

전에 caffemodel을 사용해서 face detection을 하고 얼굴이 검출된 프레임만 저장하는 작업을 했었던게 생각났다.

 

또 yolo의 detect.py 맨 위 주석을 보면

$ python path/to/detect.py --weights yolov5s.pt --source 0              # webcam
                                                         img.jpg        # image
                                                         vid.mp4        # video
                                                         path/          # directory
                                                         path/*.jpg     # glob
                                                         'https://youtu.be/Zgi9g1ksQHc'  # YouTube
                                                         'rtsp://example.com/media.mp4'  # RTSP, RTMP, HTTP stream

요로케... 유튜브 영상도 source로 줄 수 있다는걸 알 수 있다.

 

 

 

그래서 생각해낸 아이디어 💡 : 

1. 냥튜브 채널의 영상을 source로 가져오자

2. 기존 yolo모델을 사용하여 detect를 진행하자

3. cat이 detect된 프레임을 jpg로 저장하자

4. --save-txt 인자를 이용해 cat이 detect된 좌표가 저장된 txt파일을 만들자

 

 

뿌듯... 생각만 했는데도 뿌듯함..

어차피 봉식이도,, 다른 고양이도 ,, 고양이니까 기존 yolo 모델과 기능들을 활용하면 편하게 데이터셋을 만들수 있다

 

 

 

 

---

 

 

일단 결과를 저장하는 부분부터 살펴보았다.

 

# Save results (image with detections)
if save_img:
    if dataset.mode == 'image':
        cv2.imwrite(save_path, im0)
    else:  # 'video' or 'stream'
        if vid_path[i] != save_path:  # new video
            vid_path[i] = save_path
            if isinstance(vid_writer[i], cv2.VideoWriter):
                vid_writer[i].release()  # release previous video writer
            if vid_cap:  # video
                fps = vid_cap.get(cv2.CAP_PROP_FPS)
                w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            else:  # stream
                fps, w, h = 30, im0.shape[1], im0.shape[0]
            save_path = str(Path(save_path).with_suffix('.mp4'))  # force *.mp4 suffix on results videos
            vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
        vid_writer[i].write(im0)

 

185라인에 있는 if save_img 구문이다. 필요한 부분만 건드려보자.

 

save_path 부분을 보면 suffix를 mp4로 주고있는데 이 부분을 지운다.

vid_writer부분도 동영상 저장과 관련된 부분이기 때문에 지웠다.

 

cv2.imwrite를 사용하여 이미지를 저장하는 코드를 넣는다.

 

 

결과 : 

# Save results (image with detections)
if save_img:
    if dataset.mode == 'image':
        cv2.imwrite(save_path, im0)
    else:  # 'video' or 'stream'
        if vid_path[i] != save_path:  # new video
            vid_path[i] = save_path
            if isinstance(vid_writer[i], cv2.VideoWriter):
                vid_writer[i].release()  # release previous video writer
            if vid_cap:  # video
                fps = vid_cap.get(cv2.CAP_PROP_FPS)
                w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            else:  # stream
                fps, w, h = 30, im0.shape[1], im0.shape[0]
            #save_path = str(Path(save_path).with_suffix('.mp4'))  # force *.mp4 suffix on results videos
            #vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
        #vid_writer[i].write(im0)
        cv2.imwrite(save_path, im0)

 

 

 

 

이제 save_path를 어떻게 설정할지 힌트를 얻기 위해 txt_path를 설정해주는 부분에 가보자

 

save_path = str(save_dir / p.name)  # im.jpg
txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # im.txt

148라인이다.

기본적인 save_path 값은 이미지를 source로 한 작업의 결과물을 jpg파일로 저장하기 위한 값이다.

이걸 txt_path랑 비스무리하게 바꿔주면 된다

txt_path는 frame값을 파일명 제일 뒤에 붙여줘서 이름이 계속 바뀌게 되어있다.

 

 

결과:

#save_path = str(save_dir / p.name)  # im.jpg
save_path = str(save_dir / 'images' / p.stem) + f'_{frame}' + '.jpg'
txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # im.txt

필요없는 if else부분을 빼고 디렉토리명을 images로 바꿔주었다.

 

 

당연히 images 디렉토리를 만드는 코드도 넣어주어야한다. txt_path의 labels 디렉토리를 만드는 부분을 참고하자.

 

# Directories
save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir

87라인 (save_dir은 detect 결과가 저장되는 디렉토리를 만들때 자동으로 뒤에 붙는 숫자가 증가하도록 되어있는 부분... detect폴더 안에 exp, exp1, exp2... 이런식으로 생성되는... 지금은 관계없다.)

 

labels는 save_txt 인자 값에 따라 디렉토리를 만들지 말지가 정해지는데 난 무조건 만들어줘야해서 if부분을 빼버렸다.

 

결과:

(save_dir / 'labels').mkdir(parents=True, exist_ok=True)
(save_dir / 'images').mkdir(parents=True, exist_ok=True)

images 디렉토리를 만들어주는 부분까지 완료했다.

 

 

 

이제 'if 고양이 detect: 사진저장' 하는 부분을 만들어줘야한다. write result부분으로 가줬다.

# Write results
for *xyxy, conf, cls in reversed(det):
    if save_txt:  # Write to file
        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
        line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
        with open(txt_path + '.txt', 'a') as f:
            f.write(('%g ' * len(line)).rstrip() % line + '\n')

167라인

txt파일을 write하는 부분이다. save-txt인자를 넣어주면 txt파일을 생성하는데,

이걸 고양이가 detect 됐을 때~를 나타내는 if문으로 바꿔준다.

 

또 하나.. detect된 객체의 class도 같이 write해주고 있는데 이걸 1로 고정해주었다.

(다음에 학습시킬 때, 0이 봉식이고 1을 새로운 고양이로 하기 위해서..)

 

 

 

결과:

for *xyxy, conf, cls in reversed(det):
    #if save_txt:  # Write to file
    if cls == 15:
        save_cat = True
        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
        line = (1, *xywh)
        #line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
        with open(txt_path + '.txt', 'a') as f:
            f.write(('%g ' * len(line)).rstrip() % line + '\n')

cls값은 15가 고양이다. 

coco128.yaml파일의 names 부분을 보면 알 수 있다.

(따로 --data 인자를 넣지 않는다면 기본이 coco128.yaml이다)

 

coco128.yaml의 Classes부분

 

 

고양이가 detect 됐는지 아닌지를 판별하기 위한 save_cat 변수를 넣어주었다.

save_cat = False

선언은 txt_path 밑줄에 해주어서 For문을 돌면서 기본값은 False이고 cls가 15일 때만 True값을 가지게 해주었다. 

 

그리고 아까 수정했던 if save_img 를 if save_cat으로 수정해주었다.

if save_cat:

 

이제 고양이가 detect된 프레임 이미지를 저장하고 바운딩 박스의 좌표를 저장해준다!!

 

 

근데 이제 yolo는 그 바운딩박스를 이미지에 그려주고... class name과 confidence값도 출력해준다

 

이렇게요...

 

 

 

이걸 지워주자

 

            annotator.box_label(xyxy, label, color=colors(c, True))
            if save_crop:
                save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)

# Stream results
im0 = annotator.result()
if view_img:

174라인 anotator 부분이다.

annotator.. 박스 그리고 label 표시하고...

im0에 annotate 완료한 이미지를 넣어주는데 필요없으니까 없애준다.

 

 

 

결과:

            #annotator.box_label(xyxy, label, color=colors(c, True))
            if save_crop:
                save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)

# Stream results
#im0 = annotator.result()
if view_img:

 

 

 

 

휴 이제 완성이다... 테스트해보자~

 

영상은 hahaha님의 무 영상을 썼다.

https://www.youtube.com/watch?v=DHhaFAkyWMM 

 

tmi : 무는 봉식이처럼 복막염 진단을 받았고 최근 완치됐기 때문에... 정감이 간다. 흑흑 세상 모든 냥이들아 아프지마...

 

 

python detect.py --weights yolov5s.pt --img 640 --conf 0.5 --source https://www.youtube.com/watch?v=DHhaFAkyWMM

> 이거는 코드를 계속 수정하고 테스트하느라 로컬에서 진행했다. < 휴...코랩은 런타임이 자꾸 끊겨서...

 

 

 

 

 

 

 

흑흑 성공

이제 이걸로 모델 학습 시켜봐야겠다.