Below is an accuracy-focused two-stage training script. It first trains the detection head with most of the backbone frozen, then unfreezes the entire model and fine-tunes it at a lower learning rate. Ultralytics specifically recommends two-stage fine-tuning for microscopy and other domains that differ substantially from COCO. (Ultralytics Docs)

It uses the official YOLO26s loss gains and moderate augmentation rather than the aggressive COCO pretraining augmentation. (Ultralytics Docs) It also enables the current cls_pw class-balancing mechanism. Ultralytics computes inverse-frequency class weights from the training annotations and applies them to classification loss; 0.5 uses damped square-root inverse weighting instead of potentially unstable full inverse weighting. (GitHub)

Install or update the packages first:

bash
pip install -U ultralytics albumentations

train.py

python
from pathlib import Path

import albumentations as A
import torch
from ultralytics import YOLO


# ============================================================================
# FIXED PATHS AND SETTINGS
# ============================================================================

PROJECT_ROOT = Path(
    r"C:\Users\BIOCELIZE003\Desktop\besho_detection\Besho_train"
)

DATA_YAML = PROJECT_ROOT / "custom_dataset.yaml"
INITIAL_MODEL = "yolo26s.pt"

OUTPUT_PROJECT = PROJECT_ROOT / "runs_sperm_yolo26s"

STAGE1_NAME = "stage1_frozen_backbone"
STAGE2_NAME = "stage2_full_finetuning"

IMAGE_SIZE = 1024

# Uses approximately 75% of available GPU memory.
# Ultralytics determines the largest safe batch automatically.
BATCH_SIZE = 0.75

DEVICE = 0
WORKERS = 4
SEED = 42


# ============================================================================
# MICROSCOPY-SPECIFIC ALBUMENTATIONS
# ============================================================================
#
# Keep these mild. Strong blur/noise/color augmentation can destroy the subtle
# morphology separating c1-c5.
#
# These transformations are added to Ultralytics' native augmentations.
# ============================================================================

MICROSCOPY_AUGMENTATIONS = [
    # Simulates differences in microscope contrast and illumination.
    A.CLAHE(
        clip_limit=2.0,
        tile_grid_size=(8, 8),
        p=0.20,
    ),

    # Mild defocus variation. Strong blur would hide sperm morphology.
    A.Blur(
        blur_limit=3,
        p=0.10,
    ),

    # Small brightness and contrast differences between microscope sessions.
    A.RandomBrightnessContrast(
        brightness_limit=0.10,
        contrast_limit=0.15,
        p=0.15,
    ),
]


# ============================================================================
# SHARED TRAINING CONFIGURATION
# ============================================================================

COMMON_TRAIN_SETTINGS = {
    # Dataset and image resolution
    "data": str(DATA_YAML),
    "imgsz": IMAGE_SIZE,
    "batch": BATCH_SIZE,

    # Hardware
    "device": DEVICE,
    "workers": WORKERS,
    "amp": True,
    "cache": "disk",

    # Reproducibility
    "seed": SEED,
    "deterministic": True,

    # Optimizer
    #
    # AdamW is stable for small fine-tuning datasets.
    # lr0 is configured separately for each stage.
    "optimizer": "AdamW",
    "momentum": 0.90,
    "weight_decay": 0.00027,
    "cos_lr": True,
    "lrf": 0.05,
    "warmup_epochs": 3.0,
    "warmup_momentum": 0.80,
    "warmup_bias_lr": 0.01,

    # Official YOLO26s loss gains
    "box": 9.83,
    "cls": 0.65,
    "dfl": 0.96,

    # Class imbalance handling
    #
    # 0.0 = disabled
    # 0.5 = square-root inverse frequency, recommended starting point here
    # 1.0 = full inverse frequency, likely too aggressive for your imbalance
    "cls_pw": 0.50,

    # Geometry augmentation
    #
    # Sperm may appear in any orientation, so arbitrary rotation is useful.
    # Translation and scaling are kept moderate to protect small objects.
    "degrees": 180.0,
    "translate": 0.08,
    "scale": 0.30,
    "shear": 0.0,
    "perspective": 0.0,
    "flipud": 0.50,
    "fliplr": 0.50,

    # Microscope color/intensity augmentation
    "hsv_h": 0.01,
    "hsv_s": 0.20,
    "hsv_v": 0.20,
    "bgr": 0.0,

    # Composition augmentations
    #
    # Moderate mosaic helps small datasets, but heavy mosaic can make already
    # small sperm objects too tiny.
    "mosaic": 0.40,
    "mixup": 0.0,
    "cutmix": 0.0,
    "copy_paste": 0.0,

    # Custom Albumentations
    "augmentations": MICROSCOPY_AUGMENTATIONS,

    # Preserve consistent high resolution for small-object detection.
    "multi_scale": 0.0,

    # Validation and output
    "val": True,
    "plots": True,
    "save": True,
    "save_period": 10,
    "max_det": 500,
    "verbose": True,

    # Output directory
    "project": str(OUTPUT_PROJECT),
    "exist_ok": True,
}


def check_environment() -> None:
    """Validate important files and GPU availability before training."""

    if not DATA_YAML.exists():
        raise FileNotFoundError(
            f"Dataset YAML was not found:\n{DATA_YAML}\n\n"
            "Save custom_dataset.yaml at this location or change DATA_YAML."
        )

    if not torch.cuda.is_available():
        raise RuntimeError(
            "A CUDA GPU was not detected. Training YOLO26s at 1024 pixels "
            "on CPU would be extremely slow."
        )

    gpu_name = torch.cuda.get_device_name(DEVICE)

    print("=" * 80)
    print("YOLO26s sperm detection training")
    print(f"Dataset:   {DATA_YAML}")
    print(f"GPU:       {gpu_name}")
    print(f"Resolution:{IMAGE_SIZE}")
    print(f"Outputs:   {OUTPUT_PROJECT}")
    print("=" * 80)


def train_stage1() -> Path:
    """
    Stage 1: train the neck and detection heads while freezing most of the
    backbone. This lets the new six-class head stabilize before full training.
    """

    print("\nSTAGE 1: training detection head and neck\n")

    model = YOLO(INITIAL_MODEL)

    model.train(
        **COMMON_TRAIN_SETTINGS,

        name=STAGE1_NAME,

        # Short head-adaptation stage
        epochs=30,
        patience=30,

        # Freeze layers 0-9, covering most of the YOLO26 backbone.
        freeze=10,

        # Stable initial fine-tuning learning rate
        lr0=0.001,

        # Disable mosaic during the final five stage-1 epochs.
        close_mosaic=5,
    )

    stage1_best = (
        OUTPUT_PROJECT
        / STAGE1_NAME
        / "weights"
        / "best.pt"
    )

    if not stage1_best.exists():
        raise FileNotFoundError(
            f"Stage-1 best checkpoint was not created:\n{stage1_best}"
        )

    return stage1_best


