YOLO v8 예제를 통해 사용해 보기

Spacenet·2023년 11월 6일
YOLOYOLOv8
1
post-thumbnail

YOLO v8

YOLO v8 Git : https://github.com/ultralytics/ultralytics
YOLO v8 Document : https://docs.ultralytics.com/

Git에서 YOLO v8을 다운 받고 원하는 위치에 압축을 푼다 (프로젝트 폴더 등)

사전 준비

데이터 셋 준비

연습용 데이터 셋 : https://universe.roboflow.com/ds/6bCd8cGyu0?key=UjNo8A0wm5
철탑에서 이음매를 찾는 데이터이다

이 데이터 셋을 아래 위치에 압축을 풀고 넣어 둔다

ultralytics-main\ultralytics\cfg\

데이터 셋 폴더 이름을 cell_towers_dataset 이라고 변경함

cfg는 config(설정)폴더인데 설정 파일과 데이터 셋이 같이 있으면 설정 파일에서 데이터 셋의 경로를 자동으로 인식하기 때문에 보통 여기에 두는걸 추천한다

데이터셋 구조는 이미지 데이터 셋 기준으로

  • 데이터 폴더
    • train
      • images
      • labels
    • test
      • images
      • labels
    • vaild
      • images
      • labels

이런 구조이여야 한다

데이터셋 yaml 파일 생성

~\cfg\datasets 폴더에 cell_towers.yaml 파일을 생성한다
다운받은 데이터 셋의 data.yaml 파일을 보면

train: ../train/images
val: ../valid/images
test: ../test/images

nc: 2
names: ['joint', 'side']

roboflow:
  workspace: roboflow-100
  project: cell-towers
  version: 2
  license: CC BY 4.0
  url: https://universe.roboflow.com/roboflow-100/cell-towers/dataset/2

이렇게 v5용으로 작성 되있는데 names가 joint와 side다
이걸 기준으로 v8용으로 작성하면

train: ../cell_towers_dataset/train/images/
val: ../cell_towers_dataset/valid/images/

# Classes
names:
  0: joint
  1: side

이렇게 작성하면 된다

하이퍼파라미터 설정

~\cfg폴더에 default.yaml을 열어보면 밑쪽에 Hyperparameters 부분이 있다
YOLO v8은 하이퍼파라미터의 기본값을 여기에서 설정한다
원하는 대로 설정하면 된다

어그멘테이션 설정

~\ultralytics\data폴더에 보면 augment.py파일이 있는데 해당 파일을 열어서 ctrl+f로 Albumentations 클래스를 찾아보자
그러면 __init__ 함수에 T라고 된 리스트가 있는데 이게 augmentions 다

T = [
    A.Blur(p=0.01),
    A.MedianBlur(p=0.01),
    A.ToGray(p=0.01),
    A.CLAHE(p=0.01),
    A.RandomBrightnessContrast(p=0.0),
    A.RandomGamma(p=0.0),
    A.RandomShadow(p=0.8),
    A.ImageCompression(quality_lower=75, p=0.0)]  # transforms

여기서 본인이 원하는대로 수정하거나 추가하면 된다

옵티마이저 설정

~\ultralytics\engine폴더의 trainer.py파일을 열고 build_optimizer 함수를 찾아보자
그러면

if name in ('Adam', 'Adamax', 'AdamW', 'NAdam', 'RAdam'):
    optimizer = getattr(optim, name, optim.Adam)(g[2], lr=lr, betas=(momentum, 0.999), weight_decay=0.0)
elif name == 'RMSProp':
    optimizer = optim.RMSprop(g[2], lr=lr, momentum=momentum)
elif name == 'SGD':
    optimizer = optim.SGD(g[2], lr=lr, momentum=momentum, nesterov=True)
else:
    raise NotImplementedError(
        f"Optimizer '{name}' not found in list of available optimizers "
        f'[Adam, AdamW, NAdam, RAdam, RMSProp, SGD, auto].'
        'To request support for addition optimizers please visit https://github.com/ultralytics/ultralytics.')

이 코드가 보이는데 여기서 optimizer를 추가하거나 설정하면 된다

데이터로더 설정

