Scripts for YOLO Model Fine Tuning

prepare_dataset.py

""" take a dataset of images and labels in the YOLO format and break them into training and validation sets """

import argparse
import shutil
import yaml
from pathlib import Path
from sklearn.model_selection import train_test_split

# parse the cli args
parser = argparse.ArgumentParser(description='Prepare Dataset')
parser.add_argument('-d', '--dir', default="dataset", help="Path to the dataset")

args = parser.parse_args()

# load that paths
source_path = Path(args.dir)
dest_path = source_path.with_stem(source_path.stem + "_yolo")

source_images = source_path / "images"
source_labels = source_path / "labels"

train_images_dir = dest_path / "images" / "train"
val_images_dir = dest_path / "images" / "val"
train_labels_dir = dest_path / "labels" / "train"
val_labels_dir = dest_path / "labels" / "val"

# make sure the directories exist
train_images_dir.mkdir(parents=True, exist_ok=True)
val_images_dir.mkdir(parents=True, exist_ok=True)
train_labels_dir.mkdir(parents=True, exist_ok=True)
val_labels_dir.mkdir(parents=True, exist_ok=True)

# get a list of all the images
image_files = sorted((source_images).iterdir())

# Split the indices into train and validation sets
indices = list(range(len(image_files)))
train_indices, _ = train_test_split(indices, test_size=0.2, random_state=42)
train_files = {image_files[i].stem for i in train_indices}  # Use base filenames for faster lookups

# move the files into the appropriate folders
for image_file in image_files:
    # create the name of the label file
    label_file = source_labels / image_file.with_suffix('.txt').name

    # decide which folder to move them in
    image_move_path = train_images_dir if image_file.stem in train_files else val_images_dir
    label_move_path = train_labels_dir if image_file.stem in train_files else val_labels_dir

    # copy the files
    print(f"Moving {image_file}")
    shutil.copy(image_file, image_move_path / image_file.name)
    shutil.copy(label_file, label_move_path / label_file.name)

# open the class file to get a list of classes
with open(source_path / "classes.txt") as class_file:
    classes_str = class_file.read().strip()

# create class array
classes = [x.strip() for x in classes_str.split("\n")]
print(f"Found {len(classes)} classes")

# create a yaml file describing the structure
print("Writing dataset.yaml file")
yaml_file = Path("dataset.yaml")
yaml_data = {
    "train": str(train_images_dir),
    "val": str(val_images_dir),
    "nc": len(classes),
    "names": classes
}

with yaml_file.open('w') as f:
    yaml.dump(yaml_data, f)

requirements.txt

--extra-index-url  https://download.pytorch.org/whl/cu126
Pillow
scikit-learn
torch
ultralytics-opencv-headless

train.py

""" fine tune a YOLO model using a new dataset """

import argparse
from pathlib import Path
from ultralytics import YOLO

# parse command line arguments
parser = argparse.ArgumentParser(description='Transfer Learning YOLO Model')
parser.add_argument('-m', '--model', default="yolo11s.pt", help="model checkpoint to load")
parser.add_argument('-e', '--epochs', default=10, type=int, help="number of epochs to train for")
parser.add_argument('-d', '--dataset', required=True, type=str, help="dataset path")
parser.add_argument('-f', '--freeze', default=0, type=int, help="numer of layers to freeze")

args = parser.parse_args()

# load the model
model = YOLO(args.model)

# train the model using the dataset information
results = model.train(data=args.dataset, epochs=args.epochs, project="trash", imgsz=960,
                     cos_lr=True, mosaic=0, freeze=args.freeze)

# print the results
print(results.results_dict)