def train_stage2(stage1_best: Path) -> Path:
    """
    Stage 2: unfreeze the full network and adapt the pretrained visual features
    to microscopy using a substantially lower learning rate.
    """

    print("\nSTAGE 2: full-network microscopy fine-tuning\n")

    model = YOLO(str(stage1_best))

    model.train(
        **COMMON_TRAIN_SETTINGS,

        name=STAGE2_NAME,

        # Maximum stage-2 duration. Early stopping normally ends sooner.
        epochs=170,
        patience=40,

        # Unfreeze the complete network.
        freeze=0,

        # Lower learning rate protects the useful stage-1 representation.
        lr0=0.0003,

        # Finish with 20 clean epochs without mosaic.
        close_mosaic=20,
    )

    stage2_best = (
        OUTPUT_PROJECT
        / STAGE2_NAME
        / "weights"
        / "best.pt"
    )

    if not stage2_best.exists():
        raise FileNotFoundError(
            f"Stage-2 best checkpoint was not created:\n{stage2_best}"
        )

    return stage2_best


def validate_best_model(best_checkpoint: Path) -> None:
    """Run final high-resolution validation and print per-class mAP."""

    print("\nFINAL VALIDATION\n")

    model = YOLO(str(best_checkpoint))

    metrics = model.val(
        data=str(DATA_YAML),
        split="val",
        imgsz=IMAGE_SIZE,
        batch=4,
        device=DEVICE,
        max_det=500,
        plots=True,
        verbose=True,
    )

    print("\n" + "=" * 80)
    print(f"Best model: {best_checkpoint}")
    print(f"Overall mAP50:    {metrics.box.map50:.6f}")
    print(f"Overall mAP50-95: {metrics.box.map:.6f}")
    print("\nPer-class mAP50-95:")

    class_names = model.names
    class_maps = metrics.box.maps

    for class_id, class_map in enumerate(class_maps):
        class_name = class_names[class_id]
        print(f"  {class_id}: {class_name:<12} {class_map:.6f}")

    print("=" * 80)


def main() -> None:
    check_environment()

    stage1_best = train_stage1()
    stage2_best = train_stage2(stage1_best)

    validate_best_model(stage2_best)


if __name__ == "__main__":
    # Required on Windows when DataLoader workers are enabled.
    main()

Ultralytics supports supplying custom Albumentations transformations through the Python training API. (Ultralytics Docs) copy_paste is intentionally disabled here because the current configuration defines it for segmentation copy-paste, while your dataset uses detection boxes. (GitHub)

Recommended custom_dataset.yaml

Save this as:

text
C:\Users\BIOCELIZE003\Desktop\besho_detection\Besho_train\custom_dataset.yaml

Use forward slashes in YAML:

yaml
path: C:/Users/BIOCELIZE003/Desktop/besho_detection/Besho_train/dataset

train: train
val: val

names:
  0: sperm
  1: c1
  2: c2
  3: c3
  4: c4
  5: c5

This assumes each image and its matching .txt annotation are together inside dataset/train and dataset/val.

For the standard separated structure:

text
dataset/
├── images/
│   ├── train/
│   └── val/
└── labels/
    ├── train/
    └── val/

use:

yaml
path: C:/Users/BIOCELIZE003/Desktop/besho_detection/Besho_train/dataset

train: images/train
val: images/val

names:
  0: sperm
  1: c1
  2: c2
  3: c3
  4: c4
  5: c5

Critical dataset considerations

Your validation set must contain meaningful numbers of instances of every class, not merely images. If c4 or c5 appears only once or twice in validation, its AP will be unstable and selecting the “best” checkpoint becomes unreliable.

Also split images by biological sample, slide, donor, recording, or microscope session—not by randomly separating neighboring frames. Near-duplicate frames in training and validation can produce unrealistically high metrics.

The final model will be here:

text
C:\Users\BIOCELIZE003\Desktop\besho_detection\Besho_train\runs_sperm_yolo26s\stage2_full_finetuning\weights\best.pt

Yes—1280 is appropriate for your 1280×720 microscopy images, especially because sperm may occupy relatively few pixels. At imgsz=1024, each image is effectively reduced to roughly 1024×576 before padding; at 1280, the original horizontal detail is preserved. Ultralytics notes that larger training sizes may improve small-object accuracy, with increased memory and computation cost. (Ultralytics Docs)

The 30 epochs are only Stage 1

Your log:

text
22/30 ... 1024

is the frozen-backbone stage. The script is configured for:

text
Stage 1: 30 epochs, backbone frozen
Stage 2: up to 170 epochs, everything unfrozen
Total maximum: 200 epochs

So no, the complete training is not only 30 epochs. The official Ultralytics fine-tuning guide also demonstrates frozen training followed by lower-learning-rate full-network training. (Ultralytics Docs)

Your current Stage 1 is already at epoch 22, so do not restart it. Let the 30 epochs finish.

Best approach for your current run

Use:

text
Stage 1: 1024 resolution, 30 epochs
Stage 2: 1280 resolution, up to 170 epochs

This is a good progressive-resizing strategy: Stage 1 learns the new classification head inexpensively, while Stage 2 performs the important full-network microscopy adaptation at maximum resolution.

However, changing the Python file while it is running will not alter the current process. Your current script will automatically start Stage 2 at 1024 because IMAGE_SIZE=1024 was already loaded.

After Stage 1 finishes, stop the program before Stage 2 starts and run this dedicated Stage 2 script.

train_stage2_1280.py

python
from pathlib import Path

import albumentations as A
import torch
from ultralytics import YOLO


ROOT = Path(
    r"C:\Users\BIOCELIZE003\Desktop\besho_detection\Besho_train"
)

DATA_YAML = ROOT / "custom_dataset.yaml"

STAGE1_BEST = (
    ROOT
    / "runs_sperm_yolo26s"
    / "stage1_frozen_backbone"
    / "weights"
    / "best.pt"
)

OUTPUT_PROJECT = ROOT / "runs_sperm_yolo26s"

DEVICE = 0
IMAGE_SIZE = 1280

# Good starting point for 1280 training.
# Batch 8 is usually preferable to allowing auto-batch to fall to batch 2.
BATCH_SIZE = 8


