Как повысить точность YOLO11?
Разбираюсь с нейросетями и машинным обучением. Решил сделать небольшой проект для обнаружения лего минифигурок. Немного уростил задачу - пытаюсь не узнать какая где минифигурка, а просто найти места, где они есть на изображении.
Пробовал это реализовать двумя подходами.
Первый - это собрал 300 реальных изображений с открытых источников, где несколько минифигурок. Обвел их в LabelStudio вручную и собрал из них датасет. Обучил с такими параметрами yolo detect train data=/content/dataset/data.yaml model=yolo11n.pt epochs=30 imgsz=640
.
Второй - взял несколько изображений минифигурок (около 200). Поизменял их и сделал 2000 изображений с несколькими минифигурками. Получилось что-то вроде изображения ниже. Тренировал идентично команде выше.
Но почему-то в обоих случаях при тестах модель не может определить минифигурки даже на простых фонах. Почему она тут не нашла одну? И почему получилась такая низкая уверенность? Думал, что с таким фоном должна быть точность 0.98+. Если минифигурки перевернуты, то вообще ничего не находит.
Можете подсказать, что делаю не так? Ниже приложил полный код для второго варианта с генерацией датасета. Или мне просто нужно больше изображений?
import os
import random
from glob import glob
import cv2
import numpy as np
import yaml
from tqdm import tqdm
def add_flip(image, chance=0.5):
return cv2.flip(image, 1) if random.random() < chance else image
def add_scale(image, scale_strength=(0.2, 0.5)):
h, w = image.shape[:2]
s = random.uniform(*scale_strength)
return cv2.resize(image, (max(1, int(w * s)), max(1, int(h * s))), interpolation=cv2.INTER_LINEAR)
def add_rotation(image, rotation_strength=(-180, 180)):
h, w = image.shape[:2]
angle = random.uniform(*rotation_strength)
M = cv2.getRotationMatrix2D((w // 2, h // 2), angle, 1.0)
cos, sin = abs(M[0, 0]), abs(M[0, 1])
new_w, new_h = int(h * sin + w * cos), int(h * cos + w * sin)
M[0, 2] += (new_w - w) // 2
M[1, 2] += (new_h - h) // 2
rotated = cv2.warpAffine(image, M, (new_w, new_h), borderValue=(0, 0, 0, 0))
# === crop to alpha bounding box ===
alpha = rotated[:, :, 3]
ys, xs = np.where(alpha > 0)
if len(xs) == 0 or len(ys) == 0:
return rotated # fully transparent, return as is
x1, x2 = xs.min(), xs.max()
y1, y2 = ys.min(), ys.max()
cropped = rotated[y1:y2 + 1, x1:x2 + 1]
return cropped
def add_blur_whole(img, blur_strength=(0, 10)):
k = random.randint(*blur_strength)
if k > 0:
if k % 2 == 0:
k += 1
img = cv2.GaussianBlur(img, (k, k), 0)
return img
def add_noise_whole(img, noise_strength=(0.0, 0.05)):
noise = np.random.normal(0, 255 * random.uniform(*noise_strength), img.shape).astype(np.float32)
noisy = np.clip(img.astype(np.float32) + noise, 0, 255).astype(np.uint8)
return noisy
def generate_background(width, height):
choice = random.choice(["multi_color", "gradient", "blocks"])
if choice == "multi_color":
bg = np.zeros((height, width, 3), dtype=np.uint8)
for i in range(random.randint(2, 6)):
color = [random.randint(0, 255) for _ in range(3)]
x1, y1 = random.randint(0, width - 1), random.randint(0, height - 1)
x2, y2 = random.randint(x1, width), random.randint(y1, height)
cv2.rectangle(bg, (x1, y1), (x2, y2), color, -1)
elif choice == "gradient":
x = np.linspace(0, 1, width)
y = np.linspace(0, 1, height)
X, Y = np.meshgrid(x, y)
c1, c2 = np.random.randint(0, 255, 3), np.random.randint(0, 255, 3)
grad = (X[..., None] * c1 + (1 - X[..., None]) * c2).astype(np.uint8)
bg = np.tile(grad, (1, 1, 1))
elif choice == "blocks":
bg = np.zeros((height, width, 3), dtype=np.uint8)
block_size = random.randint(20, 80)
for y in range(0, height, block_size):
for x in range(0, width, block_size):
color = [random.randint(0, 255) for _ in range(3)]
bg[y:y + block_size, x:x + block_size] = color
return bg
# ============ PASTE WITH ALPHA ============
def paste_image(bg, fg, x, y):
h, w = fg.shape[:2]
x1, y1 = max(0, x), max(0, y)
x2, y2 = min(bg.shape[1], x + w), min(bg.shape[0], y + h)
fg_x1, fg_y1 = max(0, -x), max(0, -y)
fg_x2, fg_y2 = fg_x1 + (x2 - x1), fg_y1 + (y2 - y1)
if x1 >= x2 or y1 >= y2:
return bg, None
roi = bg[y1:y2, x1:x2]
fg_crop = fg[fg_y1:fg_y2, fg_x1:fg_x2]
fg_rgb, fg_a = fg_crop[:, :, :3], fg_crop[:, :, 3] / 255.0
for c in range(3):
roi[:, :, c] = roi[:, :, c] * (1 - fg_a) + fg_rgb[:, :, c] * fg_a
bg[y1:y2, x1:x2] = roi
return bg, (x1, y1, x2, y2)
def iou(box1, box2):
x1 = max(box1[0], box2[0])
y1 = max(box1[1], box2[1])
x2 = min(box1[2], box2[2])
y2 = min(box1[3], box2[3])
inter = max(0, x2 - x1) * max(0, y2 - y1)
area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
return inter / max(area1, area2) if max(area1, area2) > 0 else 0
def generate_dataset(
src_dir, out_dir="dataset", n_images=2000, img_size=(800, 800), val_split=0.2
):
img_train = os.path.join(out_dir, "images/train")
img_val = os.path.join(out_dir, "images/val")
lbl_train = os.path.join(out_dir, "labels/train")
lbl_val = os.path.join(out_dir, "labels/val")
for d in [img_train, img_val, lbl_train, lbl_val]:
os.makedirs(d, exist_ok=True)
minifigs = glob(os.path.join(src_dir, "*.png"))
n_val = int(n_images * val_split)
for idx in tqdm(range(n_images)):
bg = generate_background(*img_size)
boxes = []
n_objs = random.randint(1, 20)
for _ in range(n_objs):
img = cv2.imread(random.choice(minifigs), cv2.IMREAD_UNCHANGED)
if img is None:
continue
img = add_flip(img)
img = add_scale(img)
img = add_rotation(img)
h, w = img.shape[:2]
# random placement with retries
for attempt in range(20):
x = random.randint(-w // 4, img_size[0] - 3 * w // 4)
y = random.randint(-h // 4, img_size[1] - 3 * h // 4)
_, bbox = paste_image(bg.copy(), img, x, y)
if bbox is None:
continue
bx1, by1, bx2, by2 = bbox
visible_area = (bx2 - bx1) * (by2 - by1)
full_area = w * h
if visible_area < 0.5 * full_area:
continue
ok = True
for b in boxes:
if iou(b, bbox) > 0.5:
ok = False
break
if ok:
bg, bbox = paste_image(bg, img, x, y)
boxes.append(bbox)
break
bg = add_blur_whole(bg)
bg = add_noise_whole(bg)
# split train/val
if idx < n_val:
img_dir, lbl_dir = img_val, lbl_val
else:
img_dir, lbl_dir = img_train, lbl_train
fname = f"{idx:05d}"
cv2.imwrite(os.path.join(img_dir, fname + ".jpg"), bg)
with open(os.path.join(lbl_dir, fname + ".txt"), "w") as f:
for (x1, y1, x2, y2) in boxes:
x_center = (x1 + x2) / 2 / img_size[0]
y_center = (y1 + y2) / 2 / img_size[1]
bw = (x2 - x1) / img_size[0]
bh = (y2 - y1) / img_size[1]
f.write(f"0 {x_center:.6f} {y_center:.6f} {bw:.6f} {bh:.6f}\n")
data = {
"train": os.path.join(out_dir, "images/train"),
"val": os.path.join(out_dir, "images/val"),
"nc": 1,
"names": ["minifigure"]
}
with open(os.path.join(out_dir, "data.yaml"), "w") as f:
yaml.dump(data, f, default_flow_style=False)
UPD1:
Попробовал увеличить количество эпох до 100 в первом варианте с синтетическими данными. Это повысило точность. Сейчас распознаются фигурки на простых фонах с вероятностью 0.8-0.91 даже если они перевернуты. Но всё равно кажется, что этого не достаточно. Попробовал 200 эпох, но модель начала переобучаться. Еще попробовал 5000 синтетических изображений и 100 эпох, но модель также переобучилась.