152 lines
5.0 KiB
Python
152 lines
5.0 KiB
Python
"""
|
|
STEP 3 - train.py
|
|
================
|
|
|
|
Fine-tunes YOLOv8-nano on your augmented dataset.
|
|
Pre-trained COCO weights give you a huve head-start - the backbone
|
|
already knows edges, textures, and shapes; it only needs to learn
|
|
what Famico (and other products) look like.
|
|
|
|
Install:
|
|
pip install ultralytics
|
|
|
|
Usage:
|
|
python train.py
|
|
|
|
Outputs:
|
|
runs/detect/coffee_v1/weights/best.pt <- use this for detection
|
|
"""
|
|
|
|
import os
|
|
import torch
|
|
from pathlib import Path
|
|
from ultralytics import YOLO
|
|
|
|
# ── Configuration ───────────────────────────────────────────────────
|
|
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
|
|
|
|
DATASET_YAML = "dataset.yml" # Generated by prepare_dataset.py
|
|
MODEL = "yolov8n.pt" # nano - fast, good for few-shot
|
|
# alternativse to try: yolov8s.pt (better acc)
|
|
# yolov8m.pt (even better, needs GPU)
|
|
|
|
PROJECT = "runs/detect"
|
|
RUN_NAME = "coffee_v1"
|
|
IMG_SIZE = 640
|
|
|
|
|
|
# ── Training hyperparameters ───────────────────────────────────────────────────
|
|
#
|
|
# Tuned for: small dataset (few-shot), packaged goods, retail shelf.
|
|
#
|
|
EPOCHS = 150
|
|
PATIENCE = 30 # Early Stopping - stops if no improvement
|
|
BATCH = 8 # Lower if you get OOM on CPU/small GPU.
|
|
LR0 = 0.005 # Initial LR (Learing Rate) - lower than default (0.01)
|
|
|
|
LRF = 0.01 # final LR = LR0 * LRF
|
|
MOMENTUM = 0.938
|
|
WEIGHT_DECAY = 0.0005
|
|
WARMUP_EPOCHS = 5 # Longer warmup stabilises few-shot training.
|
|
IOU_THRESH = 0.5 # IoU threshold for NMS.
|
|
CONF_THRESH = 0.25 # confidence threshold at val time.
|
|
|
|
|
|
# ── YOLOv8 built-in augmentation (on top of our albumentatios pass) ────────────
|
|
#
|
|
# These run during tarining batches - complementary to augment.py.
|
|
# Mosaic is the single most impactful augmentation for small datasets.
|
|
#
|
|
MOSAIC = 1.0 # mosaic probability
|
|
MIXUP = 0.1 # mixup probability
|
|
COPY_PASTE = 0.1
|
|
DEGREES = 5.0 # degrees ±5°
|
|
TRANSLATE = 0.1
|
|
SCALE = 0.5
|
|
FLIPUD = 0.0
|
|
FLIPLR = 0.5
|
|
HSV_H = 0.015
|
|
HSV_S = 0.7
|
|
HSV_V = 0.4
|
|
|
|
def check_enviroment():
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
if device == "cpu":
|
|
print("[INFO] No GPU detected - trainning on CPU")
|
|
print(" Expected time: ~60-90min for 150 epochs on a typical laptop.")
|
|
print(" Tips: use Google Colab or Kaggle they got free GPU, if training is too slow.\n")
|
|
else:
|
|
gpu= torch.cuda.get_device_name(0)
|
|
print(f"[INFO] GPU detected: {gpu}")
|
|
print(" Expected time: ~5-15min.\n")
|
|
return device
|
|
|
|
def train():
|
|
device = check_enviroment()
|
|
|
|
if not Path(DATASET_YAML).exists():
|
|
print(f"[ERROR] {DATASET_YAML} not found.")
|
|
print(" Run prepare_dataset.py first.")
|
|
return
|
|
|
|
model = YOLO(MODEL)
|
|
print(f"Starting training: {EPOCHS} epochs, img size {IMG_SIZE}")
|
|
print(f"Dataset: {DATASET_YAML}")
|
|
print(f"Output: {PROJECT}/{RUN_NAME}/weights/best.pt\n")
|
|
|
|
results = model.train(
|
|
data = DATASET_YAML,
|
|
epochs = EPOCHS,
|
|
patience = PATIENCE,
|
|
imgsz = IMG_SIZE,
|
|
batch = BATCH,
|
|
device = device,
|
|
project = PROJECT,
|
|
name = RUN_NAME,
|
|
exist_ok = True,
|
|
|
|
lr0 = LR0,
|
|
lrf = LRF,
|
|
momentum = MOMENTUM,
|
|
weight_decay = WEIGHT_DECAY,
|
|
warmup_epochs = WARMUP_EPOCHS,
|
|
|
|
iou = IOU_THRESH,
|
|
conf = CONF_THRESH,
|
|
|
|
# Built-in augmentation
|
|
mosaic = MOSAIC,
|
|
mixup = MIXUP,
|
|
copy_paste = COPY_PASTE,
|
|
degrees = DEGREES,
|
|
translate = TRANSLATE,
|
|
scale = SCALE,
|
|
flipud = FLIPUD,
|
|
fliplr = FLIPLR,
|
|
hsv_h = HSV_H,
|
|
hsv_s = HSV_S,
|
|
hsv_v = HSV_V,
|
|
|
|
save = True,
|
|
save_period = 10, # Checkpoint every 10 epochs
|
|
plots = True, # Saves training curvse as PNGs
|
|
verbose = True,
|
|
seed = 42,
|
|
)
|
|
|
|
best_weights = Path(PROJECT) / RUN_NAME / "weights" / "best.pt"
|
|
print("\n" + "-" * 60)
|
|
print(f"Training complete.")
|
|
print(f"Best weights saved -> {best_weights}")
|
|
print(f"mAP50: {results.results_dict.get('metrics/mAP50(B)', 'N/A'):.3f}")
|
|
print(f"mAP50-95: {results.results_dict.get('metrics/mAP50-95(B)', 'N/A'):.3f}")
|
|
print("-" * 60)
|
|
print("\nRun next: python detect.py --source scene.jpg")
|
|
|
|
return best_weights
|
|
|
|
if __name__ == "__main__":
|
|
train()
|
|
|
|
|