MICROSCOPY_AUGMENTATIONS = [
    A.CLAHE(
        clip_limit=2.0,
        tile_grid_size=(8, 8),
        p=0.20,
    ),
    A.Blur(
        blur_limit=3,
        p=0.10,
    ),
    A.RandomBrightnessContrast(
        brightness_limit=0.10,
        contrast_limit=0.15,
        p=0.15,
    ),
]


def main() -> None:
    if not DATA_YAML.exists():
        raise FileNotFoundError(f"Dataset YAML not found: {DATA_YAML}")

    if not STAGE1_BEST.exists():
        raise FileNotFoundError(f"Stage-1 checkpoint not found: {STAGE1_BEST}")

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA GPU was not detected.")

    print(f"GPU: {torch.cuda.get_device_name(DEVICE)}")
    print(f"Starting checkpoint: {STAGE1_BEST}")
    print(f"Training resolution: {IMAGE_SIZE}")
    print(f"Batch size: {BATCH_SIZE}")

    model = YOLO(str(STAGE1_BEST))

    model.train(
        data=str(DATA_YAML),

        project=str(OUTPUT_PROJECT),
        name="stage2_full_finetuning_1280",
        exist_ok=True,

        # Full-resolution fine-tuning
        imgsz=IMAGE_SIZE,
        batch=BATCH_SIZE,

        # Full-network training
        epochs=170,
        patience=40,
        freeze=0,

        # Hardware
        device=DEVICE,
        workers=4,
        amp=True,
        cache="disk",

        # Reproducibility
        seed=42,
        deterministic=True,

        # Optimizer
        optimizer="AdamW",
        lr0=0.0003,
        lrf=0.05,
        momentum=0.90,
        weight_decay=0.00027,
        cos_lr=True,

        warmup_epochs=3.0,
        warmup_momentum=0.80,
        warmup_bias_lr=0.01,

        # YOLO26s loss gains
        box=9.83,
        cls=0.65,
        dfl=0.96,

        # Imbalanced class weighting
        cls_pw=0.50,

        # Geometry augmentation
        degrees=180.0,
        translate=0.08,
        scale=0.30,
        shear=0.0,
        perspective=0.0,
        flipud=0.50,
        fliplr=0.50,

        # Mild microscopy color augmentation
        hsv_h=0.01,
        hsv_s=0.20,
        hsv_v=0.20,
        bgr=0.0,

        # Composition augmentation
        mosaic=0.40,
        close_mosaic=20,
        mixup=0.0,
        cutmix=0.0,
        copy_paste=0.0,

        augmentations=MICROSCOPY_AUGMENTATIONS,

        multi_scale=0.0,
        rect=False,

        # Validation and saving
        val=True,
        plots=True,
        save=True,
        save_period=10,
        max_det=500,
        verbose=True,
    )

    best_model_path = (
        OUTPUT_PROJECT
        / "stage2_full_finetuning_1280"
        / "weights"
        / "best.pt"
    )

    print(f"\nBest model should be saved at:\n{best_model_path}")


if __name__ == "__main__":
    main()

About batch size

Your log shows 138 batches for approximately 600 images, suggesting that the current effective batch size may be around four. At 1280, automatic batch selection could lower it further. Ultralytics supports fixed batches or automatic GPU-memory-based batch selection. (Ultralytics Docs)

Start with:

python
BATCH_SIZE = 8

For a GPU with about 24 GB or more, test:

python
BATCH_SIZE = 16

Batch 8 is already sufficient. A larger batch does not automatically mean higher accuracy.

Are the displayed losses good?

These values are not obviously problematic:

text
box_loss ≈ 1.87
cls_loss ≈ 0.43
dfl_loss ≈ 0.0046

They appear stable, but training losses cannot tell whether your model is accurate. The important output is the validation row after each epoch:

text
Precision
Recall
mAP50
mAP50-95

Most importantly, inspect the per-class results for:

text
c1, c2, c3, c4, c5

Because sperm class 0 dominates the dataset, a model can obtain good-looking overall precision while still performing poorly on rare classes. Use the final best.pt, not last.pt; early stopping may finish Stage 2 before epoch 170 if validation performance stops improving. Ultralytics uses patience to control how long training continues without validation improvement. (Ultralytics Docs)

The error occurs because your current train.py is the Stage 2-only script, so it expects an existing Stage 1 checkpoint.

Replace the entire train.py with the script below. It starts a completely new run from yolo26s.pt, creates its own Stage 1 checkpoint, and automatically proceeds to full Stage 2 training at 1280 resolution.

This is a fresh training run, but it uses COCO-pretrained yolo26s.pt weights. With only about 600 training images, this is substantially safer than true random initialization with yolo26s.yaml; random initialization generally needs much more data. (Ultralytics Docs)

Install/update first:

bat
pip install -U ultralytics albumentations pyyaml

Complete replacement train.py

python
from __future__ import annotations

import math
import random
from pathlib import Path

import albumentations as A
import torch
import yaml

from ultralytics import YOLO
from ultralytics.data.utils import IMG_FORMATS, img2label_paths


# =============================================================================
# FIXED CONFIGURATION
# =============================================================================

ROOT = Path(
    r"C:\Users\BIOCELIZE003\Desktop\besho_detection\Besho_train"
)

DATA_YAML = ROOT / "custom_dataset.yaml"

# Starts a new fine-tuning run from official pretrained weights.
PRETRAINED_MODEL = "yolo26s.pt"

PROJECT_DIR = ROOT / "runs_sperm_yolo26s"

BALANCED_TRAIN_LIST = ROOT / "balanced_train_images.txt"
BALANCED_DATA_YAML = ROOT / "custom_dataset_balanced.yaml"

NUMBER_OF_CLASSES = 6

IMAGE_SIZE = 1280

# Automatic batch selection using approximately 80% GPU memory.
# Ultralytics can reduce it automatically if the first epoch runs out of memory.
BATCH_SIZE = 0.80

DEVICE = 0
WORKERS = 4
SEED = 42

# Damped oversampling prevents extremely rare classes from being duplicated
# an excessive number of times.
MAX_IMAGE_REPETITIONS = 4


# =============================================================================
# MICROSCOPY AUGMENTATIONS
# =============================================================================

MICROSCOPY_AUGMENTATIONS = [
    # Improve visibility under varying microscope illumination.
    A.CLAHE(
        clip_limit=2.0,
        tile_grid_size=(8, 8),
        p=0.20,
    ),

    # Mild microscope defocus. Keep probability and strength low because
    # strong blur may erase differences between c1-c5.
    A.Blur(
        blur_limit=3,
        p=0.08,
    ),

    # Mild illumination and contrast variation.
    A.RandomBrightnessContrast(
        brightness_limit=0.10,
        contrast_limit=0.15,
        p=0.20,
    ),
]