커스텀 데이터로더를 설정하고 싶다면 ~\ultralytics\data 폴더의 build.py파일로 가보자
여기서 맨 밑의 load_inference_source 함수를 보면

# Dataloader
if tensor:
    dataset = LoadTensor(source)
elif in_memory:
    dataset = source
elif webcam:
    dataset = LoadStreams(source, imgsz=imgsz, vid_stride=vid_stride, buffer=buffer)
elif screenshot:
    dataset = LoadScreenshots(source, imgsz=imgsz)
elif from_img:
    dataset = LoadPilAndNumpy(source, imgsz=imgsz)
else:
    dataset = LoadImages(source, imgsz=imgsz, vid_stride=vid_stride)

이 코드에서 설정하면 된다

학습 하기

YOLO v8 폴더에 cell_tower_train.py 파일을 생성한다
ex) ultralytics-main\cell_tower_train.py

다음과 같이 작성한다

from ultralytics import YOLO

if __name__ == "__main__":
    model = YOLO("yolov8m.pt") # YOLO 모델 설정
    model.train(data="./cell_towers.yaml", epochs=100, batch=32, lrf=0.001)

그리고 실행 시키면 학습 시작!

epochs, batch, lrf는 본인이 원하는대로 작성하면 된다
상세내용은 https://docs.ultralytics.com/modes/train/#arguments

재시작 하는 법

혹시 모종의 이유로 학습이 중단되고 다시 학습할려고 할때 처음부터 시작하면 그동안 학습했던 시간이 날아가게 된다
이럴때 cell_tower_train.py에 다음 코드를 추가하여 실행하자

# !!!이전 코드는 주석 처리!!!
# resume
model = YOLO("./runs/detect/train/weights/last.pt") # .pt 파일 확인 요망
model.train(resume=True)

그러면 중단된 부분 부터 학습이 다시 시작된다

결과물

train 후 ~\runs\detect\에 보면 train 폴더가 생성돼 있다
여기에 결과물들이 저장된다

label.cache

학습 실행 후 데이터 셋 폴더를 보면 train폴더와 vaild폴더에 라벨폴더의 캐시 파일이 생성돼 있다

yolo v8은 라벨 정보를 빠르게 불러오기 위해 캐시 파일을 생성하는데 만약 학습 실행 후 라벨이나 데이터에 변화가 있다면 이 캐시 파일을 삭제 후 학습 진행 해야 한다

cache 파일이 자동 생성되는 건 v3부터 였다고 한다

nvidia-smi -l

cuda를 통해 gpu로 학습을 한다면 명령 프롬프트에 nvidia-smi -l를 실행시켜 보자
약 5초마다 현재 NVIDIA-SMI, Driver, CUDA 버전을 비롯한 그래픽 카드의 여러 사용 정보를 확인 할 수 있다

테스트 하기

cell_towers_test.py을 하나 생성한다
위치는 cell_towers_train.py과 동일하게 생성했다

from ultralytics import YOLO

model = YOLO("./runs/detect/train2/weights/best.pt")

results = model.predict(
    "./ultralytics/cfg/cell_towers_dataset/test/images/DJI_20210911155632_0255_Z_JPG_jpg.rf.96cc5edb922418007dc24cd044209970.jpg",
    save=False,
    imgsz=640,
    conf=0.5,
    device="cuda",
)
for r in results:
    print(r.boxes)

일단 이미지 한개만 테스트 해보자

arguments에 대한 자세한 정보는 https://docs.ultralytics.com/modes/predict/#inference-arguments

실행시키면 다음과 같이 출력되는데

209970.jpg: 480x640 4 joints, 64.5ms
Speed: 15.6ms preprocess, 64.5ms inference, 6.0ms postprocess per image at shape (1, 3, 480, 640)
ultralytics.engine.results.Boxes object with attributes:

cls: tensor([0., 0., 0., 0.], device='cuda:0')
conf: tensor([0.7859, 0.7701, 0.7644, 0.6778], device='cuda:0')
data: tensor([[8.2697e+02, ~~~, 0.0000e+00]], device='cuda:0')
~~~
xywh: tensor([[ 988.1722, ~~~, 399.2513]], device='cuda:0')
xywhn: tensor([[0.1906, ~~~, 0.1027]], device='cuda:0')
xyxy: tensor([[ 826.9733, ~~~, 1222.8093]], device='cuda:0')
xyxyn: tensor([[0.1595, ~~~, 0.3145]], device='cuda:0')

