yolov8
安装环境
装pytorch
conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.6 -c pytorch -c conda-forge
安装用到的包
requirements.txt
ultralytics
tqdm
scikit-learn
pyaml
shapely
pyclipper
scikit-image
imgaug
lmdb
tqdm
rapidfuzz
cython
Pillow
pyyaml
requests
pip install -r requirements.txt
数据标注
得到via_region_data.csv,格式
{
"42.243.238.32-202406181512034.jpg418013": {
"fileref": "",
"size": 418013,
"filename": "42.243.238.32-202406181512034.jpg",
"base64_img_data": "",
"file_attributes": {},
"regions":
[{
"shape_attributes": {
"name": "rect",
"x": 324,
"y": 179,
"width": 279,
"height": 108
},
"region_attributes": {}
},
{
"shape_attributes": {
"name": "rect",
"x": 766,
"y": 248,
"width": 227,
"height": 154
},
"region_attributes": {}
}]
},
}
其中(x,y)是左上角坐标, (x + width, y + height)是右下角坐标
经过统计:
平均:12142158
最小:7681360
最大:1440*2560
总 890, 408 background, 482 front
train 712, 325 background
val 89, 39 background
test 89, 41+? background
转换格式
转成yolov8的格式, 格式为
label_id x y w h
其中
x
,
y
,
w
,
h
∈
[
0
,
1
]
x, y, w, h \in \left[0, 1\right]
x,y,w,h∈[0,1]
转换
#!/usr/bin/env python
# -*- encoding: utf-8 -*-
import json
import os
import cv2
from tqdm import tqdm
img_dir = r"E:\ccc\origin_images"
json_path = r"E:\ccc\via_region_data.json"
label_dir = r"E:\ccc\origin_labels"
if __name__ == "__main__":
with open(json_path, "r") as f:
label_data: dict = json.load(f)
os.makedirs(label_dir, exist_ok=True)
cnt = 0
tot = 0
for label in tqdm(label_data.values()):
filename = label["filename"]
img = cv2.imread(os.path.join(img_dir, filename))
h, w, _ = img.shape
category = 0
empty_region = True
with open(
os.path.join(label_dir, os.path.splitext(filename)[0] + ".txt"), "w"
) as f:
for region in label["regions"]:
empty_region = False
x = region["shape_attributes"]["x"]
y = region["shape_attributes"]["y"]
width = region["shape_attributes"]["width"]
height = region["shape_attributes"]["height"]
x_center = (x + (x + width)) / 2.0
y_center = (y + (y + height)) / 2.0
x_center_norm = x_center / w
y_center_norm = y_center / h
width_norm = width / w
height_norm = height / h
f.write(
"{} {:.6f} {:.6f} {:.6f} {:.6f}\n".format(
category, x_center_norm, y_center_norm, width_norm, height_norm
)
)
if empty_region:
cnt += 1
tot += 1
print("non_empty: {}; empty: {}; total: {}".format(tot - cnt, cnt, tot))
划分数据集
#!/usr/bin/env python
# -*- encoding: utf-8 -*-
import os
import shutil
from glob import glob
import yaml
from sklearn.model_selection import train_test_split
from tqdm import tqdm
img_dir = r"E:\ccc\origin_images"
label_dir = r"E:\ccc\origin_labels"
train_ratio = 0.8
val_ratio = 0.1
test_ratio = 0.1
root_path = r"E:\ccc"
img_train_path = r"images\train"
img_val_path = r"images\val"
img_test_path = r"images\test"
label_train_path = r"labels\train"
label_val_path = r"labels\val"
label_test_path = r"labels\test"
yaml_path = "config.yaml"
def copy_files(label_paths, img_dir, target_label_dir, origin_img_dir):
os.makedirs(img_dir, exist_ok=True)
os.makedirs(target_label_dir, exist_ok=True)
for label in tqdm(label_paths):
img = os.path.join(
origin_img_dir, os.path.splitext(os.path.basename(label))[0] + ".jpg"
)
shutil.copy(img, img_dir)
shutil.copy(label, target_label_dir)
if __name__ == "__main__":
# img_paths = list(glob(os.path.join(img_dir, '*')))
label_paths = list(glob(os.path.join(label_dir, "*")))
# img_paths.sort()
label_paths.sort()
is_empty = []
for path in label_paths:
with open(path, "r") as f:
data = f.read().strip()
if not data:
is_empty.append(True)
else:
is_empty.append(False)
train_val_label, test_label, train_val_is_empty, test_is_empty = train_test_split(
label_paths, is_empty, test_size=test_ratio, shuffle=True
)
train_label, val_label, train_is_empty, val_is_empty = train_test_split(
train_val_label,
train_val_is_empty,
test_size=val_ratio / (train_ratio + val_ratio),
shuffle=True,
)
train_img = [os.path.splitext(path)[0] + ".jpg" for path in train_label]
val_img = [os.path.splitext(path)[0] + ".jpg" for path in val_label]
test_img = [os.path.splitext(path)[0] + ".jpg" for path in test_label]
copy_files(
train_label,
os.path.join(root_path, img_train_path),
os.path.join(root_path, label_train_path),
img_dir,
)
copy_files(
val_label,
os.path.join(root_path, img_val_path),
os.path.join(root_path, label_val_path),
img_dir,
)
copy_files(
test_label,
os.path.join(root_path, img_test_path),
os.path.join(root_path, label_test_path),
img_dir,
)
config = {
"path": root_path,
"train": img_train_path,
"val": img_val_path,
"test": img_test_path,
"names": {0: "box"},
}
yaml_dir = os.path.dirname(yaml_path)
if yaml_dir:
os.makedirs(yaml_dir, exist_ok=True)
with open(yaml_path, "w", encoding="utf-8") as f:
yaml.dump(config, f)
会自动划分数据集和生成需要的yaml
参考:coco128.yaml
训练yolo
#!/usr/bin/env python
# -*- encoding: utf-8 -*-
from ultralytics import YOLO
if __name__ == '__main__':
# Load a model
# model = YOLO("yolov8n.yaml") # build a new model from YAML
model = YOLO("yolov8n.pt") # load a pretrained model (recommended for training)
# model = YOLO("yolov8n.yaml").load("yolov8n.pt") # build from YAML and transfer weights
# Train the model
# results = model.train(data="config.yaml", epochs=100, imgsz=640, device="cpu", single_cls=True, save=True, workers=0)
results = model.train(data="config.yaml", epochs=100, imgsz=640, device=[0], single_cls=True, save=True, workers=8)
测试yolo
#!/usr/bin/env python
# -*- encoding: utf-8 -*-
from ultralytics import YOLO
if __name__ == "__main__":
model = YOLO("runs/detect/train/weights/best.pt") # load a custom model
# Validate the model
metrics = model.val(
data="config.yaml", split="val", plots=True
) # no arguments needed, dataset and settings remembered
# metrics = model.val(
# data="config.yaml", split="test", plots=True
# ) # no arguments needed, dataset and settings remembered
# print(metrics)
print("map50-95: ", metrics.box.map) # map50-95
print("map50: ", metrics.box.map50) # map50
print("map75: ", metrics.box.map75) # map75
# print('map50-95 of each category: ',metrics.box.maps) # a list contains map50-95 of each category
预测yolo
#!/usr/bin/env python
# -*- encoding: utf-8 -*-
import os
from glob import glob
import cv2
from ultralytics import YOLO
model_path = "runs/detect/train/weights/best.pt"
imgs_dir = "/mnt/data/ccc/Eudemons_final"
gt_dir = "/mnt/data/ccc/Eudemons_v3/labels/test"
crop_dir = "/mnt/data/ccc/crop"
save_dir = "/mnt/data/ccc/overlay"
draw_predict = True
# draw_gt = True
draw_gt = False
pad = 0.035
def get_name_v1(filename, dir_name, idx):
return os.path.join(
dir_name, "{}{}.jpg".format(filename, f"_{idx}" if idx > 0 else "")
)
def get_name_v2(path, dir_name, idx):
filename, ext = os.path.splitext(os.path.basename(path))
mac = os.path.basename(os.path.dirname(path))
return os.path.join(
dir_name, "{}_{}{}{}".format(mac, filename, f"_{idx}" if idx > 0 else "", ext)
)
if __name__ == "__main__":
model = YOLO(model_path) # load a custom model
paths = list(glob(os.path.join(imgs_dir, "*", "*.jpg")))
draw_predict = draw_predict and save_dir
draw_gt = draw_gt and gt_dir and os.path.exists(gt_dir) and save_dir
if draw_predict:
os.makedirs(save_dir, exist_ok=True)
os.makedirs(crop_dir, exist_ok=True)
# results = model(os.path.join(imgs_dir, "*.jpg"), stream=True)
results = model(os.path.join(imgs_dir, "*", "*.jpg"), stream=True)
for result in results:
# result.save(os.path.basename(result.path))
path = result.path
filename = os.path.splitext(os.path.basename(path))[0]
h, w = result.orig_shape
orig_img = result.orig_img
boxes = result.boxes
boxes_num = boxes.shape[0]
confidence = boxes.conf.tolist()
xyxy = boxes.xyxy.tolist()
for idx, (x1, y1, x2, y2) in enumerate(xyxy):
x_pad = (x2 - x1) * pad
y_pad = (y2 - y1) * pad
x1 -= x_pad
y1 -= y_pad
x2 += x_pad
y2 += y_pad
x1 = max(0, int(x1))
y1 = max(0, int(y1))
x2 = min(w, int(x2))
y2 = min(h, int(y2))
cv2.imwrite(
# get_name_v1(filename, crop_dir, idx),
get_name_v2(path, crop_dir, idx),
orig_img[y1:y2, x1:x2],
)
if draw_predict:
for x1, y1, x2, y2 in xyxy:
x1 = max(0, int(x1))
y1 = max(0, int(y1))
x2 = min(w - 1, int(x2))
y2 = min(h - 1, int(y2))
orig_img = cv2.rectangle(
orig_img,
(x1, y1),
(x2, y2),
color=(0, 0, 255),
lineType=cv2.LINE_AA,
)
if draw_gt:
with open(os.path.join(gt_dir, filename + ".txt"), "r") as f:
for line in f:
line = line.strip()
if not line:
continue
_, x, y, width, height = line.split(" ")
x = float(x)
y = float(y)
width = float(width)
height = float(height)
orig_img = cv2.rectangle(
orig_img,
(int((x - width / 2) * w), int((y - height / 2) * h)),
(int((x + width / 2) * w), int((y + height / 2) * h)),
color=(0, 255, 0),
lineType=cv2.LINE_AA,
)
cv2.imwrite(
# os.path.join(save_dir, filename + ".jpg"),
get_name_v2(path, save_dir, 0),
orig_img,
)