# =============================================================================
# DATASET BALANCING
# =============================================================================

def resolve_dataset_paths(
    config: dict,
    yaml_path: Path,
    split_name: str,
) -> list[Path]:
    """Resolve train or val paths from a YOLO dataset YAML."""

    dataset_root = config.get("path", yaml_path.parent)
    dataset_root = Path(dataset_root)

    if not dataset_root.is_absolute():
        dataset_root = yaml_path.parent / dataset_root

    split_value = config[split_name]

    if not isinstance(split_value, list):
        split_value = [split_value]

    resolved_paths: list[Path] = []

    for value in split_value:
        path = Path(value)

        if not path.is_absolute():
            path = dataset_root / path

        resolved_paths.append(path.resolve())

    return resolved_paths


def collect_images(dataset_paths: list[Path]) -> list[Path]:
    """Collect image paths from directories or YOLO image-list files."""

    images: list[Path] = []

    for dataset_path in dataset_paths:
        if dataset_path.is_dir():
            for file_path in dataset_path.rglob("*"):
                if (
                    file_path.is_file()
                    and file_path.suffix.lower().lstrip(".") in IMG_FORMATS
                ):
                    images.append(file_path.resolve())

        elif dataset_path.is_file():
            # Support train.txt-style YOLO image lists.
            for line in dataset_path.read_text(
                encoding="utf-8"
            ).splitlines():
                line = line.strip()

                if not line:
                    continue

                image_path = Path(line)

                if not image_path.is_absolute():
                    image_path = dataset_path.parent / image_path

                if (
                    image_path.suffix.lower().lstrip(".")
                    in IMG_FORMATS
                ):
                    images.append(image_path.resolve())

        else:
            raise FileNotFoundError(
                f"Training image path does not exist:\n{dataset_path}"
            )

    images = sorted(images)

    if not images:
        raise RuntimeError(
            "No training images were found. Check the train path in "
            f"{DATA_YAML}"
        )

    return images


def read_image_classes(image_path: Path) -> list[int]:
    """Return all valid class IDs found in one image label file."""

    # Use Ultralytics' own image-to-label path conversion.
    label_path = Path(
        img2label_paths([str(image_path)])[0]
    )

    if not label_path.exists():
        # This may be a valid background image.
        return []

    classes: list[int] = []

    for line_number, line in enumerate(
        label_path.read_text(encoding="utf-8").splitlines(),
        start=1,
    ):
        line = line.strip()

        if not line:
            continue

        parts = line.split()

        if len(parts) < 5:
            raise ValueError(
                f"Invalid annotation in:\n{label_path}\n"
                f"Line {line_number}: {line}"
            )

        try:
            class_id = int(float(parts[0]))
        except ValueError as error:
            raise ValueError(
                f"Invalid class ID in:\n{label_path}\n"
                f"Line {line_number}: {line}"
            ) from error

        if not 0 <= class_id < NUMBER_OF_CLASSES:
            raise ValueError(
                f"Class ID {class_id} is outside the expected range "
                f"0-{NUMBER_OF_CLASSES - 1} in:\n{label_path}"
            )

        classes.append(class_id)

    return classes


def create_balanced_dataset_yaml() -> Path:
    """
    Create a training image list that repeats images containing rare classes.

    Repetition uses square-root inverse frequency and is capped. This improves
    rare-class exposure without allowing a tiny class to dominate training.
    """

    config = yaml.safe_load(
        DATA_YAML.read_text(encoding="utf-8")
    )

    train_paths = resolve_dataset_paths(
        config=config,
        yaml_path=DATA_YAML,
        split_name="train",
    )

    images = collect_images(train_paths)

    class_counts = [0] * NUMBER_OF_CLASSES
    image_class_sets: dict[Path, set[int]] = {}

    for image_path in images:
        image_classes = read_image_classes(image_path)
        image_class_sets[image_path] = set(image_classes)

        for class_id in image_classes:
            class_counts[class_id] += 1

    total_instances = sum(class_counts)

    if total_instances == 0:
        raise RuntimeError(
            "No valid YOLO annotations were found.\n\n"
            "The standard structure should normally be:\n"
            "dataset/images/train\n"
            "dataset/images/val\n"
            "dataset/labels/train\n"
            "dataset/labels/val"
        )

    missing_classes = [
        class_id
        for class_id, count in enumerate(class_counts)
        if count == 0
    ]

    if missing_classes:
        raise RuntimeError(
            "These classes have zero training instances: "
            f"{missing_classes}. All six classes must have annotations."
        )

    majority_count = max(class_counts)

    repetition_factors: list[int] = []

    for count in class_counts:
        inverse_frequency = majority_count / count

        # Square-root inverse frequency is less aggressive than full inverse
        # frequency and is safer for a small dataset.
        factor = int(round(math.sqrt(inverse_frequency)))
        factor = max(1, min(factor, MAX_IMAGE_REPETITIONS))

        repetition_factors.append(factor)

    balanced_images: list[str] = []

    for image_path in images:
        classes_in_image = image_class_sets[image_path]

        if classes_in_image:
            repeat_count = max(
                repetition_factors[class_id]
                for class_id in classes_in_image
            )
        else:
            # Preserve background images once.
            repeat_count = 1

        balanced_images.extend(
            [image_path.as_posix()] * repeat_count
        )

    random_generator = random.Random(SEED)
    random_generator.shuffle(balanced_images)

    BALANCED_TRAIN_LIST.write_text(
        "\n".join(balanced_images) + "\n",
        encoding="utf-8",
    )

    balanced_config = dict(config)
    balanced_config["train"] = BALANCED_TRAIN_LIST.as_posix()

    BALANCED_DATA_YAML.write_text(
        yaml.safe_dump(
            balanced_config,
            sort_keys=False,
            allow_unicode=True,
        ),
        encoding="utf-8",
    )

    class_names = config.get(
        "names",
        {index: str(index) for index in range(NUMBER_OF_CLASSES)},
    )

    print("\n" + "=" * 80)
    print("TRAINING DATASET STATISTICS")
    print("=" * 80)
    print(f"Original training images: {len(images)}")
    print(f"Balanced training entries: {len(balanced_images)}")
    print(f"Total annotated objects: {total_instances}")
    print()

    for class_id in range(NUMBER_OF_CLASSES):
        if isinstance(class_names, dict):
            class_name = class_names.get(
                class_id,
                class_names.get(str(class_id), str(class_id)),
            )
        else:
            class_name = class_names[class_id]

        print(
            f"Class {class_id} ({class_name}): "
            f"{class_counts[class_id]} objects, "
            f"image repetition factor up to "
            f"{repetition_factors[class_id]}x"
        )

    print("=" * 80 + "\n")

    return BALANCED_DATA_YAML


