YOLOv5

!git clone https://github.com/ultralytics/yolov5  # clone repo
%cd yolov5
%pip install -qr requirements.txt  # install dependencies

import torch
from IPython.display import Image, clear_output  # to display images

clear_output()
print(f"Setup complete. Using torch {torch.__version__} ({torch.cuda.get_device_properties(0).name if torch.cuda.is_available() else 'CPU'})")

#https://roboflow.com/에서 이미지 데이터 load
!curl -L "https://public.roboflow.com/ds/uUe1ngFbSE?key=8Y8K5oBcki" > roboflow.zip; unzip roboflow.zip; rm roboflow.zip

from glob import glob
from sklearn.model_selection import train_test_split
import yaml

#jpg파일 리스트에 담기
%cd /
img_list = glob('/content/dataset/export/*.jpg')

print(len(img_list))


#train, test 분리
train_img_list, val_img_list = train_test_split(img_list, test_size=0.2, random_state=2000)

print(len(train_img_list), len(val_img_list))


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

with open('/content/dataset/export/val.txt', 'w') as f:
  f.write('\n'.join(val_img_list)+'\n')
  
%cat /content/dataset/data.yaml
    

with open('/content/dataset/data.yaml', 'r') as f:
  data = yaml.load(f)

print(data)

data['train'] = '/content/dataset/export/train.txt'
data['val'] = '/content/dataset/export/val.txt'

with open('/content/dataset/data.yaml', 'w') as f:
  yaml.dump(data, f)

print(data)

#모델 학습
%cd /content/swin-transformer-pytorch/yolov5/

!python train.py --img 416 --batch 16 --epochs 50 --data /content/dataset/data.yaml --cfg ./models/yolov5s.yaml --weights yolov5s.pt --name gun_yolov5s_results

from utils.plots import plot_results 
#plot_results(save_dir='runs/train/exp')  # plot all results*.txt as results.png
Image(filename='/content/swin-transformer-pytorch/yolov5/runs/train/gun_yolov5s_results/results.png', width=800)

Image(filename='/content/swin-transformer-pytorch/yolov5/runs/train/gun_yolov5s_results/train_batch0.jpg', width=800)  # train batch 0 mosaics and labels
Image(filename='/content/swin-transformer-pytorch/yolov5/runs/train/gun_yolov5s_results/test_batch0_labels.jpg', width=800)  # test batch 0 labels
Image(filename='/content/swin-transformer-pytorch/yolov5/runs/train/gun_yolov5s_results/test_batch0_pred.jpg', width=800)  # test batch 0 predictions