여기에서 xywh, xywhn, xyxy, xyxyn이 바운딩 박스 정보이고 원하는 걸 선택해서 사용하면 된다
지금은 xyxy를 사용해 보겠다

for r in results:
    boxes = r.boxes.xyxy
    
    for box in boxes:
        print(box)

추가 후 실행 시키면

tensor([ 826.9733,  669.3661, 1149.3712, 1028.0715], device='cuda:0')
tensor([3326.1255, 3454.4353, 3738.6965, 3785.2217], device='cuda:0')
tensor([2202.3970, 3361.6133, 2527.1470, 3713.1497], device='cuda:0')
tensor([3099.8506,  823.5580, 3548.7449, 1222.8093], device='cuda:0')

박스가 4개에 각 박스마다 xy값이 2개씩 나온다

다음과 같이 작성해보자

for r in results:
    image_path = r.path # 현재 이미지의 path
    boxes = r.boxes.xyxy # 현재 이미지의 bbox의 xy좌표값들
    cls = r.boxes.cls # 현재 이미지의 bbox의 class들
    conf = r.boxes.conf # 현재 이미지의 bbox의 conf값
    cls_dict = r.names # 지금 예제는 {0: 'joint', 1: 'side'}
    
    # boxes, cls, conf 개수는 같기 때문에 zip으로 한번 묶어준다
    for box, cls_number, conf in zip(boxes, cls, conf):
        conf_number = float(conf.item())
        cls_number_int = int(cls_number.item())
        cls_name = cls_dict[cls_number_int]
        x1, y1, x2, y2 = box
        x1_int = int(x1.item())
        y1_int = int(y1.item())
        x2_int = int(x2.item())
        y2_int = int(y2.item())
        print(x1_int, y1_int, x2_int, y2_int, cls_name)

실행하면

826 669 1149 1028 joint
3326 3454 3738 3785 joint
2202 3361 2527 3713 joint
3099 823 3548 1222 joint

각각의 xy좌표와 클래스 이름이 나온다
이제 이 xy좌표를 사용해 cv2로 이미지를 출력해 보자

for r in results:
    image_path = r.path # 현재 이미지의 path
    boxes = r.boxes.xyxy # 현재 이미지의 bbox의 xy좌표값들
    cls = r.boxes.cls # 현재 이미지의 bbox의 class들
    conf = r.boxes.conf # 현재 이미지의 bbox의 conf값
    cls_dict = r.names # 지금 예제는 {0: 'joint', 1: 'side'}
    
    import cv2

    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    h, w, c = image.shape

    image = cv2.resize(image, (640, 640)) # 출력할 이미지 사이즈 조정

    for box, cls_number, conf in zip(boxes, cls, conf):
        conf_number = float(conf.item())
        cls_number_int = int(cls_number.item())
        cls_name = cls_dict[cls_number_int]
        x1, y1, x2, y2 = box
        x1_int = int(x1.item())
        y1_int = int(y1.item())
        x2_int = int(x2.item())
        y2_int = int(y2.item())
        print(x1_int, y1_int, x2_int, y2_int, cls_name)

        # 출력할 이미지 사이즈를 조정했기 때문에 좌표값도 같이 조정 한다
        scale_factor_x = 640 / w
        scale_factor_y = 640 / h
        x1_scale = int(x1_int * scale_factor_x)
        y1_scale = int(y1_int * scale_factor_y)
        x2_scale = int(x2_int * scale_factor_x)
        y2_scale = int(y2_int * scale_factor_y)

        image = cv2.rectangle(
            image, (x1_scale, y1_scale), (x2_scale, y2_scale), (0, 225, 0), 6
        )

    # cv2.imwrite("./test.jpg", image) # 이미지 저장
    cv2.imshow("Test", image)
    cv2.waitKey(0)


잘 나온다!

profile
Spacenet
내가 볼려고 씀
이전 포스트

첫글

0개의 댓글

관련 채용 정보