# =============================================================================
# TRAINING
# =============================================================================

def get_best_checkpoint(model: YOLO) -> Path:
    """Return the best checkpoint generated by the latest training run."""

    trainer = model.trainer

    if trainer is None:
        raise RuntimeError("Ultralytics trainer was not created.")

    best_path = getattr(trainer, "best", None)

    if best_path:
        best_path = Path(best_path)
    else:
        best_path = Path(trainer.save_dir) / "weights" / "best.pt"

    if not best_path.exists():
        raise FileNotFoundError(
            f"Best checkpoint was not created:\n{best_path}"
        )

    return best_path


def common_training_settings(
    balanced_yaml: Path,
) -> dict:
    """Shared settings for both training stages."""

    return {
        "data": str(balanced_yaml),

        # Full original image width.
        "imgsz": IMAGE_SIZE,

        # Automatically uses about 80% GPU memory.
        "batch": BATCH_SIZE,

        # Hardware
        "device": DEVICE,
        "workers": WORKERS,
        "amp": True,
        "channels_last": True,
        "cache": "disk",

        # Reproducibility
        "seed": SEED,
        "deterministic": True,

        # Optimizer
        "optimizer": "AdamW",
        "momentum": 0.90,
        "weight_decay": 0.00027,
        "cos_lr": True,
        "lrf": 0.01,
        "warmup_epochs": 3.0,
        "warmup_momentum": 0.80,
        "warmup_bias_lr": 0.01,

        # Official YOLO26s loss gains
        "box": 9.83,
        "cls": 0.65,
        "dfl": 0.96,

        # Sperm may appear at arbitrary orientations.
        "degrees": 180.0,
        "translate": 0.05,
        "scale": 0.20,
        "shear": 0.0,
        "perspective": 0.0,
        "flipud": 0.50,
        "fliplr": 0.50,

        # Mild color variation only. Strong HSV changes can destroy subtle
        # morphology or staining information.
        "hsv_h": 0.005,
        "hsv_s": 0.10,
        "hsv_v": 0.15,
        "bgr": 0.0,

        # Moderate mosaic for a dataset under 1,000 images.
        "mosaic": 0.35,
        "mixup": 0.0,
        "cutmix": 0.0,
        "copy_paste": 0.0,

        # Custom microscopy transforms
        "augmentations": MICROSCOPY_AUGMENTATIONS,

        # Keep a consistent 1280 resolution.
        "multi_scale": 0.0,
        "rect": False,

        # Output and validation
        "project": str(PROJECT_DIR),
        "exist_ok": False,
        "save": True,
        "save_period": 10,
        "val": True,
        "plots": True,
        "max_det": 500,
        "verbose": True,
    }


def train_stage1(balanced_yaml: Path) -> Path:
    """
    Stage 1: initialize from yolo26s.pt and stabilize the custom detection head.
    """

    print("\n" + "=" * 80)
    print("STAGE 1: FROZEN-BACKBONE TRAINING")
    print("=" * 80)

    model = YOLO(PRETRAINED_MODEL)

    settings = common_training_settings(balanced_yaml)

    model.train(
        **settings,
        name="stage1_frozen_1280",

        # This is only the head-adaptation stage.
        epochs=20,
        patience=20,

        # Freeze most of the YOLO26 backbone.
        freeze=10,

        # Stable learning rate for a small dataset.
        lr0=0.001,

        # Remove mosaic near the end of Stage 1.
        close_mosaic=5,
    )

    best_path = get_best_checkpoint(model)

    print(f"\nStage 1 best checkpoint:\n{best_path}\n")

    return best_path


def train_stage2(
    stage1_checkpoint: Path,
    balanced_yaml: Path,
) -> Path:
    """
    Stage 2: unfreeze the complete network and adapt all features to microscopy.
    """

    print("\n" + "=" * 80)
    print("STAGE 2: FULL-NETWORK FINE-TUNING")
    print("=" * 80)

    model = YOLO(str(stage1_checkpoint))

    settings = common_training_settings(balanced_yaml)

    model.train(
        **settings,
        name="stage2_full_1280",

        # This is the main training stage.
        epochs=180,

        # Early stopping protects against overfitting.
        patience=45,

        # Train the complete backbone, neck and detection head.
        freeze=0,

        # Lower LR after head stabilization.
        lr0=0.0003,

        # Final 20 epochs use normal images without mosaic.
        close_mosaic=20,
    )

    best_path = get_best_checkpoint(model)

    print(f"\nStage 2 best checkpoint:\n{best_path}\n")

    return best_path


def final_validation(best_checkpoint: Path) -> None:
    """Validate the final checkpoint on the original validation split."""

    print("\n" + "=" * 80)
    print("FINAL VALIDATION")
    print("=" * 80)

    model = YOLO(str(best_checkpoint))

    metrics = model.val(
        data=str(DATA_YAML),
        split="val",
        imgsz=IMAGE_SIZE,
        batch=8,
        device=DEVICE,
        workers=WORKERS,
        max_det=500,
        plots=True,
        verbose=True,
    )

    print("\n" + "=" * 80)
    print(f"FINAL BEST MODEL:\n{best_checkpoint}")
    print(f"\nmAP50:    {metrics.box.map50:.6f}")
    print(f"mAP50-95: {metrics.box.map:.6f}")

    if hasattr(metrics.box, "maps"):
        print("\nPer-class mAP50-95:")

        for class_id, class_map in enumerate(metrics.box.maps):
            class_name = model.names[class_id]

            print(
                f"  {class_id}: "
                f"{class_name:<12} "
                f"{float(class_map):.6f}"
            )

    print("=" * 80)


