Yolov9 Jupyter에서 돌려보기

환경

CPU: AMD 5900x

GPU: AMD RX 7900xtx 24GB

Memory: 128GB

OS: Ubuntu 20.04

GPU Driver: ROCm 6.0.2

Python: Python 3.9

 

공식 문서

https://docs.ultralytics.com/models/yolov9/#conclusion

YOLOv9

 

 

1. Install Ultralytics Package & import

 

pip install ultralytics

# import packages
import os
import shutil
from glob import glob
import yaml
from ultralytics import YOLO
from tqdm import tqdm

 

2. Model Download

 

공식 문서 링크에서 사용할 모델 다운로드

 

3. Data Preparation

 

 - 학습에 사용할 데이터 경로를 yaml 파일로 정의

 - images 폴더에는 이미지가 배치되고 labels 폴더에는 상응하는 label들이 배치

 - data.yaml(예시)

# 클래스 명 정의
names:
- Car
# 클래스 수 정의
nc: 1
# 학습 및 검증 데이터 셋 경로
train: /data/Car/train_base_yolo/images/
val: /data/Car/valid_base_yolo/images/

 

4. Cfg Preparation

 

 - 디폴트 cfg 파일을 아래의 링크를 통해 준비

https://github.com/ultralytics/ultralytics/blob/main/ultralytics/cfg/default.yaml

 - epochs, batch, imgsz, save_period, cache, device, worker 등을 기호에 맞게 수정

 - cfg.yaml (예시)

# Ultralytics YOLO 🚀, AGPL-3.0 license
# Default training settings and hyperparameters for medium-augmentation COCO training

task: detect # (str) YOLO task, i.e. detect, segment, classify, pose
mode: train # (str) YOLO mode, i.e. train, val, predict, export, track, benchmark

# Train settings -------------------------------------------------------------------------------------------------------
model: # (str, optional) path to model file, i.e. yolov8n.pt, yolov8n.yaml
data: # (str, optional) path to data file, i.e. coco8.yaml
epochs: 100 # (int) number of epochs to train for
time: # (float, optional) number of hours to train for, overrides epochs if supplied
patience: 100 # (int) epochs to wait for no observable improvement for early stopping of training
batch: 16 # (int) number of images per batch (-1 for AutoBatch)
imgsz: 640 # (int | list) input images size as int for train and val modes, or list[w,h] for predict and export modes
save: True # (bool) save train checkpoints and predict results
save_period: -1 # (int) Save checkpoint every x epochs (disabled if < 1)
cache: False # (bool) True/ram, disk or False. Use cache for data loading
device: # (int | str | list, optional) device to run on, i.e. cuda device=0 or device=0,1,2,3 or device=cpu
workers: 8 # (int) number of worker threads for data loading (per RANK if DDP)
project: # (str, optional) project name

 

5. Train 준비

 

model_path = "./models_yolo/yolov9c" # 위에서 받은 모델이 위치한 경로
dataset_name = "data.yaml" # 데이터 셋 경로가 정의된 yaml 파일의 경로
model = YOLO(f"{model_path}.pt") # Model Initiation
cfg = 'cfg.yaml' # cfg 파일의 경로

 

6. Train

 

model.train(project="yolo9c_test", data=dataset_name, cfg=cfg, epochs=3)

 

output:

 

 

7. Inference

 

 - 학습을 시작하면 project 명으로 된 폴더가 생성이 되며 weigth 등이 저장

 - 저장된 모델을 load 하여 inference

 - 공식 문서를 참고하면 시각화까지 지원

 - 이 글에서는 여러장의 이미지에 대해 inference 결과를 yolo format의 txt파일로 저장

 

# 이미지와 모델 path 정의
imgs_to_update = glob(f'/data/to_inf/*/*.jpg')
finetuned_model = YOLO(f'./yolo9c_test/train/weights/last.pt')

# Yolo format txt 결과 저장, conf 값 조정하여 confidence 조정
for img in tqdm(imgs_to_update):
    results = finetuned_model(img, save_txt=None, conf=0.5, verbose=False)  # Assuming 'model' processes the image and returns detection results
    with open(img[:-4] + '.txt', 'w') as file:
        try:
            # Ensure that 'results' contains all detections and iterate over each one
            for detection in results:
                # Assuming each 'detection' has a 'boxes' attribute with all bounding boxes
                for box in detection.boxes:
                    cls = box.cls.cpu().numpy()[0]  # Class for the current box
                    # Extract normalized (xywhn) bounding box coordinates
                    x, y, w, h = box.xywhn.cpu().numpy()[0]
                    # Write class and bounding box coordinates to the file
                    file.write(f"{cls} {x} {y} {w} {h}\n")
        except Exception as e:
            print(f"Error processing image {img}: {e}")
            file.write('')

 

 

끝.