feat: consolidate backend and docker-compose setup
This commit is contained in:
Rafhan Mazaya Fathurrahman committed
commit
ff3753a745
306 files changed
+35450
No files matched your search
@@ -0,0 +1,351 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Ultralytics YOLO Classification Training Script
|
||||
Trains a product-packaging classifier from class folders in `foto-kemasan-v2`.
|
||||
|
||||
Each subfolder under `foto-kemasan-v2/` is one product class; images live directly
|
||||
inside that folder.
|
||||
|
||||
Usage (from repo root or this directory):
|
||||
# 1) Train the model (defaults to foto-kemasan-v2, 100 epochs)
|
||||
uv run python pfm-web-app/public/produk-pfm/train_classifier.py train --imgsz 224
|
||||
|
||||
# 2) Run prediction on an image using the trained weights
|
||||
uv run python pfm-web-app/public/produk-pfm/train_classifier.py predict \\
|
||||
--image "pfm-web-app/public/produk-pfm/foto-kemasan-v2/15030101 FIESTA CRINKLE CUT 500 GR/WhatsApp Image 2026-05-28 at 11.46.31.jpeg"
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import shutil
|
||||
import random
|
||||
import argparse
|
||||
from datetime import date
|
||||
from pathlib import Path
|
||||
import torch
|
||||
|
||||
try:
|
||||
from ultralytics import YOLO
|
||||
except ImportError:
|
||||
print("Error: 'ultralytics' library not found. Please install it using: uv add ultralytics")
|
||||
sys.exit(1)
|
||||
|
||||
SCRIPT_DIR = Path(__file__).resolve().parent
|
||||
DEFAULT_DATASET_DIR = SCRIPT_DIR / "foto-kemasan-v2"
|
||||
DEFAULT_SPLIT_DIR = SCRIPT_DIR / "yolo_dataset"
|
||||
DEFAULT_MODEL = SCRIPT_DIR / "yolo26n-cls.pt"
|
||||
DEFAULT_MODELS_DIR = SCRIPT_DIR / "models"
|
||||
DEFAULT_PROJECT = SCRIPT_DIR / "runs" / "classify"
|
||||
DEFAULT_EPOCHS = 100
|
||||
|
||||
|
||||
def classifier_output_path(epochs: int = DEFAULT_EPOCHS, run_date: date | None = None) -> Path:
|
||||
"""Build the dated classifier artifact path under models/."""
|
||||
run_date = run_date or date.today()
|
||||
return DEFAULT_MODELS_DIR / f"produk-pfm-classifier-26n-{epochs}e-{run_date:%Y-%m-%d}.pt"
|
||||
|
||||
|
||||
def _classifier_date_from_name(path: Path) -> date | None:
|
||||
match = re.search(
|
||||
r"produk-pfm-classifier-26n-\d+e-(\d{4}-\d{2}-\d{2})\.pt$",
|
||||
path.name,
|
||||
)
|
||||
if not match:
|
||||
return None
|
||||
year, month, day = (int(part) for part in match.group(1).split("-"))
|
||||
return date(year, month, day)
|
||||
|
||||
|
||||
def latest_classifier_weights(models_dir: Path = DEFAULT_MODELS_DIR) -> Path:
|
||||
"""Return the newest produk-pfm-classifier weights in models/, if any."""
|
||||
if not models_dir.is_dir():
|
||||
return classifier_output_path()
|
||||
|
||||
candidates = list(models_dir.glob("produk-pfm-classifier-26n-*e-*.pt"))
|
||||
if not candidates:
|
||||
return classifier_output_path()
|
||||
|
||||
def sort_key(path: Path) -> tuple[date, float]:
|
||||
name_date = _classifier_date_from_name(path) or date.min
|
||||
return (name_date, path.stat().st_mtime)
|
||||
|
||||
return max(candidates, key=sort_key)
|
||||
|
||||
VALID_IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
|
||||
|
||||
|
||||
def is_image_file(path: Path) -> bool:
|
||||
return path.is_file() and path.suffix.lower() in VALID_IMAGE_EXTENSIONS
|
||||
|
||||
|
||||
def split_dataset(src_dir: Path, dest_dir: Path, split_ratio: float = 0.8, seed: int = 42):
|
||||
"""
|
||||
Split class folders from src_dir into train/val folders in dest_dir.
|
||||
Ensures every class with 2+ images keeps at least one image in validation.
|
||||
"""
|
||||
random.seed(seed)
|
||||
|
||||
train_dir = dest_dir / "train"
|
||||
val_dir = dest_dir / "val"
|
||||
|
||||
if dest_dir.exists():
|
||||
print(f"Cleaning existing split directory: {dest_dir}")
|
||||
shutil.rmtree(dest_dir)
|
||||
|
||||
train_dir.mkdir(parents=True, exist_ok=True)
|
||||
val_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
exclude_dirs = {dest_dir.name, "train", "val"}
|
||||
class_dirs = [d for d in src_dir.iterdir() if d.is_dir() and d.name not in exclude_dirs]
|
||||
class_dirs.sort()
|
||||
|
||||
print(f"Found {len(class_dirs)} product classes in {src_dir}")
|
||||
|
||||
total_train = 0
|
||||
total_val = 0
|
||||
|
||||
for c_dir in class_dirs:
|
||||
class_name = c_dir.name
|
||||
images = sorted(
|
||||
[f for f in c_dir.iterdir() if is_image_file(f)],
|
||||
key=lambda p: p.name,
|
||||
)
|
||||
random.shuffle(images)
|
||||
|
||||
num_images = len(images)
|
||||
if num_images == 0:
|
||||
print(f"Warning: Class '{class_name}' has 0 images. Skipping.")
|
||||
continue
|
||||
|
||||
class_train_dir = train_dir / class_name
|
||||
class_val_dir = val_dir / class_name
|
||||
class_train_dir.mkdir(parents=True, exist_ok=True)
|
||||
class_val_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if num_images == 1:
|
||||
train_images = images
|
||||
val_images = images
|
||||
elif num_images == 2:
|
||||
train_images = [images[0]]
|
||||
val_images = [images[1]]
|
||||
else:
|
||||
split_idx = max(1, int(num_images * split_ratio))
|
||||
split_idx = min(split_idx, num_images - 1)
|
||||
train_images = images[:split_idx]
|
||||
val_images = images[split_idx:]
|
||||
|
||||
for img in train_images:
|
||||
shutil.copy(img, class_train_dir / img.name)
|
||||
total_train += 1
|
||||
|
||||
for img in val_images:
|
||||
shutil.copy(img, class_val_dir / img.name)
|
||||
total_val += 1
|
||||
|
||||
print(
|
||||
f" Class '{class_name}': {len(train_images)} train, "
|
||||
f"{len(val_images)} val (total: {num_images})"
|
||||
)
|
||||
|
||||
print(f"Dataset split completed: {total_train} train images, {total_val} validation images.")
|
||||
print(f"Split dataset located at: {dest_dir.absolute()}")
|
||||
|
||||
|
||||
def train_model(args):
|
||||
"""Handles training the YOLO classification model."""
|
||||
src_path = Path(args.src_dir).resolve()
|
||||
dest_path = Path(args.split_dir).resolve()
|
||||
|
||||
if not src_path.is_dir():
|
||||
print(f"Error: Source dataset directory not found: {src_path}")
|
||||
sys.exit(1)
|
||||
|
||||
print(f"--- Preparing Dataset from {src_path} ---")
|
||||
split_dataset(src_path, dest_path, split_ratio=args.split_ratio)
|
||||
|
||||
model_path = Path(args.model).resolve()
|
||||
print(f"\n--- Initializing YOLO Model ({model_path}) ---")
|
||||
model = YOLO(str(model_path))
|
||||
|
||||
if args.device:
|
||||
device = args.device
|
||||
else:
|
||||
device = "0" if torch.cuda.is_available() else "cpu"
|
||||
print(f"Using device: {device}")
|
||||
|
||||
print("\n--- Starting Training ---")
|
||||
results = model.train(
|
||||
data=str(dest_path),
|
||||
epochs=args.epochs,
|
||||
imgsz=args.imgsz,
|
||||
batch=args.batch,
|
||||
device=device,
|
||||
project=str(Path(args.project).resolve()),
|
||||
name=args.name,
|
||||
exist_ok=True,
|
||||
workers=args.workers,
|
||||
lr0=args.lr,
|
||||
optimizer=args.optimizer,
|
||||
seed=42,
|
||||
)
|
||||
|
||||
best_weights = Path(results.save_dir) / "weights" / "best.pt"
|
||||
output_path = Path(args.output).resolve() if args.output else classifier_output_path(args.epochs)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy2(best_weights, output_path)
|
||||
|
||||
print("\nTraining completed successfully!")
|
||||
print(f"Run weights saved at: {best_weights}")
|
||||
print(f"Published model saved at: {output_path}")
|
||||
|
||||
if args.export:
|
||||
print("\n--- Exporting model to ONNX format ---")
|
||||
try:
|
||||
export_model = YOLO(str(output_path))
|
||||
onnx_path = Path(export_model.export(format="onnx"))
|
||||
dated_onnx = output_path.with_suffix(".onnx")
|
||||
if onnx_path.resolve() != dated_onnx.resolve():
|
||||
shutil.copy2(onnx_path, dated_onnx)
|
||||
print(f"Model exported successfully to: {dated_onnx}")
|
||||
except Exception as e:
|
||||
print(f"Warning: ONNX export failed: {e}")
|
||||
|
||||
print("\nYou can run predictions with:")
|
||||
print(f" uv run python {Path(__file__).name} predict --image <image_path> --model {output_path}")
|
||||
|
||||
|
||||
def predict_image(args):
|
||||
"""Runs classification inference on a single image."""
|
||||
model_path = Path(args.model).resolve()
|
||||
image_path = Path(args.image).resolve()
|
||||
|
||||
if not model_path.exists():
|
||||
print(f"Error: Model weights not found at {model_path}")
|
||||
sys.exit(1)
|
||||
|
||||
if not image_path.exists():
|
||||
print(f"Error: Target image file not found at {image_path}")
|
||||
sys.exit(1)
|
||||
|
||||
print(f"Loading model from {model_path}...")
|
||||
model = YOLO(str(model_path))
|
||||
|
||||
print(f"Running prediction on {image_path}...")
|
||||
results = model(str(image_path))
|
||||
|
||||
for result in results:
|
||||
probs = result.probs
|
||||
top1_idx = probs.top1
|
||||
top1_conf = float(probs.top1conf)
|
||||
top1_name = result.names[top1_idx]
|
||||
|
||||
print("\n=== Classification Results ===")
|
||||
print(f"Top-1 Prediction: {top1_name} (Confidence: {top1_conf:.4f})")
|
||||
print("\nAll Probabilities:")
|
||||
|
||||
sorted_probs = sorted(
|
||||
[(result.names[i], float(val)) for i, val in enumerate(probs.data)],
|
||||
key=lambda x: x[1],
|
||||
reverse=True,
|
||||
)
|
||||
for name, score in sorted_probs:
|
||||
print(f" {name}: {score:.4f}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Ultralytics YOLO classification utility for produk-pfm packaging photos."
|
||||
)
|
||||
subparsers = parser.add_subparsers(dest="command", required=True, help="Command to run")
|
||||
|
||||
train_parser = subparsers.add_parser("train", help="Train a classification model")
|
||||
train_parser.add_argument(
|
||||
"--src-dir",
|
||||
type=str,
|
||||
default=str(DEFAULT_DATASET_DIR),
|
||||
help=f"Source dataset directory with one class folder per product (default: {DEFAULT_DATASET_DIR.name})",
|
||||
)
|
||||
train_parser.add_argument(
|
||||
"--split-dir",
|
||||
type=str,
|
||||
default=str(DEFAULT_SPLIT_DIR),
|
||||
help="Output split dataset directory",
|
||||
)
|
||||
train_parser.add_argument(
|
||||
"--split-ratio",
|
||||
type=float,
|
||||
default=0.8,
|
||||
help="Train/val split ratio for classes with 3+ images (default: 0.8)",
|
||||
)
|
||||
train_parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default=str(DEFAULT_MODEL),
|
||||
help="Pretrained model (e.g. yolo26n-cls.pt, yolo11n-cls.pt, yolov8n-cls.pt)",
|
||||
)
|
||||
train_parser.add_argument(
|
||||
"--epochs",
|
||||
type=int,
|
||||
default=DEFAULT_EPOCHS,
|
||||
help=f"Number of training epochs (default: {DEFAULT_EPOCHS})",
|
||||
)
|
||||
train_parser.add_argument(
|
||||
"--output",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Published .pt output path (default: "
|
||||
"models/produk-pfm-classifier-26n-{epochs}e-{YYYY-MM-DD}.pt)"
|
||||
),
|
||||
)
|
||||
train_parser.add_argument("--imgsz", type=int, default=224, help="Target image size for classification")
|
||||
train_parser.add_argument("--batch", type=int, default=8, help="Batch size for training")
|
||||
train_parser.add_argument(
|
||||
"--device",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Device to run on (e.g. 0 or 'cpu'). Default is GPU if available.",
|
||||
)
|
||||
train_parser.add_argument(
|
||||
"--project",
|
||||
type=str,
|
||||
default=str(DEFAULT_PROJECT),
|
||||
help="Project output folder name",
|
||||
)
|
||||
train_parser.add_argument("--name", type=str, default="train", help="Experiment name")
|
||||
train_parser.add_argument("--workers", type=int, default=4, help="Number of data loading workers")
|
||||
train_parser.add_argument("--lr", type=float, default=0.01, help="Initial learning rate")
|
||||
train_parser.add_argument(
|
||||
"--optimizer",
|
||||
type=str,
|
||||
default="auto",
|
||||
choices=["SGD", "Adam", "AdamW", "RMSProp", "auto"],
|
||||
help="Optimizer to use",
|
||||
)
|
||||
train_parser.add_argument(
|
||||
"--export",
|
||||
action="store_true",
|
||||
default=True,
|
||||
help="Export model to ONNX after training",
|
||||
)
|
||||
|
||||
predict_parser = subparsers.add_parser("predict", help="Predict class of an image")
|
||||
predict_parser.add_argument("--image", type=str, required=True, help="Path to image file")
|
||||
predict_parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default=str(latest_classifier_weights()),
|
||||
help="Path to trained YOLO .pt model weights (default: newest models/produk-pfm-classifier-*.pt)",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.command == "train":
|
||||
train_model(args)
|
||||
elif args.command == "predict":
|
||||
predict_image(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in new issue
Block a user