def main() -> None:
    if not DATA_YAML.exists():
        raise FileNotFoundError(
            f"Dataset YAML was not found:\n{DATA_YAML}"
        )

    if not torch.cuda.is_available():
        raise RuntimeError(
            "CUDA GPU was not detected. Check your PyTorch CUDA installation."
        )

    print("\n" + "=" * 80)
    print("YOLO26S SPERM DETECTION TRAINING")
    print("=" * 80)
    print(f"GPU: {torch.cuda.get_device_name(DEVICE)}")
    print(f"Resolution: {IMAGE_SIZE}")
    print(f"Dataset YAML: {DATA_YAML}")
    print(f"Output project: {PROJECT_DIR}")
    print("=" * 80)

    balanced_yaml = create_balanced_dataset_yaml()

    stage1_best = train_stage1(balanced_yaml)

    stage2_best = train_stage2(
        stage1_checkpoint=stage1_best,
        balanced_yaml=balanced_yaml,
    )

    final_validation(stage2_best)


if __name__ == "__main__":
    # Required for Windows multiprocessing DataLoader workers.
    main()

Run it normally:

bat
cd C:\Users\BIOCELIZE003\Desktop\besho_detection\Besho_train
python train.py

What this run will do

It performs:

text
Stage 1: 20 epochs at 1280, backbone mostly frozen
Stage 2: up to 180 epochs at 1280, full model unfrozen
Maximum total: 200 epochs

Stage 2 has patience=45, so it can stop before 180 epochs when validation accuracy no longer improves. Epoch count should be selected using validation behavior rather than assuming every run needs the full maximum. (Ultralytics Docs)

The batch=0.80 setting instructs Ultralytics to select a batch using approximately 80% of GPU memory. On a single GPU, current Ultralytics can also retry with a smaller batch after an out-of-memory error during the first epoch. (Ultralytics Docs)

The script creates:

text
balanced_train_images.txt
custom_dataset_balanced.yaml

Rare-class images are repeated up to four times, while validation remains unchanged. The repetitions receive different random augmentations during training.

I also removed cls_pw. It is not a normal public Ultralytics training setting. Current official documentation says structural class-weighted loss requires a custom loss/model/trainer implementation; passing unsupported internal weighting arguments can raise an invalid-argument error. (Ultralytics Docs)

Your final model will be inside the directory printed after training, under:

text
runs_sperm_yolo26s\stage2_full_1280*\weights\best.pt

The crash is caused by this unsupported line:

python
"channels_last": True,

Your installed Ultralytics version predates that setting. The replacement below removes it and also checks every training option against the configuration supported by your installed version, preventing another invalid-argument crash. Current Ultralytics configurations include options that older releases may not recognize, including recently added settings such as channels_last and cls_pw. (GitHub)

Replace the entire contents of train.py with this:

python
from __future__ import annotations

import math
import random
import shutil
from pathlib import Path

import torch
import yaml
from ultralytics import YOLO
from ultralytics.data.utils import IMG_FORMATS, img2label_paths

try:
    # Allows the script to detect which arguments your installed
    # Ultralytics version actually supports.
    from ultralytics.cfg import DEFAULT_CFG_DICT
except ImportError:
    DEFAULT_CFG_DICT = {}


# =============================================================================
# FIXED CONFIGURATION
# =============================================================================

ROOT = Path(
    r"C:\Users\BIOCELIZE003\Desktop\besho_detection\Besho_train"
)

ORIGINAL_DATA_YAML = ROOT / "custom_dataset.yaml"

BALANCED_IMAGE_LIST = ROOT / "balanced_train_images.txt"
BALANCED_DATA_YAML = ROOT / "custom_dataset_balanced.yaml"

PROJECT_DIR = ROOT / "runs_sperm_yolo26s"

MODEL_FILE = "yolo26s.pt"

DEVICE = 0
IMAGE_SIZE = 1280

# RTX 3060 12 GB should normally handle YOLO26s at 1280 with batch 8.
# Change this to 4 only if CUDA reports an out-of-memory error.
BATCH_SIZE = 8

WORKERS = 4
SEED = 42

NUMBER_OF_CLASSES = 6
MAX_IMAGE_REPETITIONS = 4

STAGE1_NAME = "stage1_frozen_1280"
STAGE2_NAME = "stage2_full_1280"

# Delete only these two run directories before starting.
START_COMPLETELY_FRESH = True


# =============================================================================
# ULTRALYTICS VERSION COMPATIBILITY
# =============================================================================

def supported_arguments(arguments: dict) -> dict:
    """
    Remove training arguments that the installed Ultralytics version does not
    support. This prevents errors such as:
        'channels_last' is not a valid YOLO argument
    """

    if not DEFAULT_CFG_DICT:
        return arguments

    allowed_keys = set(DEFAULT_CFG_DICT.keys())

    accepted = {}
    skipped = {}

    for key, value in arguments.items():
        if key in allowed_keys:
            accepted[key] = value
        else:
            skipped[key] = value

    if skipped:
        print("\nArguments skipped because this Ultralytics version")
        print("does not support them:")

        for key in skipped:
            print(f"  - {key}")

        print()

    return accepted


def run_training(model: YOLO, arguments: dict):
    """Train using only arguments supported by this installation."""

    arguments = supported_arguments(arguments)
    return model.train(**arguments)


# =============================================================================
# DATASET BALANCING
# =============================================================================

def resolve_split_paths(
    dataset_config: dict,
    yaml_path: Path,
    split_name: str,
) -> list[Path]:
    """Resolve train or validation paths from the dataset YAML."""

    dataset_root_value = dataset_config.get("path", yaml_path.parent)
    dataset_root = Path(dataset_root_value)

    if not dataset_root.is_absolute():
        dataset_root = yaml_path.parent / dataset_root

    split_value = dataset_config[split_name]

    if not isinstance(split_value, list):
        split_value = [split_value]

    resolved = []

    for item in split_value:
        item_path = Path(item)

        if not item_path.is_absolute():
            item_path = dataset_root / item_path

        resolved.append(item_path.resolve())

    return resolved


def collect_images(paths: list[Path]) -> list[Path]:
    """Find all training images from directories or image-list files."""

    images = []

    for path in paths:
        if path.is_dir():
            for file_path in path.rglob("*"):
                extension = file_path.suffix.lower().lstrip(".")

                if file_path.is_file() and extension in IMG_FORMATS:
                    images.append(file_path.resolve())

        elif path.is_file():
            for line in path.read_text(
                encoding="utf-8"
            ).splitlines():
                line = line.strip()

                if not line:
                    continue

                image_path = Path(line)

                if not image_path.is_absolute():
                    image_path = path.parent / image_path

                extension = image_path.suffix.lower().lstrip(".")

                if extension in IMG_FORMATS:
                    images.append(image_path.resolve())

        else:
            raise FileNotFoundError(
                f"Dataset path does not exist:\n{path}"
            )

    images = sorted(images)

    if not images:
        raise RuntimeError(
            "No training images were found. Check the train path in:\n"
            f"{ORIGINAL_DATA_YAML}"
        )

    return images


