chore: normalize line endings (CRLF -> LF)
No content changes: git diff --ignore-all-space over these files is empty. The churn came from editing on Windows against a repo checked out with LF.
This commit is contained in:
fhanyuh committed
1 parent
15566a6951
commit
caf8e98378
315 files changed
+86950
-86950
No files matched your search
@@ -1,132 +1,132 @@
|
||||
#!/usr/bin/env python3
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
import pickle
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torchvision import transforms
|
||||
from pathlib import Path
|
||||
|
||||
# Setup directories
|
||||
SCRIPT_DIR = Path(__file__).resolve().parent
|
||||
DEFAULT_DATASET_DIR = SCRIPT_DIR / "foto-kemasan-v2"
|
||||
DEFAULT_MODELS_DIR = SCRIPT_DIR / "models"
|
||||
DEFAULT_INDEX_PATH = DEFAULT_MODELS_DIR / "dinov2_index.pkl"
|
||||
|
||||
# Allowed image extensions
|
||||
IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".webp", ".bmp")
|
||||
|
||||
# DINOv2 Image preprocessing
|
||||
DINOV2_TRANSFORMS = transforms.Compose([
|
||||
transforms.Resize((224, 224)),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
||||
])
|
||||
|
||||
|
||||
def get_embedding(dinov2_model, image: Image.Image, device):
|
||||
if image.mode != "RGB":
|
||||
image = image.convert("RGB")
|
||||
|
||||
tensor = DINOV2_TRANSFORMS(image).unsqueeze(0).to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
embedding = dinov2_model(tensor)
|
||||
# L2 normalization for dot product similarity
|
||||
embedding = embedding / embedding.norm(dim=-1, keepdim=True)
|
||||
|
||||
return embedding.squeeze(0).cpu().numpy()
|
||||
|
||||
|
||||
def run_indexing(src_dir=DEFAULT_DATASET_DIR, out_path=DEFAULT_INDEX_PATH):
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
print(f"Using device: {device}")
|
||||
|
||||
src_path = Path(src_dir).resolve()
|
||||
out_file_path = Path(out_path).resolve()
|
||||
|
||||
if not src_path.is_dir():
|
||||
print(f"Error: Source dataset directory not found: {src_path}")
|
||||
return False
|
||||
|
||||
out_file_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Load DINOv2 Model from Torch Hub
|
||||
print("Loading DINOv2 model (dinov2_vits14)...")
|
||||
t0 = time.perf_counter()
|
||||
dinov2_model = torch.hub.load("facebookresearch/dinov2", "dinov2_vits14").to(device)
|
||||
dinov2_model.eval()
|
||||
print(f"DINOv2 loaded in {time.perf_counter() - t0:.2f}s")
|
||||
|
||||
# Scan dataset directory
|
||||
class_dirs = [d for d in src_path.iterdir() if d.is_dir()]
|
||||
class_dirs.sort()
|
||||
|
||||
embeddings_list = []
|
||||
metadata_list = []
|
||||
|
||||
total_images = 0
|
||||
indexed_images = 0
|
||||
|
||||
for c_dir in class_dirs:
|
||||
class_name = c_dir.name
|
||||
|
||||
images = sorted(
|
||||
[f for f in c_dir.iterdir() if f.suffix.lower() in IMAGE_EXTS],
|
||||
key=lambda p: p.name
|
||||
)
|
||||
|
||||
if not images:
|
||||
continue
|
||||
|
||||
print(f"Processing class: {class_name} ({len(images)} images)")
|
||||
total_images += len(images)
|
||||
|
||||
for img_file in images:
|
||||
try:
|
||||
# Load image
|
||||
image = Image.open(img_file).convert("RGB")
|
||||
|
||||
# Extract DINOv2 embedding (using whole image as reference photo)
|
||||
embedding = get_embedding(dinov2_model, image, device)
|
||||
|
||||
embeddings_list.append(embedding)
|
||||
metadata_list.append({
|
||||
"class_name": class_name,
|
||||
"image_path": str(img_file.relative_to(src_path.parent)),
|
||||
"file_name": img_file.name
|
||||
})
|
||||
indexed_images += 1
|
||||
|
||||
except Exception as e:
|
||||
print(f" [Error] Failed to process {img_file.name}: {e}")
|
||||
|
||||
# Save the index
|
||||
if embeddings_list:
|
||||
embeddings_arr = np.vstack(embeddings_list)
|
||||
index_data = {
|
||||
"embeddings": embeddings_arr,
|
||||
"metadata": metadata_list
|
||||
}
|
||||
|
||||
with open(out_file_path, "wb") as f:
|
||||
pickle.dump(index_data, f)
|
||||
|
||||
print(f"\nSuccess! Indexed {indexed_images}/{total_images} images.")
|
||||
print(f"DINOv2 Vector Index saved to: {out_file_path}")
|
||||
return True
|
||||
else:
|
||||
print("\n[Warning] No images were successfully indexed.")
|
||||
return False
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
parser = argparse.ArgumentParser(description="Build DINOv2 image vector index for PFM products")
|
||||
parser.add_argument("--src-dir", default=str(DEFAULT_DATASET_DIR), help="Source directory of classes")
|
||||
parser.add_argument("--output", default=str(DEFAULT_INDEX_PATH), help="Output pickle index path")
|
||||
args = parser.parse_args()
|
||||
|
||||
run_indexing(src_dir=args.src_dir, out_path=args.output)
|
||||
#!/usr/bin/env python3
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
import pickle
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torchvision import transforms
|
||||
from pathlib import Path
|
||||
|
||||
# Setup directories
|
||||
SCRIPT_DIR = Path(__file__).resolve().parent
|
||||
DEFAULT_DATASET_DIR = SCRIPT_DIR / "foto-kemasan-v2"
|
||||
DEFAULT_MODELS_DIR = SCRIPT_DIR / "models"
|
||||
DEFAULT_INDEX_PATH = DEFAULT_MODELS_DIR / "dinov2_index.pkl"
|
||||
|
||||
# Allowed image extensions
|
||||
IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".webp", ".bmp")
|
||||
|
||||
# DINOv2 Image preprocessing
|
||||
DINOV2_TRANSFORMS = transforms.Compose([
|
||||
transforms.Resize((224, 224)),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
||||
])
|
||||
|
||||
|
||||
def get_embedding(dinov2_model, image: Image.Image, device):
|
||||
if image.mode != "RGB":
|
||||
image = image.convert("RGB")
|
||||
|
||||
tensor = DINOV2_TRANSFORMS(image).unsqueeze(0).to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
embedding = dinov2_model(tensor)
|
||||
# L2 normalization for dot product similarity
|
||||
embedding = embedding / embedding.norm(dim=-1, keepdim=True)
|
||||
|
||||
return embedding.squeeze(0).cpu().numpy()
|
||||
|
||||
|
||||
def run_indexing(src_dir=DEFAULT_DATASET_DIR, out_path=DEFAULT_INDEX_PATH):
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
print(f"Using device: {device}")
|
||||
|
||||
src_path = Path(src_dir).resolve()
|
||||
out_file_path = Path(out_path).resolve()
|
||||
|
||||
if not src_path.is_dir():
|
||||
print(f"Error: Source dataset directory not found: {src_path}")
|
||||
return False
|
||||
|
||||
out_file_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Load DINOv2 Model from Torch Hub
|
||||
print("Loading DINOv2 model (dinov2_vits14)...")
|
||||
t0 = time.perf_counter()
|
||||
dinov2_model = torch.hub.load("facebookresearch/dinov2", "dinov2_vits14").to(device)
|
||||
dinov2_model.eval()
|
||||
print(f"DINOv2 loaded in {time.perf_counter() - t0:.2f}s")
|
||||
|
||||
# Scan dataset directory
|
||||
class_dirs = [d for d in src_path.iterdir() if d.is_dir()]
|
||||
class_dirs.sort()
|
||||
|
||||
embeddings_list = []
|
||||
metadata_list = []
|
||||
|
||||
total_images = 0
|
||||
indexed_images = 0
|
||||
|
||||
for c_dir in class_dirs:
|
||||
class_name = c_dir.name
|
||||
|
||||
images = sorted(
|
||||
[f for f in c_dir.iterdir() if f.suffix.lower() in IMAGE_EXTS],
|
||||
key=lambda p: p.name
|
||||
)
|
||||
|
||||
if not images:
|
||||
continue
|
||||
|
||||
print(f"Processing class: {class_name} ({len(images)} images)")
|
||||
total_images += len(images)
|
||||
|
||||
for img_file in images:
|
||||
try:
|
||||
# Load image
|
||||
image = Image.open(img_file).convert("RGB")
|
||||
|
||||
# Extract DINOv2 embedding (using whole image as reference photo)
|
||||
embedding = get_embedding(dinov2_model, image, device)
|
||||
|
||||
embeddings_list.append(embedding)
|
||||
metadata_list.append({
|
||||
"class_name": class_name,
|
||||
"image_path": str(img_file.relative_to(src_path.parent)),
|
||||
"file_name": img_file.name
|
||||
})
|
||||
indexed_images += 1
|
||||
|
||||
except Exception as e:
|
||||
print(f" [Error] Failed to process {img_file.name}: {e}")
|
||||
|
||||
# Save the index
|
||||
if embeddings_list:
|
||||
embeddings_arr = np.vstack(embeddings_list)
|
||||
index_data = {
|
||||
"embeddings": embeddings_arr,
|
||||
"metadata": metadata_list
|
||||
}
|
||||
|
||||
with open(out_file_path, "wb") as f:
|
||||
pickle.dump(index_data, f)
|
||||
|
||||
print(f"\nSuccess! Indexed {indexed_images}/{total_images} images.")
|
||||
print(f"DINOv2 Vector Index saved to: {out_file_path}")
|
||||
return True
|
||||
else:
|
||||
print("\n[Warning] No images were successfully indexed.")
|
||||
return False
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
parser = argparse.ArgumentParser(description="Build DINOv2 image vector index for PFM products")
|
||||
parser.add_argument("--src-dir", default=str(DEFAULT_DATASET_DIR), help="Source directory of classes")
|
||||
parser.add_argument("--output", default=str(DEFAULT_INDEX_PATH), help="Output pickle index path")
|
||||
args = parser.parse_args()
|
||||
|
||||
run_indexing(src_dir=args.src_dir, out_path=args.output)
|
||||
@@ -1,375 +1,375 @@
|
||||
#!/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"}
|
||||
AUG_SUFFIX_RE = re.compile(r"_aug_\d+$")
|
||||
|
||||
|
||||
def is_image_file(path: Path) -> bool:
|
||||
return path.is_file() and path.suffix.lower() in VALID_IMAGE_EXTENSIONS
|
||||
|
||||
|
||||
def _source_group_key(filename_stem: str) -> str:
|
||||
"""Strip an `_aug_<n>` suffix so an augmented image groups with its source photo."""
|
||||
return AUG_SUFFIX_RE.sub("", filename_stem)
|
||||
|
||||
|
||||
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.
|
||||
|
||||
Splits by *source photo group*, not by individual file: an augmented image
|
||||
(`photo1_aug_2.jpeg`) always stays in the same split as its source
|
||||
(`photo1.jpeg`). Splitting file-by-file would let near-duplicate images
|
||||
land on opposite sides of train/val, inflating val accuracy with
|
||||
memorization instead of measuring generalization.
|
||||
"""
|
||||
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,
|
||||
)
|
||||
|
||||
num_images = len(images)
|
||||
if num_images == 0:
|
||||
print(f"Warning: Class '{class_name}' has 0 images. Skipping.")
|
||||
continue
|
||||
|
||||
# Group by source photo (stripping any `_aug_N` suffix) so an
|
||||
# augmented image and the photo it came from always land on the same
|
||||
# side of the split.
|
||||
groups: dict[str, list[Path]] = {}
|
||||
for img in images:
|
||||
groups.setdefault(_source_group_key(img.stem), []).append(img)
|
||||
group_keys = sorted(groups.keys())
|
||||
random.shuffle(group_keys)
|
||||
|
||||
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)
|
||||
|
||||
num_groups = len(group_keys)
|
||||
if num_groups == 1:
|
||||
train_groups = group_keys
|
||||
val_groups = group_keys
|
||||
elif num_groups == 2:
|
||||
train_groups = [group_keys[0]]
|
||||
val_groups = [group_keys[1]]
|
||||
else:
|
||||
split_idx = max(1, int(num_groups * split_ratio))
|
||||
split_idx = min(split_idx, num_groups - 1)
|
||||
train_groups = group_keys[:split_idx]
|
||||
val_groups = group_keys[split_idx:]
|
||||
|
||||
train_images = [img for key in train_groups for img in groups[key]]
|
||||
val_images = [img for key in val_groups for img in groups[key]]
|
||||
|
||||
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 (from {num_groups} source photos, {num_images} files total)"
|
||||
)
|
||||
|
||||
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()
|
||||
#!/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"}
|
||||
AUG_SUFFIX_RE = re.compile(r"_aug_\d+$")
|
||||
|
||||
|
||||
def is_image_file(path: Path) -> bool:
|
||||
return path.is_file() and path.suffix.lower() in VALID_IMAGE_EXTENSIONS
|
||||
|
||||
|
||||
def _source_group_key(filename_stem: str) -> str:
|
||||
"""Strip an `_aug_<n>` suffix so an augmented image groups with its source photo."""
|
||||
return AUG_SUFFIX_RE.sub("", filename_stem)
|
||||
|
||||
|
||||
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.
|
||||
|
||||
Splits by *source photo group*, not by individual file: an augmented image
|
||||
(`photo1_aug_2.jpeg`) always stays in the same split as its source
|
||||
(`photo1.jpeg`). Splitting file-by-file would let near-duplicate images
|
||||
land on opposite sides of train/val, inflating val accuracy with
|
||||
memorization instead of measuring generalization.
|
||||
"""
|
||||
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,
|
||||
)
|
||||
|
||||
num_images = len(images)
|
||||
if num_images == 0:
|
||||
print(f"Warning: Class '{class_name}' has 0 images. Skipping.")
|
||||
continue
|
||||
|
||||
# Group by source photo (stripping any `_aug_N` suffix) so an
|
||||
# augmented image and the photo it came from always land on the same
|
||||
# side of the split.
|
||||
groups: dict[str, list[Path]] = {}
|
||||
for img in images:
|
||||
groups.setdefault(_source_group_key(img.stem), []).append(img)
|
||||
group_keys = sorted(groups.keys())
|
||||
random.shuffle(group_keys)
|
||||
|
||||
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)
|
||||
|
||||
num_groups = len(group_keys)
|
||||
if num_groups == 1:
|
||||
train_groups = group_keys
|
||||
val_groups = group_keys
|
||||
elif num_groups == 2:
|
||||
train_groups = [group_keys[0]]
|
||||
val_groups = [group_keys[1]]
|
||||
else:
|
||||
split_idx = max(1, int(num_groups * split_ratio))
|
||||
split_idx = min(split_idx, num_groups - 1)
|
||||
train_groups = group_keys[:split_idx]
|
||||
val_groups = group_keys[split_idx:]
|
||||
|
||||
train_images = [img for key in train_groups for img in groups[key]]
|
||||
val_images = [img for key in val_groups for img in groups[key]]
|
||||
|
||||
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 (from {num_groups} source photos, {num_images} files total)"
|
||||
)
|
||||
|
||||
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