Как повысить точность 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 эпох, но модель также переобучилась.


Ответы (0 шт):