def read_classes(image_path: Path) -> list[int]:
    """Read class IDs from an image's YOLO annotation file."""

    label_path = Path(
        img2label_paths([str(image_path)])[0]
    )

    # Valid background image.
    if not label_path.exists():
        return []

    class_ids = []

    lines = label_path.read_text(
        encoding="utf-8"
    ).splitlines()

    for line_number, line in enumerate(lines, start=1):
        line = line.strip()

        if not line:
            continue

        values = line.split()

        if len(values) < 5:
            raise ValueError(
                f"Invalid YOLO annotation:\n"
                f"{label_path}\n"
                f"Line {line_number}: {line}"
            )

        try:
            class_id = int(float(values[0]))
        except ValueError as error:
            raise ValueError(
                f"Invalid class ID:\n"
                f"{label_path}\n"
                f"Line {line_number}: {line}"
            ) from error

        if not 0 <= class_id < NUMBER_OF_CLASSES:
            raise ValueError(
                f"Invalid class {class_id} in:\n"
                f"{label_path}\n\n"
                f"Expected classes 0-{NUMBER_OF_CLASSES - 1}."
            )

        class_ids.append(class_id)

    return class_ids


def class_name(
    names: dict | list,
    class_id: int,
) -> str:
    """Return a class name from either YAML names format."""

    if isinstance(names, list):
        return str(names[class_id])

    return str(
        names.get(
            class_id,
            names.get(str(class_id), class_id),
        )
    )


def create_balanced_dataset() -> Path:
    """
    Repeat images containing rare classes.

    Square-root inverse-frequency balancing is used instead of full inverse
    frequency, because class 2 has only 31 objects and could otherwise be
    severely overfitted.
    """

    config = yaml.safe_load(
        ORIGINAL_DATA_YAML.read_text(encoding="utf-8")
    )

    train_paths = resolve_split_paths(
        dataset_config=config,
        yaml_path=ORIGINAL_DATA_YAML,
        split_name="train",
    )

    images = collect_images(train_paths)

    class_counts = [0] * NUMBER_OF_CLASSES
    image_classes: dict[Path, set[int]] = {}

    background_images = 0

    for image_path in images:
        classes = read_classes(image_path)

        unique_classes = set(classes)
        image_classes[image_path] = unique_classes

        if not classes:
            background_images += 1

        for class_id in classes:
            class_counts[class_id] += 1

    total_objects = sum(class_counts)

    if total_objects == 0:
        raise RuntimeError(
            "No annotated objects were found."
        )

    missing_classes = [
        class_id
        for class_id, count in enumerate(class_counts)
        if count == 0
    ]

    if missing_classes:
        raise RuntimeError(
            f"Classes with zero annotations: {missing_classes}"
        )

    majority_count = max(class_counts)

    repetition_factors = []

    for count in class_counts:
        inverse_frequency = majority_count / count

        # Damped inverse-frequency balancing.
        factor = round(math.sqrt(inverse_frequency))

        factor = max(
            1,
            min(int(factor), MAX_IMAGE_REPETITIONS),
        )

        repetition_factors.append(factor)

    balanced_entries = []

    for image_path in images:
        classes = image_classes[image_path]

        if classes:
            repeat_count = max(
                repetition_factors[class_id]
                for class_id in classes
            )
        else:
            repeat_count = 1

        balanced_entries.extend(
            [image_path.as_posix()] * repeat_count
        )

    random_generator = random.Random(SEED)
    random_generator.shuffle(balanced_entries)

    BALANCED_IMAGE_LIST.write_text(
        "\n".join(balanced_entries) + "\n",
        encoding="utf-8",
    )

    balanced_config = dict(config)

    # Absolute path to the generated image list.
    balanced_config["train"] = (
        BALANCED_IMAGE_LIST.as_posix()
    )

    BALANCED_DATA_YAML.write_text(
        yaml.safe_dump(
            balanced_config,
            sort_keys=False,
            allow_unicode=True,
        ),
        encoding="utf-8",
    )

    names = config.get(
        "names",
        {
            class_id: str(class_id)
            for class_id in range(NUMBER_OF_CLASSES)
        },
    )

    print("\n" + "=" * 80)
    print("TRAINING DATASET STATISTICS")
    print("=" * 80)

    print(f"Original training images: {len(images)}")
    print(f"Background images: {background_images}")
    print(f"Balanced training entries: {len(balanced_entries)}")
    print(f"Total annotated objects: {total_objects}")
    print()

    for class_id, count in enumerate(class_counts):
        print(
            f"Class {class_id} "
            f"({class_name(names, class_id)}): "
            f"{count} objects, "
            f"maximum repetition "
            f"{repetition_factors[class_id]}x"
        )

    print("=" * 80)

    return BALANCED_DATA_YAML


# =============================================================================
# TRAINING SETTINGS
# =============================================================================

def common_training_arguments(
    data_yaml: Path,
    run_name: str,
) -> dict:
    """Parameters shared by both training stages."""

    arguments = {
        "data": str(data_yaml),
        "project": str(PROJECT_DIR),
        "name": run_name,
        "exist_ok": True,

        # Image and batch configuration
        "imgsz": IMAGE_SIZE,
        "batch": BATCH_SIZE,

        # Hardware
        "device": DEVICE,
        "workers": WORKERS,
        "amp": True,
        "cache": "disk",

        # Repeatability
        "seed": SEED,
        "deterministic": True,

        # Optimizer
        "optimizer": "AdamW",
        "momentum": 0.90,
        "weight_decay": 0.00027,
        "cos_lr": True,
        "lrf": 0.05,

        # Warmup
        "warmup_epochs": 3.0,
        "warmup_momentum": 0.80,
        "warmup_bias_lr": 0.01,

        # YOLO26s loss gains
        "box": 9.83,
        "cls": 0.65,
        "dfl": 0.96,

        # Microscopy geometry augmentation
        "degrees": 180.0,
        "translate": 0.08,
        "scale": 0.25,
        "shear": 0.0,
        "perspective": 0.0,
        "flipud": 0.50,
        "fliplr": 0.50,

        # Mild brightness/color variation
        "hsv_h": 0.005,
        "hsv_s": 0.10,
        "hsv_v": 0.15,

        # Moderate mosaic because sperm objects may be small
        "mosaic": 0.35,
        "mixup": 0.0,
        "copy_paste": 0.0,

        # These are passed only if the installed version supports them
        "bgr": 0.0,
        "cutmix": 0.0,

        # Validation and output
        "rect": False,
        "val": True,
        "plots": True,
        "save": True,
        "save_period": 10,
        "max_det": 500,
        "verbose": True,
    }

    # Newer Ultralytics versions support class-frequency weighting.
    # Older installations will continue using dataset oversampling.
    if "cls_pw" in DEFAULT_CFG_DICT:
        arguments["cls_pw"] = 0.50

    return arguments


