Week-4 YOLOv7 사용하기 (detection 첫걸음)

축제, 시험기간 사이에 끼여 classification을 더 진행하자니 재미없을 것 같고, mmdetection을 하자니 너무 무거운 주제가 될 것 같아서 yolo 사용법을 익히며 살짝 재미를 볼까 한다.

colab이 아닌 개인GPU를 사용한다면 실시한 detection도 체험할 수 있는데 그건 나중에 사용법만 남겨두조가 한다.

 

사용하기에 앞서, wandb에 회원가입을 해야한다.

https://kr.wandb.ai/

Weights & Biases – Developer tools for ML

가입하고 오자 !

 

wandb는 훈련에 사용되는 acc, loss 를 추적하게 그래프를 그려주는 라이브러리고, 서버에 저장까지해주니, 훈련 돌려놓고 잊고 있어도 보기좋게 정보를 정리해준다. 게다가 팀 단위로 자료를 저장할 수도 있다.

사용한다고 알아서 차트를 그려주는 건 아니고 log에 추적할 값을 하나하나 입력해줘야하지만, yolo에서는 wandb와 연동돼있어서 알아서 그려준다.

yolov7_train.ipynb

0.50MB

colab에서 써야하니 파일을 첨부하지만, 되도록 git pull을 써보자 !

 

위에서부터 하나씩 실행하면 쭉 실행된다.

 

!git clone https://github.com/WongKinYiu/yolov7.git

명령어로 yolov7 깃허브를 clone한다

https://github.com/WongKinYiu/yolov7?tab=readme-ov-file

GitHub - WongKinYiu/yolov7: Implementation of paper - YOLOv7: Trainable bag-of-freebies sets new state-of-the-art for real-time

이 사이트인데, yolov7에 대한 다양한 정보가 있다.

 

흘러가는 이야기이지만, yolo의 최초 개발자는 yolov5까지만 제작했고, 좋지 않은 목적으로 이용될 것을 염려하여 더이상의 개발을 멈췄기에 yolov5가 마지막 정식버전이다.

그런데 지금은 yolov8도 나왔고, 최근에 기가막힌 yolo-world까지 나왔다.

 

%cd yolov7

이 명령어로 폴더를 이동한다.

코랩은 로컬과 다르게 폴더 이동에 살짝 불편한 감이 있지만, GPU없는 본인을 탓하자.

 

!wget https://public.roboflow.com/ds/VjVsQFjtie?key=mGrHsJzXoS

이번 데이터셋은  Hard Hat Worker 이다. 현장에서 누가 헬맷 안썻는지, 누가 똑바로 안 쓰고 다니는지 detection한다.

실제로 이런 task로 현장에서 외주 맡기는 경우가 꽤 있다. 

https://public.roboflow.com/object-detection/hard-hat-workers

Hard Hat Workers Object Detection Dataset

!wget 명령어는 href 주석으로, 주소를 통해 데이터를 다운받을 수 있는 경우에 사용된다.

이 명령어로 zip 파일이 다운받아졌으면

!unzip VjVsQFjtie?key=mGrHsJzXoS

이 명령어로 unzip 해주자.

 

아까 %cd yolov7 로 이동했기 때문에 zip파일와 unzip파일은 yolov7 폴더 안에 만들어진다.

 

라이브러리 import를 위해 requirements 설치를 하자.

pip install wandb pandas opencv-python scikit-learn tensorboard tqdm

 

import os
import json
import shutil
from tqdm import tqdm

라이브러리 살짝 import 해주고

 

from glob import glob
train_img_list = glob('./train/images/*.jpg')
print(len(train_img_list))

>>5269

