!git clone https://github.com/ultralytics/yolov5 # clone repo
!pip install -r yolov5/requirements.txt # install dependencies
%cd yolov5
import torch
print('Setup complete. Using torch %s %s' % (torch.__version__, torch.cuda.get_device_properties(0) if torch.cuda.is_available() else 'CPU'))
from IPython.display import Image
import pandas as pd
import random
import os
from shutil import copyfile
train = pd.read_csv('train.csv')
train
## Make directory structure
os.mkdir('/working/data')
os.mkdir('/working/data/images')
os.mkdir('/working/data/labels')
os.mkdir('/working/data/images/train')
os.mkdir('/working/data/images/valid')
os.mkdir('/working/data/labels/train')
os.mkdir('/working/data/labels/valid')
# copy images + labels to the training location, with 30% validation data
# can be done by hand as well
for image_id in list(set(train.image_id)):
image_data = train[train['image_id'] == image_id]
image_bboxes = image_data['bbox']
image_bboxes = image_bboxes.apply(lambda x: x.strip('[').strip(']').split(', '))
if random.random() > 0.3:
path = 'train'
else:
path = 'valid'
with open('/working/data/labels/{}/'.format(path) + image_id + '.txt', 'w+') as file:
for bbox in image_bboxes:
xc, yc, w, h = bbox
x_center_n = (float(xc) + float(w) / 2) / 1024.
y_center_n = (float(yc) + float(h) / 2) / 1024.
width_n = float(w) / 1024.
height_n = float(h) / 1024.
line = ' '.join(('0', str(x_center_n), str(y_center_n), str(width_n), str(height_n))) + '\n'
file.write(line)
copyfile('/input/global-wheat-detection/train/' + image_id + '.jpg', '/working/data/images/{}/'.format(path) + image_id + '.jpg')
The first one is easy, it is simply a copy of the yolo s (small), but with nc = 1, because we have only 1 class
with open('/working/new_train_yaml', 'w+') as file:
file.write(
"""
# parameters
nc: 1 # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.50 # layer channel multiple
# anchors
anchors:
- [10,13, 16,30, 33,23] # P3/8
- [30,61, 62,45, 59,119] # P4/16
- [116,90, 156,198, 373,326] # P5/32
# YOLOv5 backbone
backbone:
# [from, number, module, args]
[[-1, 1, Focus, [64, 3]], # 0-P1/2
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
[-1, 3, BottleneckCSP, [128]],
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
[-1, 9, BottleneckCSP, [256]],
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
[-1, 9, BottleneckCSP, [512]],
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
[-1, 1, SPP, [1024, [5, 9, 13]]],
[-1, 3, BottleneckCSP, [1024, False]], # 9
]
# YOLOv5 head
head:
[[-1, 1, Conv, [512, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 6], 1, Concat, [1]], # cat backbone P4
[-1, 3, BottleneckCSP, [512, False]], # 13
[-1, 1, Conv, [256, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 4], 1, Concat, [1]], # cat backbone P3
[-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small)
[-1, 1, Conv, [256, 3, 2]],
[[-1, 14], 1, Concat, [1]], # cat head P4
[-1, 3, BottleneckCSP, [512, False]], # 20 (P4/16-medium)
[-1, 1, Conv, [512, 3, 2]],
[[-1, 10], 1, Concat, [1]], # cat head P5
[-1, 3, BottleneckCSP, [1024, False]], # 23 (P5/32-large)
[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
]
"""
)
The second one is also easy, but it has to be in accordance with the location of the data.
with open('/working/new_data_yaml', 'w+') as file:
file.write(
"""
train: /working/data/images/train
val: /working/data/images/valid
nc: 1
names: ['wheat']
"""
)
os.chdir('/working/yolov5')
%%time
!python train.py --img 400 --batch 16 --epochs 3 --data '/working/new_data_yaml' --cfg '/working/new_train_yaml' --weights '' --name joos --nosave --cache
# find the correct weigths 1
!ls /working/yolov5/runs/
# find the correct weigths 2
!ls /working/yolov5/runs/exp0_joos
# find the correct weigths 3
!ls /working/yolov5/runs/exp0_joos/weights
# use you weigths in the detection
!python detect.py --source /input/global-wheat-detection/test --weights 'runs/exp0_joos/weights/last_joos.pt' --img 416 --conf 0.5 --save-txt
# find the output: annotated images and text files with the bounding box locations
!ls /working/yolov5/inference/output/