def get_best_checkpoint(model: YOLO) -> Path:
    """Find the best checkpoint from a completed training stage."""

    if model.trainer is None:
        raise RuntimeError(
            "Ultralytics trainer was not initialized."
        )

    best_path = Path(model.trainer.best)

    if best_path.exists():
        return best_path

    last_path = Path(model.trainer.last)

    if last_path.exists():
        print(
            "Warning: best.pt was not found. "
            "Using last.pt instead."
        )
        return last_path

    raise FileNotFoundError(
        "No checkpoint was created by training."
    )


# =============================================================================
# STAGE 1
# =============================================================================

def train_stage1(data_yaml: Path) -> Path:
    """
    Stabilize the six-class detection head while the early backbone is frozen.
    This is not the entire training run.
    """

    print("\n" + "=" * 80)
    print("STAGE 1: FROZEN-BACKBONE TRAINING")
    print("20 epochs at 1280")
    print("=" * 80)

    model = YOLO(MODEL_FILE)

    arguments = common_training_arguments(
        data_yaml=data_yaml,
        run_name=STAGE1_NAME,
    )

    arguments.update(
        {
            "epochs": 20,
            "patience": 20,
            "freeze": 10,
            "lr0": 0.001,
            "close_mosaic": 5,
        }
    )

    run_training(model, arguments)

    checkpoint = get_best_checkpoint(model)

    print("\nStage 1 checkpoint:")
    print(checkpoint)

    return checkpoint


# =============================================================================
# STAGE 2
# =============================================================================

def train_stage2(
    stage1_checkpoint: Path,
    data_yaml: Path,
) -> Path:
    """Fine-tune the complete model at a lower learning rate."""

    print("\n" + "=" * 80)
    print("STAGE 2: FULL-NETWORK FINE-TUNING")
    print("Up to 180 epochs at 1280")
    print("=" * 80)

    model = YOLO(str(stage1_checkpoint))

    arguments = common_training_arguments(
        data_yaml=data_yaml,
        run_name=STAGE2_NAME,
    )

    arguments.update(
        {
            "epochs": 180,
            "patience": 45,
            "freeze": 0,
            "lr0": 0.0003,
            "close_mosaic": 20,
        }
    )

    run_training(model, arguments)

    checkpoint = get_best_checkpoint(model)

    print("\nStage 2 checkpoint:")
    print(checkpoint)

    return checkpoint


# =============================================================================
# FINAL VALIDATION
# =============================================================================

def validate_final_model(checkpoint: Path) -> None:
    """Validate against the original, non-oversampled validation set."""

    print("\n" + "=" * 80)
    print("FINAL VALIDATION")
    print("=" * 80)

    model = YOLO(str(checkpoint))

    validation_arguments = {
        "data": str(ORIGINAL_DATA_YAML),
        "split": "val",
        "imgsz": IMAGE_SIZE,
        "batch": BATCH_SIZE,
        "device": DEVICE,
        "workers": WORKERS,
        "max_det": 500,
        "plots": True,
        "verbose": True,
    }

    validation_arguments = supported_arguments(
        validation_arguments
    )

    metrics = model.val(**validation_arguments)

    print("\n" + "=" * 80)
    print("FINAL BEST MODEL")
    print(checkpoint)

    print(f"\nOverall mAP50:    {metrics.box.map50:.6f}")
    print(f"Overall mAP50-95: {metrics.box.map:.6f}")

    if hasattr(metrics.box, "maps"):
        print("\nPer-class mAP50-95:")

        for class_id, class_map in enumerate(
            metrics.box.maps
        ):
            print(
                f"  {class_id}: "
                f"{model.names[class_id]:<12} "
                f"{float(class_map):.6f}"
            )

    print("=" * 80)


# =============================================================================
# MAIN
# =============================================================================

def main() -> None:
    if not ORIGINAL_DATA_YAML.exists():
        raise FileNotFoundError(
            "Dataset YAML was not found:\n"
            f"{ORIGINAL_DATA_YAML}"
        )

    if not torch.cuda.is_available():
        raise RuntimeError(
            "CUDA GPU was not detected. "
            "Check your PyTorch CUDA installation."
        )

    print("\n" + "=" * 80)
    print("YOLO26S SPERM DETECTION TRAINING")
    print("=" * 80)

    print(f"GPU: {torch.cuda.get_device_name(DEVICE)}")
    print(f"Resolution: {IMAGE_SIZE}")
    print(f"Batch size: {BATCH_SIZE}")
    print(f"Dataset: {ORIGINAL_DATA_YAML}")
    print(f"Output: {PROJECT_DIR}")

    print("=" * 80)

    if START_COMPLETELY_FRESH:
        for run_name in [STAGE1_NAME, STAGE2_NAME]:
            run_directory = PROJECT_DIR / run_name

            if run_directory.exists():
                print(
                    f"Removing previous run directory:\n"
                    f"{run_directory}"
                )
                shutil.rmtree(run_directory)

    balanced_yaml = create_balanced_dataset()

    stage1_checkpoint = train_stage1(
        data_yaml=balanced_yaml
    )

    stage2_checkpoint = train_stage2(
        stage1_checkpoint=stage1_checkpoint,
        data_yaml=balanced_yaml,
    )

    validate_final_model(stage2_checkpoint)


if __name__ == "__main__":
    # Required for Windows DataLoader multiprocessing.
    main()

Run:

bat
cd C:\Users\BIOCELIZE003\Desktop\besho_detection\Besho_train
python train.py

The run consists of:

text
Stage 1: 20 frozen-backbone epochs
Stage 2: up to 180 full-model epochs
Maximum: 200 epochs

It will use early stopping during Stage 2, so it may finish before 180 epochs when validation performance stops improving. The native augmentation options used here—rotation, translation, scale, flips, HSV adjustment, Mosaic and MixUp—are standard Ultralytics training settings. (GitHub)

Your final checkpoint will be:

text
C:\Users\BIOCELIZE003\Desktop\besho_detection\Besho_train\
runs_sperm_yolo26s\stage2_full_1280\weights\best.pt

Should CUDA report an out-of-memory error, change only:

python
BATCH_SIZE = 8

to:

python
BATCH_SIZE = 4