이제 train 데이터를 불러오자. glob 라이브러리로 와일드카드를 이용할 수 있다. *.jpg는 jpg형식의 모든 파일을 불러오겠다는 뜻이다. 만약 ./images/*s.jpg 로 한다면. images 폴더 안에 s로 끝나는 모든 jpg 파일을 불러오게된다. 

 

from sklearn.model_selection import train_test_split
train_list, valid_list = train_test_split(train_img_list, test_size=0.2, random_state = 2000)
print(len(train_list), len(valid_list))

sklearn의 random split으로 8:2 로 나눠준다.

 

with open('./train/train.txt', 'w') as f:
    f.write('\n'.join(train_img_list) + '\n')

이건 앞으로 아주아주 자주 사용하게 될 문법이다.

파일을 읽고, 생성하는 문법이다.

오늘은 txt파일 생성을 하지만,

특히나 detection, segmentation을 할 땐 json파일을 자주 정형하게될 건데, 그건 다음 시간을 기대하면 된다 !

 

'w'는 쓰기, 'r' 읽기 전용이다. as f 로 선언했으므로 train.txt.의 정보가 f에 저장되고, f를 통해 train.txt를 "쓰기" 할 거다.

join 명령어로 train_img_list를 묶어주는데 '\n' 를 중간중간에 넣어주니, 한 줄씩 데이터가 저장된다.

그리고 그 정보를 f를 통해 train.txt에 저장해준다.

즉, train_img_list의 정보를 한 줄 씩 train.txt에 쓰고, 저장한다.

 

# train data에 80%, valid에 20% 배정

with open('./train.txt', 'w') as f:
    f.write('\n'.join(train_list) + '\n')

with open('./valid.txt', 'w') as f:
    f.write('\n'.join(valid_list) + '\n')

아까 train.txt는 사실 필요없어서 삭제해도 되고,
train, valid를 나눠주자.

이렇게 txt에 저장해주는 이유는, train 폴더에 모든 train, valid 이미지가 저장돼있고,
train, valid.txt 에 저장된 파일명을 통해 dataloader가 불러가게만 하면 된다.

1. train, valid 폴더를 따로 관리해도 되고

2. train 폴더만 두고, train, valid.txt.로 호출 방식만 변경하는 것

두 가지 방식이 있다.  물론 모델마다 사용법이 다르지만 yolo에서는 나는 (2)번 방법을 선호한다.

( 혹시나 데이터셋 다시 섞으려면 폴더 새로 만들기 귀찮거든... )

 

import yaml

text = """
names : ['head', 'helmet', 'person']
nc: 3
train: ./train.txt
val: ./valid.txt
test : ./test/images
"""

with open('./custom_data.yaml','w') as f:
  f.write(text)

이번에는 yaml 파일을 새로 만들어줘야한다.

여기엔 class개수, class 이름, train, valid, test data 위치를 저장한다. yolo의 방식이다.

나는 train, valid, test를 txt위치로 저장해두고 알아서 데이터를 찾아가게 만들었다.

 

이제 train, valid, test 데이터셋 생성, yaml를 만들면서 데이터셋 생성은 완료했다.

 

!pip -q install wandb

 

이제 wandb를 설치하고

 

import wandb
wandb.login()

wandb에 로그인하자.

wandb API key는  QuickStarrt에 가면 볼 수 있다.

https://wandb.ai/authorize

Sign In with Auth0

나는 좀 더 검색하기 좋게 여기를 첨부하두겠다 !

 

!wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7x.pt

이거는 yolov7-x 모델의 weight 파일이다. 이게 있어야 훈련을 돌릴 수 있다.

 

이전에 첨부해둔 yolov7 깃허브에 들어가면 

 

이 표가 있는데, 여기서 파란글씨가 href 정보다. 이걸 누르면 .pt(=pre trained weight) file를 다운로드받을 수 있고, wget을 써서 내려받을 수도 있다.

 

우린 yolov7-x를 사용해보자.

 

그리고 이제 훈련을 돌려보자.

이때 표를 참고하여 이미지 사이즈는 640 x 640 으로 한다.

!python train.py --device 0 --batch-size 10 --epochs 2 --adam --img 640 640 --data ./custom_data.yaml --hyp ./data/hyp.scratch.custom.yaml --cfg ./cfg/training/yolov7x.yaml --weights yolov7x.pt --name yolov7x-custom

 

 

 

훈련이 다 됐다면 새로운 .pt파일이 생성됐을 것이다.

그걸 토대로 inference를 진행하면 된다.

!python _detect.py --device 0 --img-size 640 --conf-thres 0.01 --iou-thres 0.45 --nosave --weights ./yolov7x.pt --source ./test/images --name yolov7_x1_test_data

여기서 실행파일이 detect.py가 아닌, _detect.py인데, 대회 기간에 csv파일로 합쳐야할 일이 있어서 detect.py 의 출력부분을 조금 수정한 코드이다. 이 또한 깃허브에 올라가있기 때문에 복사를하든 내려받아서 실행해보는 것으로 하자 !

 

_detect.py의 코드는 다음과 같다

# _Detec.py 코드

import argparse
import time
from pathlib import Path

import cv2
import torch
import torch.backends.cudnn as cudnn
from numpy import random
import csv

from models.experimental import attempt_load
from utils.datasets import LoadStreams, LoadImages
from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \
    scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path
from utils.plots import plot_one_box
from utils.torch_utils import select_device, load_classifier, time_synchronized, TracedModel



def detect(save_img=False):
    source, weights, view_img, save_txt, imgsz, trace = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size, not opt.no_trace
    save_img = not opt.nosave and not source.endswith('.txt')  # save inference images
    webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
        ('rtsp://', 'rtmp://', 'http://', 'https://'))

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

    # Initialize
    set_logging()
    device = select_device(opt.device)
    half = device.type != 'cpu'  # half precision only supported on CUDA

    # Load model
    model = attempt_load(weights, map_location=device)  # load FP32 model
    stride = int(model.stride.max())  # model stride
    imgsz = check_img_size(imgsz, s=stride)  # check img_size

    if trace:
        model = TracedModel(model, device, opt.img_size)

    if half:
        model.half()  # to FP16

    # Second-stage classifier
    classify = False
    if classify:
        modelc = load_classifier(name='resnet101', n=2)  # initialize
        modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']).to(device).eval()

    # Set Dataloader
    vid_path, vid_writer = None, None
    if webcam:
        view_img = check_imshow()
        cudnn.benchmark = True  # set True to speed up constant image size inference
        dataset = LoadStreams(source, img_size=imgsz, stride=stride)
    else:
        dataset = LoadImages(source, img_size=imgsz, stride=stride)

    # Get names and colors
    names = model.module.names if hasattr(model, 'module') else model.names
    colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]

    # Run inference
    if device.type != 'cpu':
        model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))  # run once
    old_img_w = old_img_h = imgsz
    old_img_b = 1

    # 추가: CSV 파일을 저장할 리스트 초기화
    results_list = [['PredictionString','image_id']]

    t0 = time.time()
    for path, img, im0s, vid_cap in dataset:
        img = torch.from_numpy(img).to(device)
        img = img.half() if half else img.float()  # uint8 to fp16/32
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        # Warmup
        if device.type != 'cpu' and (old_img_b != img.shape[0] or old_img_h != img.shape[2] or old_img_w != img.shape[3]):
            old_img_b = img.shape[0]
            old_img_h = img.shape[2]
            old_img_w = img.shape[3]
            for i in range(3):
                model(img, augment=opt.augment)[0]

        # Inference
        t1 = time_synchronized()
        with torch.no_grad():   # Calculating gradients would cause a GPU memory leak
            pred = model(img, augment=opt.augment)[0]
        t2 = time_synchronized()

        # Apply NMS
        pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
        t3 = time_synchronized()

        # Apply Classifier
        if classify:
            pred = apply_classifier(pred, modelc, img, im0s)

        # Process detections
        for i, det in enumerate(pred):  # detections per image
            results_sublist = []
            results_sublist_row = ''
            if webcam:  # batch_size >= 1
                p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count
            else:
                p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)

            p = Path(p)  # to Path
            save_path = str(save_dir / p.name)  # img.jpg
            txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # img.txt
            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
            if len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

                # Print results
                for c in det[:, -1].unique():
                    n = (det[:, -1] == c).sum()  # detections per class
                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

                # 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 opt.save_conf else (cls, *xywh)  # label format
                        with open(txt_path + '.txt', 'a') as f:
                            f.write(('%g ' * len(line)).rstrip() % line + '\n')

                    if save_img or view_img:  # Add bbox to image
                        label = f'{names[int(cls)]} {conf:.2f}'
                        plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=1)
                    
                    # 추가: CSV 파일에 정보 추가 #
                    results_sublist_row += f'{int(cls.item())} {conf.item()} '
                    for i in xyxy:
                        results_sublist_row += f'{i} '
            # CSV 정보 모으기
            img_name = str(p.name).replace("'",'')
            results_sublist = [results_sublist_row, img_name]   
            results_list.append(results_sublist)
            
            # Print time (inference + NMS)
            print(f'{s}Done. ({(1E3 * (t2 - t1)):.1f}ms) Inference, ({(1E3 * (t3 - t2)):.1f}ms) NMS')

            # Stream results
            if view_img:
                cv2.imshow(str(p), im0)
                cv2.waitKey(1)  # 1 millisecond

            # Save results (image with detections)
            if save_img:
                if dataset.mode == 'image':
                    cv2.imwrite(save_path, im0)
                    print(f" The image with the result is saved in: {save_path}")
                else:  # 'video' or 'stream'
                    if vid_path != save_path:  # new video
                        vid_path = save_path
                        if isinstance(vid_writer, cv2.VideoWriter):
                            vid_writer.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 += '.mp4'
                        vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
                    vid_writer.write(im0)

    #csv 만들기
    with open("results_custom_test.csv", 'w') as file:
        writer = csv.writer(file)
        writer.writerows(results_list)
    print(f'Done. ({time.time() - t0:.3f}s)')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', nargs='+', type=str, default='yolov7.pt', help='model.pt path(s)')
    parser.add_argument('--source', type=str, default='inference/images', help='source')  # file/folder, 0 for webcam
    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--view-img', action='store_true', help='display results')
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
    parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument('--update', action='store_true', help='update all models')
    parser.add_argument('--project', default='runs/detect', help='save results to project/name')
    parser.add_argument('--name', default='exp', help='save results to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    parser.add_argument('--no-trace', action='store_true', help='don`t trace model')
    opt = parser.parse_args()
    print(opt)
    #check_requirements(exclude=('pycocotools', 'thop'))

    with torch.no_grad():
        if opt.update:  # update all models (to fix SourceChangeWarning)
            for opt.weights in ['yolov7.pt']:
                detect()
                strip_optimizer(opt.weights)
        else:
            detect()