Object Detection YOLO

Object Detection YOLO

---

Prerequisites

---

1. Object Detection YOLO

1.1 Example

Convert Data for YOLO

import shutil
from pathlib import Path

import pandas as pd
from sklearn.model_selection import train_test_split
from tqdm import tqdm


ROOT_DIR = Path(__file__).resolve().parents[1]  # ComputerVision

RAW_DIR = ROOT_DIR / "data" / "gwhd_2021"
IMAGE_DIR = RAW_DIR / "images"
CSV_PATH = RAW_DIR / "competition_train.csv"

OUT_DIR = ROOT_DIR / "dataset_gwhd"
DATA_YAML_PATH = ROOT_DIR / "gwhd.yaml"

TRAIN_RATIO = 0.8
CLASS_ID = 0
IMG_W = 1024
IMG_H = 1024


def make_dirs():
    for split in ["train", "val"]:
        (OUT_DIR / "images" / split).mkdir(parents=True, exist_ok=True)
        (OUT_DIR / "labels" / split).mkdir(parents=True, exist_ok=True)


def find_image(image_name: str):
    path = IMAGE_DIR / image_name
    if path.exists():
        return path

    for ext in [".jpg", ".jpeg", ".png"]:
        path = IMAGE_DIR / f"{image_name}{ext}"
        if path.exists():
            return path

    return None


def parse_boxes_string(boxes_string):
    boxes = []

    if pd.isna(boxes_string):
        return boxes

    boxes_string = str(boxes_string).strip()

    if boxes_string == "" or boxes_string.lower() == "no_box":
        return boxes

    for box_text in boxes_string.split(";"):
        box_text = box_text.strip()
        if not box_text:
            continue

        values = box_text.split()

        if len(values) != 4:
            continue

        x1, y1, x2, y2 = map(float, values)
        boxes.append((x1, y1, x2, y2))

    return boxes


def xyxy_to_yolo(x1, y1, x2, y2, img_w, img_h):
    x1 = max(0.0, min(img_w, x1))
    y1 = max(0.0, min(img_h, y1))
    x2 = max(0.0, min(img_w, x2))
    y2 = max(0.0, min(img_h, y2))

    w = x2 - x1
    h = y2 - y1

    if w <= 0 or h <= 0:
        return None

    cx = x1 + w / 2.0
    cy = y1 + h / 2.0

    return cx / img_w, cy / img_h, w / img_w, h / img_h


def main():
    print("ROOT_DIR:", ROOT_DIR)
    print("CSV_PATH:", CSV_PATH)
    print("IMAGE_DIR:", IMAGE_DIR)

    make_dirs()

    df = pd.read_csv(CSV_PATH)

    print("CSV columns:", list(df.columns))
    print(df.head())

    required_cols = ["image_name", "BoxesString"]
    for col in required_cols:
        if col not in df.columns:
            raise RuntimeError(f"Required column missing: {col}")

    image_names = df["image_name"].unique()

    train_names, val_names = train_test_split(
        image_names,
        train_size=TRAIN_RATIO,
        random_state=42,
        shuffle=True,
    )

    split_map = {name: "train" for name in train_names}
    split_map.update({name: "val" for name in val_names})

    missing_images = 0
    total_boxes = 0

    for _, row in tqdm(df.iterrows(), total=len(df), desc="Converting"):
        image_name = str(row["image_name"])
        boxes_string = row["BoxesString"]

        split = split_map[image_name]

        src_img = find_image(image_name)

        if src_img is None:
            missing_images += 1
            print("Missing image:", image_name)
            continue

        dst_img = OUT_DIR / "images" / split / src_img.name
        shutil.copy2(src_img, dst_img)

        label_path = OUT_DIR / "labels" / split / f"{src_img.stem}.txt"

        boxes = parse_boxes_string(boxes_string)

        lines = []

        for x1, y1, x2, y2 in boxes:
            yolo_box = xyxy_to_yolo(x1, y1, x2, y2, IMG_W, IMG_H)

            if yolo_box is None:
                continue

            cx, cy, w, h = yolo_box
            lines.append(f"{CLASS_ID} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}")

        total_boxes += len(lines)

        with open(label_path, "w", encoding="utf-8") as f:
            f.write("\n".join(lines))

    with open(DATA_YAML_PATH, "w", encoding="utf-8") as f:
        f.write(
            f"path: {OUT_DIR.as_posix()}\n"
            "train: images/train\n"
            "val: images/val\n\n"
            "names:\n"
            "  0: wheat_head\n"
        )

    print("Done.")
    print("Missing images:", missing_images)
    print("Total boxes:", total_boxes)
    print("YOLO dataset:", OUT_DIR)
    print("YAML:", DATA_YAML_PATH)


if __name__ == "__main__":
    main()

Train YOLO

from pathlib import Path
from ultralytics import YOLO

ROOT_DIR = Path(__file__).resolve().parents[1]
DATA_YAML = ROOT_DIR / "gwhd.yaml"

model = YOLO("yolo11n.pt")

model.train(
    data=str(DATA_YAML),
    epochs=16,
    imgsz=640,
    batch=2,
    device="cpu",
    workers=0,
    project=str(ROOT_DIR / "runs"),
    name="gwhd_16epoch_640",
)

Prediect

from pathlib import Path
from ultralytics import YOLO

ROOT_DIR = Path(__file__).resolve().parents[1]

MODEL_PATH = ROOT_DIR / "runs" / "gwhd_16epoch_640" / "weights" / "best.pt"

SOURCE_PATH = ROOT_DIR / "dataset_gwhd" / "images" / "val"

model = YOLO(str(MODEL_PATH))

model.predict(
    source=str(SOURCE_PATH),
    imgsz=320,
    conf=0.25,
    save=True,
    save_txt=True,
    project=str(ROOT_DIR / "runs" / "detect"),
    name="predict_gwhd",
)