First Commit
This commit is contained in:
commit
9f3dfdace9
234
.gitignore
vendored
Normal file
234
.gitignore
vendored
Normal file
@ -0,0 +1,234 @@
|
||||
# Auto Label
|
||||
.autolabel_staging/
|
||||
_rejected_/
|
||||
_unreviewed_/
|
||||
|
||||
# Crops Dataset
|
||||
crops_dataset/
|
||||
crops_dataset.zip
|
||||
|
||||
# Archived files.
|
||||
*.zip
|
||||
*.7z
|
||||
*.rar
|
||||
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[codz]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
share/python-wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
|
||||
# PyInstaller
|
||||
# Usually these files are written by a python script from a template
|
||||
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Installer logs
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.nox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
*.py.cover
|
||||
*.lcov
|
||||
.hypothesis/
|
||||
.pytest_cache/
|
||||
cover/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Django stuff:
|
||||
*.log
|
||||
local_settings.py
|
||||
db.sqlite3
|
||||
db.sqlite3-journal
|
||||
|
||||
# Flask stuff:
|
||||
instance/
|
||||
.webassets-cache
|
||||
|
||||
# Scrapy stuff:
|
||||
.scrapy
|
||||
|
||||
# Sphinx documentation
|
||||
docs/_build/
|
||||
|
||||
# PyBuilder
|
||||
.pybuilder/
|
||||
target/
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# IPython
|
||||
profile_default/
|
||||
ipython_config.py
|
||||
|
||||
# pyenv
|
||||
# For a library or package, you might want to ignore these files since the code is
|
||||
# intended to run in multiple environments; otherwise, check them in:
|
||||
# .python-version
|
||||
|
||||
# pipenv
|
||||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
||||
# install all needed dependencies.
|
||||
# Pipfile.lock
|
||||
|
||||
# UV
|
||||
# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
# uv.lock
|
||||
|
||||
# poetry
|
||||
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
||||
# poetry.lock
|
||||
# poetry.toml
|
||||
|
||||
# pdm
|
||||
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
||||
# pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
|
||||
# https://pdm-project.org/en/latest/usage/project/#working-with-version-control
|
||||
# pdm.lock
|
||||
# pdm.toml
|
||||
.pdm-python
|
||||
.pdm-build/
|
||||
|
||||
# pixi
|
||||
# Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
|
||||
# pixi.lock
|
||||
# Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
|
||||
# in the .venv directory. It is recommended not to include this directory in version control.
|
||||
.pixi/*
|
||||
!.pixi/config.toml
|
||||
|
||||
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
||||
__pypackages__/
|
||||
|
||||
# Celery stuff
|
||||
celerybeat-schedule*
|
||||
celerybeat.pid
|
||||
|
||||
# Redis
|
||||
*.rdb
|
||||
*.aof
|
||||
*.pid
|
||||
|
||||
# RabbitMQ
|
||||
mnesia/
|
||||
rabbitmq/
|
||||
rabbitmq-data/
|
||||
|
||||
# ActiveMQ
|
||||
activemq-data/
|
||||
|
||||
# SageMath parsed files
|
||||
*.sage.py
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.envrc
|
||||
.venv
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
env.bak/
|
||||
venv.bak/
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
.dmypy.json
|
||||
dmypy.json
|
||||
|
||||
# Pyre type checker
|
||||
.pyre/
|
||||
|
||||
# pytype static type analyzer
|
||||
.pytype/
|
||||
|
||||
# Cython debug symbols
|
||||
cython_debug/
|
||||
|
||||
# PyCharm
|
||||
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
||||
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
||||
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
||||
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
||||
# .idea/
|
||||
|
||||
# Abstra
|
||||
# Abstra is an AI-powered process automation framework.
|
||||
# Ignore directories containing user credentials, local state, and settings.
|
||||
# Learn more at https://abstra.io/docs
|
||||
.abstra/
|
||||
|
||||
# Visual Studio Code
|
||||
# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
|
||||
# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
|
||||
# and can be added to the global gitignore or merged into this file. However, if you prefer,
|
||||
# you could uncomment the following to ignore the entire vscode folder
|
||||
# .vscode/
|
||||
# Temporary file for partial code execution
|
||||
tempCodeRunnerFile.py
|
||||
|
||||
# Ruff stuff:
|
||||
.ruff_cache/
|
||||
|
||||
# PyPI configuration file
|
||||
.pypirc
|
||||
|
||||
# Marimo
|
||||
marimo/_static/
|
||||
marimo/_lsp/
|
||||
__marimo__/
|
||||
|
||||
# Streamlit
|
||||
.streamlit/secrets.toml
|
||||
194
1_generate_crop_dataset.py
Normal file
194
1_generate_crop_dataset.py
Normal file
@ -0,0 +1,194 @@
|
||||
"""
|
||||
Script 1 — Crop Dataset Generator
|
||||
===================================
|
||||
Runs your fine-tuned YOLOv8n model on a folder of shelf images,
|
||||
crops every detection, and dumps everything into a single flat folder.
|
||||
|
||||
You classify ONCE, then run script 1b_split_dataset.py to create
|
||||
the train / val / test split automatically.
|
||||
|
||||
Output structure
|
||||
----------------
|
||||
data/
|
||||
unknown/
|
||||
img001_crop0000.jpg
|
||||
img001_crop0001.jpg
|
||||
img002_crop0000.jpg
|
||||
...
|
||||
|
||||
After generation
|
||||
----------------
|
||||
1. Rename / move crops into class subfolders inside data/:
|
||||
data/
|
||||
cola_can/ img001_crop0000.jpg …
|
||||
pepsi_can/ img003_crop0002.jpg …
|
||||
…
|
||||
|
||||
Use Label Studio, a file manager, or any annotation tool.
|
||||
The `unknown/` folder holds anything you haven't labelled yet.
|
||||
|
||||
2. Then run:
|
||||
python 1b_split_dataset.py --data_dir data --output_dir crops_dataset
|
||||
|
||||
Usage
|
||||
-----
|
||||
python 1_generate_crop_dataset.py \
|
||||
--source /path/to/shelf/images \
|
||||
--weights /path/to/best.pt \
|
||||
--output_dir data \
|
||||
--conf 0.25 \
|
||||
--iou 0.45 \
|
||||
--img_size 640 \
|
||||
--padding 10 \
|
||||
--min_size 32
|
||||
"""
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from ultralytics import YOLO
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
# ─────────────────────────── helpers ─────────────────────────────────────────
|
||||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser(description="Generate flat crop dataset from YOLOv8 detections")
|
||||
p.add_argument("--source", required=True , help="Directory with shelf images (jpg/png/…)")
|
||||
p.add_argument("--weights", default="detector/best.pt", help="Path to your fine-tuned YOLOv8n .pt file")
|
||||
p.add_argument("--output_dir", default="data" , help="Root output dir — crops go into <output_dir>/unknown/")
|
||||
p.add_argument("--conf", type=float, default=0.25 , help="Detection confidence threshold")
|
||||
p.add_argument("--iou", type=float, default=0.45 , help="NMS IoU threshold")
|
||||
p.add_argument("--img_size", type=int, default=640 , help="YOLO inference image size")
|
||||
p.add_argument("--padding", type=int, default=10 , help="Extra pixels around each crop (px)")
|
||||
p.add_argument("--min_size", type=int, default=32 , help="Discard crops smaller than this (px)")
|
||||
p.add_argument("--save_labels", action="store_true",
|
||||
help="Save a sidecar .txt with YOLO cls_id / confidence next to each crop")
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
def get_image_paths(source: str):
|
||||
exts = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tiff"}
|
||||
p = Path(source)
|
||||
|
||||
if p.is_file():
|
||||
# Single image passed directly e.g. --source 01.jpg
|
||||
if p.suffix.lower() not in exts:
|
||||
raise ValueError(f"File is not a supported image type: {p}")
|
||||
return [p]
|
||||
|
||||
if p.is_dir():
|
||||
# Folder — recurse for all images
|
||||
paths = sorted(f for f in p.rglob("*") if f.suffix.lower() in exts)
|
||||
if not paths:
|
||||
raise FileNotFoundError(f"No images found in directory: {p}")
|
||||
return paths
|
||||
|
||||
# Glob pattern e.g. --source "frames/*.jpg"
|
||||
paths = sorted(f for f in Path(".").glob(source) if f.suffix.lower() in exts)
|
||||
if not paths:
|
||||
raise FileNotFoundError(f"No images matched: {source}")
|
||||
return paths
|
||||
|
||||
|
||||
def crop_detection(image: np.ndarray, box, padding: int, min_size: int):
|
||||
"""
|
||||
Extract one padded crop from *image* using *box* (xyxy).
|
||||
Returns the crop array, or None if it is too small.
|
||||
"""
|
||||
x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
|
||||
h, w = image.shape[:2]
|
||||
|
||||
x1 = max(0, x1 - padding)
|
||||
y1 = max(0, y1 - padding)
|
||||
x2 = min(w, x2 + padding)
|
||||
y2 = min(h, y2 + padding)
|
||||
|
||||
if (x2 - x1) < min_size or (y2 - y1) < min_size:
|
||||
return None
|
||||
|
||||
return image[y1:y2, x1:x2]
|
||||
|
||||
|
||||
# ─────────────────────────── main ────────────────────────────────────────────
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
out_dir = Path(args.output_dir) / "unknown"
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# ── Load model ────────────────────────────────────────────────────────────
|
||||
print(f"[1/2] Loading model: {args.weights}")
|
||||
model = YOLO(args.weights)
|
||||
|
||||
# ── Detect & crop ─────────────────────────────────────────────────────────
|
||||
image_paths = get_image_paths(args.source)
|
||||
print(f"[2/2] Found {len(image_paths)} images — running detection …\n")
|
||||
|
||||
total_saved = 0
|
||||
total_skip = 0
|
||||
|
||||
for img_path in tqdm(image_paths, desc="Detecting & cropping"):
|
||||
image = cv2.imread(str(img_path))
|
||||
if image is None:
|
||||
print(f" ⚠ Cannot read {img_path.name} — skipping")
|
||||
continue
|
||||
|
||||
results = model(
|
||||
image,
|
||||
conf=args.conf,
|
||||
iou=args.iou,
|
||||
imgsz=args.img_size,
|
||||
verbose=False,
|
||||
)[0]
|
||||
|
||||
stem = img_path.stem
|
||||
for i, box in enumerate(results.boxes):
|
||||
crop = crop_detection(image, box, args.padding, args.min_size)
|
||||
|
||||
if crop is None:
|
||||
total_skip += 1
|
||||
continue
|
||||
|
||||
crop_name = f"{stem}_crop{i:04d}.jpg"
|
||||
crop_path = out_dir / crop_name
|
||||
cv2.imwrite(str(crop_path), crop)
|
||||
total_saved += 1
|
||||
|
||||
if args.save_labels:
|
||||
cls_id = int(box.cls[0])
|
||||
conf = float(box.conf[0])
|
||||
cls_name = model.names.get(cls_id, str(cls_id))
|
||||
crop_path.with_suffix(".txt").write_text(
|
||||
f"{cls_id} {cls_name} {conf:.4f}\n"
|
||||
)
|
||||
|
||||
# ── Summary ───────────────────────────────────────────────────────────────
|
||||
print("\n" + "=" * 55)
|
||||
print(" Crop generation complete!")
|
||||
print(f" Saved : {total_saved:,} crops -> {out_dir.resolve()}")
|
||||
print(f" Skipped : {total_skip:,} (below --min_size={args.min_size}px)")
|
||||
print("=" * 55)
|
||||
print("""
|
||||
Next steps:
|
||||
1. Organise crops into class subfolders inside data/:
|
||||
data/
|
||||
cola_can/ *.jpg
|
||||
pepsi_can/ *.jpg
|
||||
...
|
||||
Leave unlabelled crops in data/unknown/ — the split
|
||||
script ignores that folder automatically.
|
||||
|
||||
2. Run the splitter:
|
||||
python 1b_split_dataset.py \\
|
||||
--data_dir data \\
|
||||
--output_dir crops_dataset \\
|
||||
--val_split 0.15 \\
|
||||
--test_split 0.05
|
||||
""")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
221
1b_split_dataset.py
Normal file
221
1b_split_dataset.py
Normal file
@ -0,0 +1,221 @@
|
||||
"""
|
||||
Script 1b — Dataset Splitter
|
||||
==============================
|
||||
Takes your labelled data/ folder (one subfolder per class) and splits
|
||||
it into the train / val / test structure expected by the training script.
|
||||
|
||||
The split is STRATIFIED — every class gets the same train/val/test ratio,
|
||||
so rare classes are not accidentally wiped out of val or test.
|
||||
|
||||
Input layout (after you've labelled crops from script 1)
|
||||
---------------------------------------------------------
|
||||
data/
|
||||
cola_can/ ← labelled class folders
|
||||
img001_crop0000.jpg
|
||||
img002_crop0003.jpg
|
||||
...
|
||||
pepsi_can/
|
||||
...
|
||||
unknown/ ← IGNORED automatically (unlabelled crops)
|
||||
...
|
||||
|
||||
Output layout
|
||||
-------------
|
||||
crops_dataset/
|
||||
train/
|
||||
cola_can/ ...
|
||||
pepsi_can/ ...
|
||||
val/
|
||||
cola_can/ ...
|
||||
pepsi_can/ ...
|
||||
test/ (only if --test_split > 0)
|
||||
cola_can/ ...
|
||||
pepsi_can/ ...
|
||||
|
||||
Files are COPIED by default (--move to use move instead).
|
||||
Original data/ folder is left intact unless --move is used.
|
||||
|
||||
Usage
|
||||
-----
|
||||
python 1b_split_dataset.py \
|
||||
--data_dir data \
|
||||
--output_dir crops_dataset \
|
||||
--val_split 0.15 \
|
||||
--test_split 0.05 \
|
||||
--seed 42
|
||||
|
||||
Then run:
|
||||
python 2_train_classifier.py --data_dir crops_dataset
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import random
|
||||
import shutil
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
# ─────────────────────────── args ────────────────────────────────────────────
|
||||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser(description="Stratified train/val/test splitter")
|
||||
p.add_argument("--data_dir", required=True, help="Labelled data root (subfolders = classes)")
|
||||
p.add_argument("--output_dir", default="crops_dataset", help="Where to write the split dataset")
|
||||
p.add_argument("--val_split", type=float, default=0.15, help="Fraction for validation (default 0.15)")
|
||||
p.add_argument("--test_split", type=float, default=0.05, help="Fraction for test — 0 to skip")
|
||||
p.add_argument("--seed", type=int, default=42)
|
||||
p.add_argument("--move", action="store_true",
|
||||
help="Move files instead of copying (saves disk space)")
|
||||
p.add_argument("--min_per_class", type=int, default=2,
|
||||
help="Warn if a class has fewer than this many images after splitting")
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
# ─────────────────────────── helpers ─────────────────────────────────────────
|
||||
|
||||
IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tiff"}
|
||||
IGNORE_DIRS = {"unknown", "train", "val", "test"} # skip these silently
|
||||
|
||||
|
||||
def collect_classes(data_dir: Path) -> dict[str, list[Path]]:
|
||||
"""
|
||||
Scan *data_dir* for class subfolders and return {class_name: [image_paths]}.
|
||||
Folders named 'unknown', 'train', 'val', or 'test' are silently skipped.
|
||||
"""
|
||||
classes: dict[str, list[Path]] = {}
|
||||
for folder in sorted(data_dir.iterdir()):
|
||||
if not folder.is_dir():
|
||||
continue
|
||||
if folder.name.lower() in IGNORE_DIRS:
|
||||
print(f" ⚙ Skipping '{folder.name}/' (reserved name)")
|
||||
continue
|
||||
images = sorted(f for f in folder.iterdir() if f.suffix.lower() in IMAGE_EXTS)
|
||||
if images:
|
||||
classes[folder.name] = images
|
||||
else:
|
||||
print(f" ⚠ '{folder.name}/' has no images — skipped")
|
||||
return classes
|
||||
|
||||
|
||||
def stratified_split(
|
||||
images: list[Path],
|
||||
val_frac: float,
|
||||
test_frac: float,
|
||||
rng: random.Random,
|
||||
) -> tuple[list[Path], list[Path], list[Path]]:
|
||||
"""
|
||||
Split *images* into (train, val, test) respecting the requested fractions.
|
||||
Guarantees at least 1 image in val and test if the class is large enough.
|
||||
"""
|
||||
imgs = images[:]
|
||||
rng.shuffle(imgs)
|
||||
n = len(imgs)
|
||||
|
||||
n_test = max(1, round(n * test_frac)) if test_frac > 0 else 0
|
||||
n_val = max(1, round(n * val_frac)) if val_frac > 0 else 0
|
||||
|
||||
# Ensure we don't eat into train entirely for tiny classes
|
||||
n_test = min(n_test, max(0, n - 2))
|
||||
n_val = min(n_val, max(0, n - n_test - 1))
|
||||
|
||||
test = imgs[:n_test]
|
||||
val = imgs[n_test : n_test + n_val]
|
||||
train = imgs[n_test + n_val:]
|
||||
return train, val, test
|
||||
|
||||
|
||||
def transfer(src: Path, dst: Path, move: bool) -> None:
|
||||
dst.parent.mkdir(parents=True, exist_ok=True)
|
||||
if move:
|
||||
shutil.move(str(src), str(dst))
|
||||
else:
|
||||
shutil.copy2(str(src), str(dst))
|
||||
|
||||
|
||||
# ─────────────────────────── main ────────────────────────────────────────────
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
rng = random.Random(args.seed)
|
||||
|
||||
data_dir = Path(args.data_dir)
|
||||
out_dir = Path(args.output_dir)
|
||||
|
||||
if not data_dir.exists():
|
||||
raise FileNotFoundError(f"data_dir not found: {data_dir}")
|
||||
|
||||
# ── Collect classes ───────────────────────────────────────────────────────
|
||||
print(f"\nScanning: {data_dir.resolve()}")
|
||||
classes = collect_classes(data_dir)
|
||||
|
||||
if not classes:
|
||||
raise RuntimeError(
|
||||
f"No labelled class folders found in {data_dir}.\n"
|
||||
"Create subfolders named after your product classes and move crops into them."
|
||||
)
|
||||
|
||||
num_classes = len(classes)
|
||||
total_imgs = sum(len(v) for v in classes.values())
|
||||
print(f"\nFound {num_classes} classes | {total_imgs:,} total images\n")
|
||||
|
||||
# ── Split each class separately (stratified) ──────────────────────────────
|
||||
split_counts: dict[str, dict[str, int]] = defaultdict(lambda: {"train": 0, "val": 0, "test": 0})
|
||||
warnings = []
|
||||
|
||||
for cls_name, images in classes.items():
|
||||
train, val, test = stratified_split(images, args.val_split, args.test_split, rng)
|
||||
|
||||
for split_name, split_imgs in [("train", train), ("val", val), ("test", test)]:
|
||||
if not split_imgs:
|
||||
continue
|
||||
for src in split_imgs:
|
||||
dst = out_dir / split_name / cls_name / src.name
|
||||
transfer(src, dst, move=args.move)
|
||||
split_counts[cls_name][split_name] = len(split_imgs)
|
||||
|
||||
if len(train) < args.min_per_class:
|
||||
warnings.append(
|
||||
f" ⚠ '{cls_name}' has only {len(train)} training image(s) — "
|
||||
"consider collecting more data"
|
||||
)
|
||||
|
||||
# ── Write class manifest ──────────────────────────────────────────────────
|
||||
class_list = sorted(classes.keys())
|
||||
class_map = {i: c for i, c in enumerate(class_list)}
|
||||
manifest_path = out_dir / "class_names.json"
|
||||
manifest_path.write_text(json.dumps(class_map, indent=2))
|
||||
|
||||
# ── Print summary table ───────────────────────────────────────────────────
|
||||
col_w = max(len(c) for c in classes) + 2
|
||||
|
||||
header = f" {'Class':<{col_w}} {'Train':>7} {'Val':>7} {'Test':>7} {'Total':>7}"
|
||||
print("=" * len(header))
|
||||
print(header)
|
||||
print("-" * len(header))
|
||||
|
||||
g_train = g_val = g_test = 0
|
||||
for cls_name in sorted(split_counts):
|
||||
c = split_counts[cls_name]
|
||||
tot = c["train"] + c["val"] + c["test"]
|
||||
print(f" {cls_name:<{col_w}} {c['train']:>7} {c['val']:>7} {c['test']:>7} {tot:>7}")
|
||||
g_train += c["train"]; g_val += c["val"]; g_test += c["test"]
|
||||
|
||||
print("-" * len(header))
|
||||
print(f" {'TOTAL':<{col_w}} {g_train:>7} {g_val:>7} {g_test:>7} {g_train+g_val+g_test:>7}")
|
||||
print("=" * len(header))
|
||||
|
||||
if warnings:
|
||||
print()
|
||||
for w in warnings:
|
||||
print(w)
|
||||
|
||||
action = "moved" if args.move else "copied"
|
||||
print(f"\n {action.capitalize()} to : {out_dir.resolve()}")
|
||||
print(f" class_names.json: {manifest_path}")
|
||||
print(f"\n Next step:")
|
||||
print(f" python 2_train_classifier.py --data_dir {out_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
1
2_product-classifier.ipynb
Normal file
1
2_product-classifier.ipynb
Normal file
File diff suppressed because one or more lines are too long
359
2_train_classifier.py
Normal file
359
2_train_classifier.py
Normal file
@ -0,0 +1,359 @@
|
||||
"""
|
||||
Script 2 — Classification Model Trainer
|
||||
=========================================
|
||||
Trains a lightweight EfficientNet-B0 (ImageNet pre-trained) on the
|
||||
crop dataset produced by script 1.
|
||||
|
||||
Expected dataset layout
|
||||
-----------------------
|
||||
crops_dataset/
|
||||
train/
|
||||
cola_can/
|
||||
pepsi_can/
|
||||
lays_chips/
|
||||
…
|
||||
val/
|
||||
cola_can/
|
||||
…
|
||||
test/ (optional)
|
||||
…
|
||||
|
||||
Outputs
|
||||
-------
|
||||
runs/classify/
|
||||
best.pt ← best checkpoint (val accuracy)
|
||||
last.pt ← last checkpoint
|
||||
class_names.json ← {idx: class_name} mapping used at inference
|
||||
training_curves.png
|
||||
|
||||
Usage
|
||||
-----
|
||||
python 2_train_classifier.py \
|
||||
--data_dir crops_dataset \
|
||||
--output_dir runs/classify \
|
||||
--model efficientnet_b0 \
|
||||
--epochs 50 \
|
||||
--batch_size 64 \
|
||||
--img_size 224 \
|
||||
--lr 1e-3 \
|
||||
--patience 10 \
|
||||
--workers 4 \
|
||||
--amp
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.data import DataLoader
|
||||
from torchvision import datasets, models, transforms
|
||||
from torchvision.models import (
|
||||
EfficientNet_B0_Weights,
|
||||
EfficientNet_B2_Weights,
|
||||
MobileNet_V3_Small_Weights,
|
||||
ResNet50_1x64d_QuantizedWeights,
|
||||
)
|
||||
from torch.cuda.amp import GradScaler, autocast
|
||||
import matplotlib.pyplot as plt
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
# ─────────────────────────── args ────────────────────────────────────────────
|
||||
|
||||
SUPPORTED_MODELS = [
|
||||
"efficientnet_b0",
|
||||
"efficientnet_b2",
|
||||
"mobilenet_v3_small",
|
||||
"resnet50",
|
||||
]
|
||||
|
||||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser(description="Train a product classifier on crop dataset")
|
||||
p.add_argument("--data_dir", required=True, help="Root of crops_dataset/ with train/ val/ subfolders")
|
||||
p.add_argument("--output_dir", default="runs/classify", help="Where to save checkpoints & logs")
|
||||
p.add_argument("--model", default="efficientnet_b0", choices=SUPPORTED_MODELS)
|
||||
p.add_argument("--epochs", type=int, default=50)
|
||||
p.add_argument("--batch_size", type=int, default=64)
|
||||
p.add_argument("--img_size", type=int, default=224)
|
||||
p.add_argument("--lr", type=float, default=1e-3)
|
||||
p.add_argument("--weight_decay",type=float, default=1e-4)
|
||||
p.add_argument("--patience", type=int, default=10, help="Early-stop patience (epochs)")
|
||||
p.add_argument("--workers", type=int, default=4)
|
||||
p.add_argument("--amp", action="store_true", help="Use Automatic Mixed Precision (CUDA only)")
|
||||
p.add_argument("--freeze_backbone", action="store_true",
|
||||
help="Freeze backbone, only train the classifier head")
|
||||
p.add_argument("--seed", type=int, default=42)
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
# ─────────────────────────── model factory ───────────────────────────────────
|
||||
|
||||
def build_model(name: str, num_classes: int, freeze_backbone: bool) -> nn.Module:
|
||||
if name == "efficientnet_b0":
|
||||
m = models.efficientnet_b0(weights=EfficientNet_B0_Weights.IMAGENET1K_V1)
|
||||
if freeze_backbone:
|
||||
for p in m.features.parameters():
|
||||
p.requires_grad = False
|
||||
in_features = m.classifier[1].in_features
|
||||
m.classifier[1] = nn.Linear(in_features, num_classes)
|
||||
|
||||
elif name == "efficientnet_b2":
|
||||
m = models.efficientnet_b2(weights=EfficientNet_B2_Weights.IMAGENET1K_V1)
|
||||
if freeze_backbone:
|
||||
for p in m.features.parameters():
|
||||
p.requires_grad = False
|
||||
in_features = m.classifier[1].in_features
|
||||
m.classifier[1] = nn.Linear(in_features, num_classes)
|
||||
|
||||
elif name == "mobilenet_v3_small":
|
||||
m = models.mobilenet_v3_small(weights=MobileNet_V3_Small_Weights.IMAGENET1K_V1)
|
||||
if freeze_backbone:
|
||||
for p in m.features.parameters():
|
||||
p.requires_grad = False
|
||||
in_features = m.classifier[3].in_features
|
||||
m.classifier[3] = nn.Linear(in_features, num_classes)
|
||||
|
||||
elif name == "resnet50":
|
||||
m = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
|
||||
if freeze_backbone:
|
||||
for name_, p in m.named_parameters():
|
||||
if "fc" not in name_:
|
||||
p.requires_grad = False
|
||||
m.fc = nn.Linear(m.fc.in_features, num_classes)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unknown model: {name}")
|
||||
|
||||
return m
|
||||
|
||||
|
||||
# ─────────────────────────── data ────────────────────────────────────────────
|
||||
|
||||
def build_loaders(data_dir: Path, img_size: int, batch_size: int, workers: int):
|
||||
mean = [0.485, 0.456, 0.406]
|
||||
std = [0.229, 0.224, 0.225]
|
||||
|
||||
train_tf = transforms.Compose([
|
||||
transforms.RandomResizedCrop(img_size, scale=(0.7, 1.0)),
|
||||
transforms.RandomHorizontalFlip(),
|
||||
transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.05),
|
||||
transforms.RandomGrayscale(p=0.05),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean, std),
|
||||
])
|
||||
val_tf = transforms.Compose([
|
||||
transforms.Resize(int(img_size * 1.15)),
|
||||
transforms.CenterCrop(img_size),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean, std),
|
||||
])
|
||||
|
||||
train_set = datasets.ImageFolder(str(data_dir / "train"), transform=train_tf)
|
||||
val_set = datasets.ImageFolder(str(data_dir / "val"), transform=val_tf)
|
||||
|
||||
# ── Weighted sampler — balances classes regardless of image count per class.
|
||||
# Rare classes get oversampled, dominant ones get undersampled.
|
||||
# If your dataset is already balanced this is a no-op in effect.
|
||||
import numpy as np
|
||||
from torch.utils.data import WeightedRandomSampler
|
||||
targets = train_set.targets
|
||||
class_counts = np.bincount(targets)
|
||||
class_weights = 1.0 / class_counts
|
||||
sample_weights = class_weights[targets]
|
||||
sampler = WeightedRandomSampler(
|
||||
weights = sample_weights,
|
||||
num_samples = len(train_set),
|
||||
replacement = True,
|
||||
)
|
||||
|
||||
# ── Imbalance report ──────────────────────────────────────────────────────
|
||||
ratio = class_counts.max() / max(class_counts.min(), 1)
|
||||
if ratio > 3:
|
||||
print(f" ⚠ Class imbalance detected — {ratio:.1f}x ratio "
|
||||
f"(max {class_counts.max()} / min {class_counts.min()} images).")
|
||||
print(" WeightedRandomSampler enabled to compensate.")
|
||||
low = [train_set.classes[i] for i, c in enumerate(class_counts) if c < 20]
|
||||
if low:
|
||||
print(f" ⚠ Classes with <20 training images: {low}")
|
||||
print(" Consider collecting more data or using --freeze_backbone.")
|
||||
|
||||
train_loader = DataLoader(train_set, batch_size=batch_size, sampler=sampler,
|
||||
num_workers=workers, pin_memory=True)
|
||||
val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False,
|
||||
num_workers=workers, pin_memory=True)
|
||||
|
||||
return train_loader, val_loader, train_set.classes
|
||||
|
||||
|
||||
# ─────────────────────────── training loop ───────────────────────────────────
|
||||
|
||||
def run_epoch(model, loader, criterion, optimizer, device, scaler, train: bool):
|
||||
model.train(train)
|
||||
total_loss, correct, total = 0.0, 0, 0
|
||||
|
||||
ctx = torch.enable_grad() if train else torch.no_grad()
|
||||
with ctx:
|
||||
for imgs, labels in tqdm(loader, desc="train" if train else "val ", leave=False):
|
||||
imgs, labels = imgs.to(device), labels.to(device)
|
||||
|
||||
with autocast(enabled=(scaler is not None)):
|
||||
logits = model(imgs)
|
||||
loss = criterion(logits, labels)
|
||||
|
||||
if train:
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
if scaler:
|
||||
scaler.scale(loss).backward()
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
else:
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
total_loss += loss.item() * imgs.size(0)
|
||||
correct += (logits.argmax(1) == labels).sum().item()
|
||||
total += imgs.size(0)
|
||||
|
||||
return total_loss / total, correct / total
|
||||
|
||||
|
||||
def plot_curves(history: dict, out_path: Path):
|
||||
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
|
||||
epochs = range(1, len(history["train_loss"]) + 1)
|
||||
|
||||
ax1.plot(epochs, history["train_loss"], label="Train")
|
||||
ax1.plot(epochs, history["val_loss"], label="Val")
|
||||
ax1.set_title("Loss"); ax1.set_xlabel("Epoch"); ax1.legend()
|
||||
|
||||
ax2.plot(epochs, history["train_acc"], label="Train")
|
||||
ax2.plot(epochs, history["val_acc"], label="Val")
|
||||
ax2.set_title("Accuracy"); ax2.set_xlabel("Epoch"); ax2.legend()
|
||||
|
||||
fig.tight_layout()
|
||||
fig.savefig(str(out_path), dpi=150)
|
||||
plt.close(fig)
|
||||
print(f" Curves saved → {out_path}")
|
||||
|
||||
|
||||
# ─────────────────────────── main ────────────────────────────────────────────
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
torch.manual_seed(args.seed)
|
||||
|
||||
out_dir = Path(args.output_dir)
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
data_dir = Path(args.data_dir)
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
print(f"Device: {device}")
|
||||
|
||||
# Data
|
||||
print("Loading dataset …")
|
||||
train_loader, val_loader, class_names = build_loaders(
|
||||
data_dir, args.img_size, args.batch_size, args.workers
|
||||
)
|
||||
num_classes = len(class_names)
|
||||
print(f" Classes ({num_classes}): {class_names}")
|
||||
|
||||
# Save class names
|
||||
class_map = {i: c for i, c in enumerate(class_names)}
|
||||
with open(out_dir / "class_names.json", "w") as f:
|
||||
json.dump(class_map, f, indent=2)
|
||||
print(f" class_names.json → {out_dir / 'class_names.json'}")
|
||||
|
||||
# Model
|
||||
print(f"\nBuilding model: {args.model}")
|
||||
model = build_model(args.model, num_classes, args.freeze_backbone).to(device)
|
||||
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
total = sum(p.numel() for p in model.parameters())
|
||||
print(f" Trainable params: {trainable:,} / {total:,}")
|
||||
|
||||
# Optimiser & scheduler
|
||||
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
|
||||
optimizer = torch.optim.AdamW(
|
||||
filter(lambda p: p.requires_grad, model.parameters()),
|
||||
lr=args.lr, weight_decay=args.weight_decay
|
||||
)
|
||||
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)
|
||||
scaler = GradScaler() if (args.amp and device.type == "cuda") else None
|
||||
|
||||
# Training
|
||||
history = {"train_loss": [], "train_acc": [], "val_loss": [], "val_acc": []}
|
||||
best_val_acc = 0.0
|
||||
patience_ctr = 0
|
||||
|
||||
print(f"\n{'─'*55}")
|
||||
print(f" Starting training for {args.epochs} epochs")
|
||||
print(f"{'─'*55}")
|
||||
|
||||
for epoch in range(1, args.epochs + 1):
|
||||
t0 = time.time()
|
||||
|
||||
train_loss, train_acc = run_epoch(model, train_loader, criterion, optimizer, device, scaler, train=True)
|
||||
val_loss, val_acc = run_epoch(model, val_loader, criterion, None, device, None, train=False)
|
||||
|
||||
scheduler.step()
|
||||
elapsed = time.time() - t0
|
||||
|
||||
history["train_loss"].append(train_loss)
|
||||
history["train_acc"].append(train_acc)
|
||||
history["val_loss"].append(val_loss)
|
||||
history["val_acc"].append(val_acc)
|
||||
|
||||
improved = val_acc > best_val_acc
|
||||
tag = " ✓ best" if improved else ""
|
||||
print(
|
||||
f"Epoch {epoch:3d}/{args.epochs} | "
|
||||
f"Train loss {train_loss:.4f} acc {train_acc:.4f} | "
|
||||
f"Val loss {val_loss:.4f} acc {val_acc:.4f} | "
|
||||
f"{elapsed:.1f}s{tag}"
|
||||
)
|
||||
|
||||
# Save checkpoints
|
||||
if improved:
|
||||
best_val_acc = val_acc
|
||||
patience_ctr = 0
|
||||
torch.save({
|
||||
"epoch": epoch,
|
||||
"model_state": model.state_dict(),
|
||||
"class_names": class_names,
|
||||
"val_acc": val_acc,
|
||||
"args": vars(args),
|
||||
}, out_dir / "best.pt")
|
||||
else:
|
||||
patience_ctr += 1
|
||||
|
||||
torch.save({
|
||||
"epoch": epoch,
|
||||
"model_state": model.state_dict(),
|
||||
"class_names": class_names,
|
||||
"val_acc": val_acc,
|
||||
"args": vars(args),
|
||||
}, out_dir / "last.pt")
|
||||
|
||||
# Early stopping
|
||||
if patience_ctr >= args.patience:
|
||||
print(f"\n Early stopping triggered (no improvement for {args.patience} epochs)")
|
||||
break
|
||||
|
||||
plot_curves(history, out_dir / "training_curves.png")
|
||||
|
||||
print("\n" + "═" * 55)
|
||||
print(" Training complete!")
|
||||
print(f" Best val accuracy : {best_val_acc:.4f}")
|
||||
print(f" Checkpoints saved : {out_dir}")
|
||||
print("═" * 55)
|
||||
print("\n ➜ Next step:")
|
||||
print(" python 3_inference.py \\")
|
||||
print(f" --weights {out_dir / 'best.pt'} \\")
|
||||
print(" --source /path/to/images \\")
|
||||
print(" --detector_weights /path/to/best_yolo.pt")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
464
3_inference.py
Normal file
464
3_inference.py
Normal file
@ -0,0 +1,464 @@
|
||||
"""
|
||||
Script 3 — End-to-end Inference (Detect → Classify)
|
||||
=====================================================
|
||||
1. Runs your YOLOv8n detector on each image / video frame.
|
||||
2. Crops every detection.
|
||||
3. Feeds each crop through the classification head.
|
||||
4. Draws annotated bounding boxes with class names & confidence.
|
||||
|
||||
Supported sources
|
||||
-----------------
|
||||
• A single image (--source image.jpg)
|
||||
• A folder of images (--source /path/to/imgs)
|
||||
• A video file (--source video.mp4)
|
||||
• A webcam (--source 0)
|
||||
|
||||
Usage
|
||||
-----
|
||||
python 3_inference.py \
|
||||
--detector_weights runs/yolo/best.pt \
|
||||
--classifier_weights runs/classify/best.pt \
|
||||
--source /path/to/images_or_video \
|
||||
--det_conf 0.3 \
|
||||
--cls_conf 0.5 \
|
||||
--img_size 640 \
|
||||
--padding 8 \
|
||||
--output_dir inference_results \
|
||||
--show \
|
||||
--save_crops
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torchvision import models, transforms
|
||||
from ultralytics import YOLO
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
# ─────────────────────────── args ────────────────────────────────────────────
|
||||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser(description="YOLOv8 detect + classify inference")
|
||||
p.add_argument("--detector_weights", default="detector/best.pt", help="YOLOv8n .pt (fine-tuned)")
|
||||
p.add_argument("--classifier_weights", default="best.pt", help="Classifier best.pt from script 2")
|
||||
p.add_argument("--source", required=True, help="Image / folder / video / webcam index")
|
||||
p.add_argument("--det_conf", type=float, default=0.30, help="YOLO detection confidence")
|
||||
p.add_argument("--det_iou", type=float, default=0.45, help="YOLO NMS IoU")
|
||||
p.add_argument("--cls_conf", type=float, default=0.60, help="Min classifier confidence to label")
|
||||
p.add_argument("--img_size", type=int, default=640, help="YOLO input size")
|
||||
p.add_argument("--padding", type=int, default=8, help="Extra pixels around each crop")
|
||||
p.add_argument("--cls_img_size",type=int, default=224, help="Classifier input size")
|
||||
p.add_argument("--output_dir", default="inference_results")
|
||||
p.add_argument("--show", action="store_true", help="Display results with cv2.imshow")
|
||||
p.add_argument("--save_crops", action="store_true", help="Save individual crop images")
|
||||
p.add_argument("--no_save", action="store_true", help="Do not save annotated images/video")
|
||||
p.add_argument("--batch_cls", type=int, default=16, help="Classifier batch size per frame")
|
||||
p.add_argument("--device", default="", help="cuda / cpu / mps (auto-detect if empty)")
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
# ─────────────────────────── classifier wrapper ──────────────────────────────
|
||||
|
||||
# ── Device-agnostic OOM detection ────────────────────────────────────────────
|
||||
#
|
||||
# Each backend raises a different exception type for out-of-memory:
|
||||
#
|
||||
# CUDA → torch.cuda.OutOfMemoryError (subclass of RuntimeError)
|
||||
# MPS → RuntimeError whose message contains "out of memory"
|
||||
# CPU → MemoryError (Python built-in, raised by the OS allocator)
|
||||
#
|
||||
# We centralise the check here so nothing else needs to know the device type.
|
||||
|
||||
def _is_oom(exc: BaseException) -> bool:
|
||||
"""Return True if *exc* represents an out-of-memory condition."""
|
||||
if isinstance(exc, MemoryError): # CPU / OS
|
||||
return True
|
||||
if torch.cuda.is_available() and isinstance(exc, torch.cuda.OutOfMemoryError):
|
||||
return True # CUDA
|
||||
if isinstance(exc, RuntimeError): # MPS + fallback CUDA
|
||||
msg = str(exc).lower()
|
||||
return "out of memory" in msg or "memory" in msg and "alloc" in msg
|
||||
return False
|
||||
|
||||
|
||||
def _free_device_cache(device: torch.device) -> None:
|
||||
"""Release any cached memory held by the current device allocator."""
|
||||
if device.type == "cuda":
|
||||
torch.cuda.empty_cache()
|
||||
elif device.type == "mps":
|
||||
# torch.mps.empty_cache() was added in PyTorch 2.1 — guard for older versions
|
||||
if hasattr(torch.mps, "empty_cache"):
|
||||
torch.mps.empty_cache()
|
||||
# CPU has no cache to release
|
||||
|
||||
|
||||
def _sync_device(device: torch.device) -> None:
|
||||
"""Wait for all pending ops on *device* so OOM surfaces immediately."""
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize()
|
||||
elif device.type == "mps":
|
||||
if hasattr(torch.mps, "synchronize"):
|
||||
torch.mps.synchronize()
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
class ProductClassifier:
|
||||
"""Thin wrapper around the saved checkpoint from script 2."""
|
||||
|
||||
# Supported architectures — mirrors script 2
|
||||
_BUILDERS = {
|
||||
"efficientnet_b0": (models.efficientnet_b0, lambda m, n: setattr(m.classifier, '1', torch.nn.Linear(m.classifier[1].in_features if hasattr(m.classifier[1], 'in_features') else 1280, n))),
|
||||
"efficientnet_b2": (models.efficientnet_b2, lambda m, n: setattr(m.classifier, '1', torch.nn.Linear(m.classifier[1].in_features if hasattr(m.classifier[1], 'in_features') else 1408, n))),
|
||||
"mobilenet_v3_small": (models.mobilenet_v3_small, lambda m, n: setattr(m.classifier, '3', torch.nn.Linear(m.classifier[3].in_features if hasattr(m.classifier[3], 'in_features') else 576, n))),
|
||||
"resnet50": (models.resnet50, lambda m, n: setattr(m, 'fc', torch.nn.Linear(m.fc.in_features, n))),
|
||||
}
|
||||
|
||||
def __init__(self, weights_path: str, device: torch.device, img_size: int = 224):
|
||||
ckpt = torch.load(weights_path, map_location=device)
|
||||
self.class_names: list = ckpt["class_names"]
|
||||
num_classes = len(self.class_names)
|
||||
|
||||
model_name = ckpt.get("args", {}).get("model", "efficientnet_b0")
|
||||
builder_fn, head_fn = self._BUILDERS[model_name]
|
||||
self.model = builder_fn(weights=None)
|
||||
head_fn(self.model, num_classes)
|
||||
self.model.load_state_dict(ckpt["model_state"])
|
||||
self.model.to(device).eval()
|
||||
self.device = device
|
||||
|
||||
mean = [0.485, 0.456, 0.406]
|
||||
std = [0.229, 0.224, 0.225]
|
||||
self.transform = transforms.Compose([
|
||||
transforms.ToPILImage(),
|
||||
transforms.Resize(int(img_size * 1.15)),
|
||||
transforms.CenterCrop(img_size),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean, std),
|
||||
])
|
||||
|
||||
# Probe VRAM once at startup — stored and reused every frame.
|
||||
# Can be overridden at runtime if OOM occurs (see predict_batch).
|
||||
self._batch_size: int = self.calibrate_batch_size(img_size)
|
||||
|
||||
def calibrate_batch_size(self, img_size: int, start: int = 2, max_size: int = 512) -> int:
|
||||
"""
|
||||
Probe the device at startup to find the largest batch that fits in memory.
|
||||
|
||||
Works on CUDA, MPS, and CPU:
|
||||
• CUDA / MPS — doubles the batch size until OOM, steps back with margin.
|
||||
• CPU — RAM is virtually unlimited; returns a sensible fixed cap.
|
||||
|
||||
Returns the largest safe batch size, always ≥ 1.
|
||||
"""
|
||||
if self.device.type == "cpu":
|
||||
print(" [batch-probe] CPU detected — using fixed batch size 64")
|
||||
return 64
|
||||
|
||||
print(f" [batch-probe] Probing {self.device.type.upper()} memory "
|
||||
f"for optimal batch size (img={img_size}px) …", flush=True)
|
||||
|
||||
dummy = torch.zeros(1, 3, img_size, img_size, device=self.device)
|
||||
safe_size = 1
|
||||
probe_size = start
|
||||
|
||||
with torch.no_grad():
|
||||
while probe_size <= max_size:
|
||||
try:
|
||||
batch = dummy.expand(probe_size, -1, -1, -1)
|
||||
_ = self.model(batch)
|
||||
|
||||
# Flush pending async ops so OOM surfaces here, not later
|
||||
_sync_device(self.device)
|
||||
|
||||
safe_size = probe_size
|
||||
probe_size = probe_size * 2
|
||||
|
||||
except Exception as exc:
|
||||
if _is_oom(exc):
|
||||
break # safe_size is our answer
|
||||
raise # unexpected error — propagate
|
||||
|
||||
finally:
|
||||
_free_device_cache(self.device)
|
||||
|
||||
# 20 % headroom margin: real crops vary in size unlike the uniform dummy
|
||||
safe_size = max(1, int(safe_size * 0.8))
|
||||
print(f" [batch-probe] Settled on batch_size = {safe_size} "
|
||||
f"(probe ceiling was {probe_size // 2})")
|
||||
return safe_size
|
||||
|
||||
@torch.no_grad()
|
||||
def predict_batch(self, crops_bgr: list, batch_size: int | None = None) -> list[tuple[str, float]]:
|
||||
"""
|
||||
Classifies an arbitrary list of crops in memory-safe chunks.
|
||||
|
||||
crops_bgr : list of H×W×3 uint8 BGR numpy arrays
|
||||
batch_size : chunk size to use. Pass None to use self._batch_size
|
||||
(set by calibrate_batch_size at startup).
|
||||
Returns : list of (class_name, confidence) — same order as input
|
||||
"""
|
||||
if not crops_bgr:
|
||||
return []
|
||||
|
||||
# Resolve which batch size to use
|
||||
chunk_size = batch_size if batch_size is not None else self._batch_size
|
||||
|
||||
# Pre-process all crops on CPU (fast; no GPU memory involved yet)
|
||||
tensors = [
|
||||
self.transform(cv2.cvtColor(crop, cv2.COLOR_BGR2RGB))
|
||||
for crop in crops_bgr
|
||||
]
|
||||
|
||||
results: list[tuple[str, float]] = []
|
||||
|
||||
i = 0
|
||||
while i < len(tensors):
|
||||
chunk = tensors[i : i + chunk_size]
|
||||
batch = torch.stack(chunk).to(self.device)
|
||||
|
||||
try:
|
||||
logits = self.model(batch)
|
||||
probs = F.softmax(logits, dim=1)
|
||||
confs, idxs = probs.max(dim=1)
|
||||
|
||||
for idx, conf in zip(idxs, confs):
|
||||
results.append((self.class_names[idx.item()], conf.item()))
|
||||
|
||||
i += chunk_size # advance only on success
|
||||
|
||||
except Exception as exc:
|
||||
if not _is_oom(exc):
|
||||
raise # unexpected error — don't swallow it
|
||||
|
||||
# ── Runtime OOM safety net ────────────────────────────────
|
||||
# The probe uses a uniform dummy; real crops vary and can
|
||||
# occasionally exceed the probed limit. Halve and retry
|
||||
# WITHOUT advancing i so the same chunk is retried.
|
||||
_free_device_cache(self.device)
|
||||
|
||||
new_size = max(1, chunk_size // 2)
|
||||
print(f" [batch] Runtime OOM on {self.device.type.upper()} "
|
||||
f"— halving chunk {chunk_size} → {new_size}")
|
||||
chunk_size = new_size
|
||||
self._batch_size = new_size # persist for future frames
|
||||
|
||||
finally:
|
||||
del batch # free GPU memory immediately after each chunk
|
||||
|
||||
return results
|
||||
|
||||
|
||||
# ─────────────────────────── drawing ─────────────────────────────────────────
|
||||
|
||||
# One colour per class, generated on first use
|
||||
_COLOUR_CACHE: dict[str, tuple] = {}
|
||||
|
||||
|
||||
def class_colour(name: str) -> tuple:
|
||||
if name not in _COLOUR_CACHE:
|
||||
# OpenCV HSV: hue is 0–179 (NOT 0–360) — values above 179 overflow uint8
|
||||
h = hash(name) % 180
|
||||
col = cv2.cvtColor(
|
||||
np.array([[[h, 220, 200]]], dtype=np.uint8), cv2.COLOR_HSV2BGR
|
||||
)[0][0]
|
||||
_COLOUR_CACHE[name] = tuple(int(c) for c in col)
|
||||
return _COLOUR_CACHE[name]
|
||||
|
||||
|
||||
def draw_result(frame, box_xyxy, label: str, conf: float, cls_conf_thresh: float):
|
||||
x1, y1, x2, y2 = map(int, box_xyxy)
|
||||
colour = class_colour(label)
|
||||
cv2.rectangle(frame, (x1, y1), (x2, y2), colour, 2)
|
||||
|
||||
if conf >= cls_conf_thresh:
|
||||
text = f"{label} {conf:.2f}"
|
||||
else:
|
||||
text = "?"
|
||||
|
||||
(tw, th), bl = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.55, 1)
|
||||
cv2.rectangle(frame, (x1, y1 - th - bl - 4), (x1 + tw + 4, y1), colour, -1)
|
||||
cv2.putText(frame, text, (x1 + 2, y1 - bl - 2),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 255), 1, cv2.LINE_AA)
|
||||
return frame
|
||||
|
||||
|
||||
# ─────────────────────────── source helpers ──────────────────────────────────
|
||||
|
||||
IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tiff"}
|
||||
|
||||
|
||||
def is_image(path: Path) -> bool:
|
||||
return path.suffix.lower() in IMAGE_EXTS
|
||||
|
||||
|
||||
def is_video(path: Path) -> bool:
|
||||
return path.suffix.lower() in {".mp4", ".avi", ".mov", ".mkv", ".webm"}
|
||||
|
||||
|
||||
def iter_source(source: str):
|
||||
"""
|
||||
Yields (frame_bgr, frame_id, source_name) for every frame / image.
|
||||
"""
|
||||
try:
|
||||
cam_idx = int(source) # webcam
|
||||
cap = cv2.VideoCapture(cam_idx)
|
||||
fid = 0
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
break
|
||||
yield frame, fid, f"webcam_{cam_idx}"
|
||||
fid += 1
|
||||
cap.release()
|
||||
return
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
p = Path(source)
|
||||
if p.is_file() and is_image(p):
|
||||
img = cv2.imread(str(p))
|
||||
if img is not None:
|
||||
yield img, 0, p.stem
|
||||
elif p.is_file() and is_video(p):
|
||||
cap = cv2.VideoCapture(str(p))
|
||||
fid = 0
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
break
|
||||
yield frame, fid, p.stem
|
||||
fid += 1
|
||||
cap.release()
|
||||
elif p.is_dir():
|
||||
img_paths = sorted(q for q in p.rglob("*") if is_image(q))
|
||||
for q in tqdm(img_paths, desc="Images"):
|
||||
img = cv2.imread(str(q))
|
||||
if img is not None:
|
||||
yield img, 0, q.stem
|
||||
else:
|
||||
raise FileNotFoundError(f"Cannot open source: {source}")
|
||||
|
||||
|
||||
# ─────────────────────────── main ────────────────────────────────────────────
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
# Device
|
||||
if args.device:
|
||||
device = torch.device(args.device)
|
||||
elif torch.cuda.is_available():
|
||||
device = torch.device("cuda")
|
||||
elif torch.backends.mps.is_available():
|
||||
device = torch.device("mps")
|
||||
else:
|
||||
device = torch.device("cpu")
|
||||
print(f"Device : {device}")
|
||||
|
||||
# Load models
|
||||
print(f"Loading detector : {args.detector_weights}")
|
||||
detector = YOLO(args.detector_weights)
|
||||
|
||||
print(f"Loading classifier: {args.classifier_weights}")
|
||||
classifier = ProductClassifier(args.classifier_weights, device, args.cls_img_size)
|
||||
print(f" Classes ({len(classifier.class_names)}): {classifier.class_names}")
|
||||
|
||||
# Output directory
|
||||
out_dir = Path(args.output_dir)
|
||||
if not args.no_save:
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
crops_dir = out_dir / "crops" if args.save_crops else None
|
||||
if crops_dir:
|
||||
crops_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Stats
|
||||
stats = {"frames": 0, "detections": 0, "classified": 0, "fps": []}
|
||||
|
||||
for frame, frame_id, stem in iter_source(args.source):
|
||||
t0 = time.perf_counter()
|
||||
|
||||
# ── Detection ──────────────────────────────────────────────────────
|
||||
results = detector(
|
||||
frame,
|
||||
conf=args.det_conf,
|
||||
iou=args.det_iou,
|
||||
imgsz=args.img_size,
|
||||
verbose=False,
|
||||
)[0]
|
||||
|
||||
h, w = frame.shape[:2]
|
||||
crops, boxes = [], []
|
||||
|
||||
for box in results.boxes:
|
||||
x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
|
||||
x1 = max(0, x1 - args.padding); y1 = max(0, y1 - args.padding)
|
||||
x2 = min(w, x2 + args.padding); y2 = min(h, y2 + args.padding)
|
||||
crop = frame[y1:y2, x1:x2]
|
||||
if crop.size == 0:
|
||||
continue
|
||||
crops.append(crop)
|
||||
boxes.append((x1, y1, x2, y2))
|
||||
|
||||
# ── Classification (dynamic chunked batches) ───────────────────────
|
||||
# Batch size is auto-calibrated to VRAM at startup via probe,
|
||||
# and auto-reduces at runtime if OOM occurs (e.g. unusually dense frame).
|
||||
predictions = classifier.predict_batch(crops)
|
||||
|
||||
# ── Draw & annotate ───────────────────────────────────────────────
|
||||
annotated = frame.copy()
|
||||
for (x1, y1, x2, y2), (cls_name, cls_conf) in zip(boxes, predictions):
|
||||
draw_result(annotated, (x1, y1, x2, y2), cls_name, cls_conf, args.cls_conf)
|
||||
if cls_conf >= args.cls_conf:
|
||||
stats["classified"] += 1
|
||||
|
||||
# FPS overlay
|
||||
elapsed = time.perf_counter() - t0
|
||||
fps = 1.0 / elapsed if elapsed > 0 else 0
|
||||
stats["fps"].append(fps)
|
||||
cv2.putText(annotated, f"FPS: {fps:.1f} Det: {len(boxes)}",
|
||||
(10, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
|
||||
|
||||
# ── Save / display ────────────────────────────────────────────────
|
||||
if not args.no_save:
|
||||
out_path = out_dir / f"{stem}_frame{frame_id:06d}.jpg"
|
||||
cv2.imwrite(str(out_path), annotated)
|
||||
|
||||
if args.save_crops and crops_dir:
|
||||
for i, (crop, (cls_name, cls_conf)) in enumerate(zip(crops, predictions)):
|
||||
tag = cls_name if cls_conf >= args.cls_conf else "uncertain"
|
||||
cv2.imwrite(str(crops_dir / f"{stem}_f{frame_id:06d}_c{i:04d}_{tag}.jpg"), crop)
|
||||
|
||||
if args.show:
|
||||
cv2.imshow("Detect + Classify", annotated)
|
||||
key = cv2.waitKey(1)
|
||||
if key in (ord("q"), 27):
|
||||
break
|
||||
|
||||
stats["frames"] += 1
|
||||
stats["detections"] += len(boxes)
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
# ── Summary ───────────────────────────────────────────────────────────
|
||||
avg_fps = np.mean(stats["fps"]) if stats["fps"] else 0
|
||||
print("\n" + "═" * 55)
|
||||
print(" Inference summary")
|
||||
print(f" Frames processed : {stats['frames']}")
|
||||
print(f" Total detections : {stats['detections']}")
|
||||
print(f" Classified (≥{args.cls_conf:.0%}): {stats['classified']}")
|
||||
print(f" Avg FPS : {avg_fps:.1f}")
|
||||
if not args.no_save:
|
||||
print(f" Results saved to : {out_dir.resolve()}")
|
||||
print("═" * 55)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
628
README.md
Normal file
628
README.md
Normal file
@ -0,0 +1,628 @@
|
||||
# Product Detection & Classification Pipeline
|
||||
|
||||
A two-stage computer vision pipeline that combines a fine-tuned **YOLOv8n** detector with a **EfficientNet** classifier to detect and identify products on retail shelves.
|
||||
|
||||
```
|
||||
Shelf image → YOLOv8n (detect) → EfficientNet (classify) → labelled bounding boxes
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Table of Contents
|
||||
|
||||
1. [Project Structure](#1-project-structure)
|
||||
2. [Installation](#2-installation)
|
||||
3. [Pipeline Overview](#3-pipeline-overview)
|
||||
4. [Step-by-step Guide](#4-step-by-step-guide)
|
||||
- [Step 1 — Generate crops](#step-1--generate-crops)
|
||||
- [Step 2 — Auto-label & Review](#step-2--auto-label--review)
|
||||
- [Step 3 — Split the dataset](#step-3--split-the-dataset)
|
||||
- [Step 4 — Train the classifier](#step-4--train-the-classifier)
|
||||
- [Step 5 — Run inference](#step-5--run-inference)
|
||||
5. [API Server](#5-api-server)
|
||||
6. [Script Reference](#6-script-reference)
|
||||
7. [Troubleshooting](#7-troubleshooting)
|
||||
8. [Tips & Best Practices](#8-tips--best-practices)
|
||||
|
||||
---
|
||||
|
||||
## 1. Project Structure
|
||||
|
||||
```
|
||||
project/
|
||||
│
|
||||
├── detector/
|
||||
│ └── best.pt ← your fine-tuned YOLOv8n weights
|
||||
│
|
||||
├── runs/classify/
|
||||
│ ├── best.pt ← trained classifier weights (output of step 4)
|
||||
│ ├── last.pt
|
||||
│ ├── class_names.json ← {index: class_name} mapping
|
||||
│ ├── training_curves.png
|
||||
│ ├── per_class_accuracy.png
|
||||
│ └── confusion_matrix.png
|
||||
│
|
||||
├── data/ ← labelled crops (built incrementally)
|
||||
│ ├── cola_can/
|
||||
│ ├── pepsi_can/
|
||||
│ ├── _rejected_/ ← crops you rejected during review
|
||||
│ └── _unreviewed_/ ← crops you skipped during review
|
||||
│
|
||||
├── crops_dataset/ ← final train/val/test split (input to training)
|
||||
│ ├── train/
|
||||
│ ├── val/
|
||||
│ └── test/
|
||||
│
|
||||
├── 1_generate_crop_dataset.py ← detect & crop (no classifier needed)
|
||||
├── auto_label.py ← detect + auto-classify + browser review UI
|
||||
├── 1b_split_dataset.py ← stratified train/val/test splitter
|
||||
├── 2_train_classifier.py ← train EfficientNet classifier
|
||||
├── 3_inference.py ← end-to-end detect + classify on new images
|
||||
├── api_server.py ← FastAPI polling server
|
||||
└── train_classifier_kaggle.ipynb
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 2. Installation
|
||||
|
||||
The project uses a **Conda environment** for reproducible dependency management.
|
||||
|
||||
### Create the environment
|
||||
|
||||
```bash
|
||||
conda env create -f environment.yml
|
||||
conda activate product-detection
|
||||
```
|
||||
|
||||
### `environment.yml`
|
||||
|
||||
```yaml
|
||||
name: product-detection
|
||||
channels:
|
||||
- conda-forge
|
||||
dependencies:
|
||||
- python=3.10
|
||||
- numpy<2
|
||||
- pip
|
||||
- pip:
|
||||
- torch
|
||||
- torchvision
|
||||
- opencv-python
|
||||
- ultralytics
|
||||
- pillow==10.0.0
|
||||
- tqdm
|
||||
- pyyaml
|
||||
- simsimd
|
||||
- albumentations
|
||||
- faiss-cpu
|
||||
- fastapi
|
||||
- uvicorn[standard]
|
||||
- flask
|
||||
```
|
||||
|
||||
### GPU users
|
||||
|
||||
`torch` and `torchvision` are installed as CPU builds by default via pip inside conda. To use a GPU, replace the `torch` and `torchvision` lines in the yml with the CUDA-enabled builds before creating the environment:
|
||||
|
||||
```yaml
|
||||
- torch==2.3.0+cu121
|
||||
- torchvision==0.18.0+cu121
|
||||
--extra-index-url https://download.pytorch.org/whl/cu121
|
||||
```
|
||||
|
||||
Or install them manually after environment creation:
|
||||
|
||||
```bash
|
||||
conda activate product-detection
|
||||
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
|
||||
```
|
||||
|
||||
### Update an existing environment
|
||||
|
||||
If `environment.yml` changes, sync the existing environment without recreating it:
|
||||
|
||||
```bash
|
||||
conda env update -f environment.yml --prune
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 3. Pipeline Overview
|
||||
|
||||
There are two ways to build your dataset depending on whether you already have a trained classifier:
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────────────────────────────┐
|
||||
│ PATH A — First time (no classifier yet) │
|
||||
│ │
|
||||
│ 1_generate_crop_dataset.py → label manually → │
|
||||
│ 1b_split_dataset.py → 2_train_classifier.py │
|
||||
└─────────────────────────────────────────────────────────────────┘
|
||||
|
||||
┌─────────────────────────────────────────────────────────────────┐
|
||||
│ PATH B — You already have a classifier (recommended after v1) │
|
||||
│ │
|
||||
│ auto_label.py → review in browser → │
|
||||
│ 1b_split_dataset.py → 2_train_classifier.py │
|
||||
└─────────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
Once you have a working classifier, always use **Path B**. It is faster, more accurate, and requires far less manual labelling work.
|
||||
|
||||
---
|
||||
|
||||
## 4. Step-by-step Guide
|
||||
|
||||
### Step 1 — Generate crops
|
||||
|
||||
> Use this step if you are starting from scratch with **no classifier yet**.
|
||||
> If you already have `runs/classify/best.pt`, skip to Step 2.
|
||||
|
||||
Runs YOLOv8n on your shelf images and saves every detection as a cropped JPEG into `data/unknown/`. No classification happens here — you label manually afterwards.
|
||||
|
||||
```bash
|
||||
python 1_generate_crop_dataset.py \
|
||||
--source shelf_images/ \
|
||||
--weights detector/best.pt \
|
||||
--output_dir data \
|
||||
--conf 0.25 \
|
||||
--padding 10 \
|
||||
--min_size 32
|
||||
```
|
||||
|
||||
| Argument | Default | Description |
|
||||
|---|---|---|
|
||||
| `--source` | required | Image file, folder, or glob pattern |
|
||||
| `--weights` | required | YOLOv8n `.pt` file |
|
||||
| `--output_dir` | `data` | Root folder — crops go into `<output_dir>/unknown/` |
|
||||
| `--conf` | `0.25` | Detection confidence threshold |
|
||||
| `--iou` | `0.45` | NMS IoU threshold |
|
||||
| `--padding` | `10` | Extra pixels added around each crop |
|
||||
| `--min_size` | `32` | Discard crops smaller than this (px) |
|
||||
| `--save_labels` | off | Save a `.txt` sidecar with YOLO class info |
|
||||
|
||||
**After this step**, move crops from `data/unknown/` into named class folders:
|
||||
|
||||
```
|
||||
data/
|
||||
unknown/ ← before
|
||||
cola_can/ ← after (you create these)
|
||||
pepsi_can/
|
||||
lays_chips/
|
||||
```
|
||||
|
||||
Use any file manager, or a tool like [Label Studio](https://labelstud.io/) for faster annotation.
|
||||
|
||||
---
|
||||
|
||||
### Step 2 — Auto-label & Review
|
||||
|
||||
> Use this step once you have a trained classifier. It replaces manual labelling almost entirely.
|
||||
|
||||
Runs the full detect → crop → classify pipeline in one command, then opens a browser UI where you review, confirm, fix, or reject each prediction.
|
||||
|
||||
```bash
|
||||
python auto_label.py \
|
||||
--source shelf_images/ \
|
||||
--detector_weights detector/best.pt \
|
||||
--classifier_weights runs/classify/best.pt \
|
||||
--output_dir data \
|
||||
--cls_conf 0.50
|
||||
```
|
||||
|
||||
| Argument | Default | Description |
|
||||
|---|---|---|
|
||||
| `--source` | required | Image file, folder, or glob |
|
||||
| `--detector_weights` | required | YOLOv8n `.pt` |
|
||||
| `--classifier_weights` | required | Classifier `best.pt` |
|
||||
| `--output_dir` | `data` | Where confirmed crops are written |
|
||||
| `--det_conf` | `0.25` | Detection threshold |
|
||||
| `--cls_conf` | `0.50` | Min classifier confidence to auto-label |
|
||||
| `--port` | `5000` | Browser UI port |
|
||||
|
||||
The browser opens automatically at `http://localhost:5000`.
|
||||
|
||||
#### Review UI controls
|
||||
|
||||
| Action | How |
|
||||
|---|---|
|
||||
| Open crop detail | Click a card |
|
||||
| Confirm prediction | `C` key or **Confirm** button |
|
||||
| Reject crop | `R` key or **Reject** button |
|
||||
| Override the class | Pick from dropdown → **Apply override** |
|
||||
| Navigate crops | `←` `→` arrow keys |
|
||||
| Close modal | `Esc` |
|
||||
| Multi-select | `Ctrl+click` or `Shift+click` |
|
||||
| Bulk confirm/reject/reclassify | Select multiple → use bulk bar |
|
||||
| Confirm all visible | **✓ Confirm all visible** button |
|
||||
| Filter by class | Click a class in the left sidebar |
|
||||
| Filter by status | Toggle buttons in toolbar |
|
||||
| Filter by confidence | Drag the min-conf slider |
|
||||
|
||||
#### Status meanings
|
||||
|
||||
| Status | Meaning |
|
||||
|---|---|
|
||||
| `auto` | Classifier's prediction, not yet reviewed |
|
||||
| `confirmed` | You agreed with the prediction |
|
||||
| `reclassified` | You changed it to a different class |
|
||||
| `rejected` | Not a useful crop (bad detection, not a product of interest) |
|
||||
|
||||
#### Committing
|
||||
|
||||
Click **Commit dataset →** to write all crops to `data/`:
|
||||
|
||||
```
|
||||
data/
|
||||
cola_can/ ← confirmed + reclassified
|
||||
pepsi_can/
|
||||
_rejected_/ ← rejected crops (kept for reference)
|
||||
_unreviewed_/ ← anything still "auto" when you committed
|
||||
```
|
||||
|
||||
> **Important:** Commit is always additive. It never deletes existing files in `data/`.
|
||||
> Running `auto_label.py` on new shelf images simply grows the dataset.
|
||||
> The only exception: if you run it twice on the **exact same source image**, crop files
|
||||
> are overwritten (same content, so this is harmless).
|
||||
|
||||
---
|
||||
|
||||
### Step 3 — Split the dataset
|
||||
|
||||
Splits `data/` into `train/`, `val/`, and `test/` using a **stratified** strategy — every class gets the same ratio across all splits, so rare classes are not accidentally lost.
|
||||
|
||||
```bash
|
||||
python 1b_split_dataset.py \
|
||||
--data_dir data \
|
||||
--output_dir crops_dataset \
|
||||
--val_split 0.15 \
|
||||
--test_split 0.05
|
||||
```
|
||||
|
||||
| Argument | Default | Description |
|
||||
|---|---|---|
|
||||
| `--data_dir` | required | Labelled data root (subfolders = classes) |
|
||||
| `--output_dir` | `crops_dataset` | Where to write train/val/test |
|
||||
| `--val_split` | `0.15` | Fraction for validation |
|
||||
| `--test_split` | `0.05` | Fraction for test (`0` to skip) |
|
||||
| `--move` | off | Move files instead of copying (saves disk space) |
|
||||
| `--min_per_class` | `2` | Warn if a class has fewer images than this after splitting |
|
||||
|
||||
**Automatically skipped folders:**
|
||||
- `unknown/`, `train/`, `val/`, `test/`
|
||||
- Any folder whose name starts with `_` (e.g. `_rejected_`, `_unreviewed_`, `_ignore_`)
|
||||
|
||||
The splitter prints a table so you can spot imbalanced or thin classes before training:
|
||||
|
||||
```
|
||||
────────────────────────────────────────────────────────
|
||||
Class Train Val Test
|
||||
────────────────────────────────────────────────────────
|
||||
cola_can 72 13 5
|
||||
pepsi_can 58 10 4
|
||||
lays_chips 12 2 ⚠ val<5
|
||||
────────────────────────────────────────────────────────
|
||||
TOTAL 142 25 9
|
||||
```
|
||||
|
||||
> **Recommended minimums:** ≥ 50 training images and ≥ 10 val images per class.
|
||||
> Classes below these thresholds should get more data before training.
|
||||
|
||||
---
|
||||
|
||||
### Step 4 — Train the classifier
|
||||
|
||||
Trains an **EfficientNet-B0** (ImageNet pre-trained) on `crops_dataset/`. A `WeightedRandomSampler` is used automatically so class imbalance does not bias training.
|
||||
|
||||
```bash
|
||||
python 2_train_classifier.py \
|
||||
--data_dir crops_dataset \
|
||||
--model efficientnet_b0 \
|
||||
--epochs 50 \
|
||||
--batch_size 64 \
|
||||
--amp
|
||||
```
|
||||
|
||||
| Argument | Default | Description |
|
||||
|---|---|---|
|
||||
| `--data_dir` | required | Root with `train/` and `val/` |
|
||||
| `--output_dir` | `runs/classify` | Where checkpoints are saved |
|
||||
| `--model` | `efficientnet_b0` | `efficientnet_b0`, `efficientnet_b2`, `mobilenet_v3_small`, `resnet50` |
|
||||
| `--epochs` | `50` | Max training epochs |
|
||||
| `--batch_size` | `64` | Batch size |
|
||||
| `--img_size` | `224` | Input resolution |
|
||||
| `--lr` | `1e-3` | Learning rate |
|
||||
| `--patience` | `10` | Early stopping patience |
|
||||
| `--amp` | off | Enable Automatic Mixed Precision (GPU only) |
|
||||
| `--freeze_backbone` | off | Only train the head — recommended when you have fewer than ~50 images per class |
|
||||
|
||||
**Outputs saved to `runs/classify/`:**
|
||||
|
||||
| File | Description |
|
||||
|---|---|
|
||||
| `best.pt` | Best checkpoint by val accuracy |
|
||||
| `last.pt` | Last epoch checkpoint |
|
||||
| `class_names.json` | `{index: class_name}` used by inference and API |
|
||||
| `training_curves.png` | Loss & accuracy plot |
|
||||
|
||||
#### Option B — Kaggle notebook (recommended for free GPU access)
|
||||
|
||||
`train_classifier_kaggle.ipynb` is a self-contained notebook version of the training script, designed to run on Kaggle's free T4 GPU with no local hardware required.
|
||||
|
||||
**Setup steps:**
|
||||
|
||||
1. Zip your `crops_dataset/` folder and upload it as a Kaggle Dataset:
|
||||
- Go to [kaggle.com/datasets](https://www.kaggle.com/datasets) → **New Dataset**
|
||||
- Upload the zip, give it a name (e.g. `my-product-crops`)
|
||||
- Wait for processing to complete
|
||||
|
||||
2. Create a new Kaggle Notebook:
|
||||
- Go to [kaggle.com/code](https://www.kaggle.com/code) → **New Notebook**
|
||||
- Upload `train_classifier_kaggle.ipynb` via **File → Import Notebook**
|
||||
|
||||
3. Attach your dataset:
|
||||
- In the right panel click **+ Add Data**
|
||||
- Search for your dataset name and attach it
|
||||
- It will appear at `/kaggle/input/my-product-crops/`
|
||||
|
||||
4. Enable the GPU:
|
||||
- Right panel → **Session options → Accelerator → GPU T4 x1**
|
||||
|
||||
5. Edit **Cell 2 — Configuration** only:
|
||||
```python
|
||||
DATA_DIR = Path("/kaggle/input/my-product-crops") # ← your dataset path
|
||||
MODEL_NAME = "efficientnet_b0"
|
||||
EPOCHS = 50
|
||||
BATCH_SIZE = 64
|
||||
```
|
||||
|
||||
6. **Run All** (`Shift+Enter` through all cells, or **Run All** from the menu)
|
||||
|
||||
**What the notebook does beyond the plain script:**
|
||||
|
||||
| Cell | Extra feature |
|
||||
|---|---|
|
||||
| Cell 2b | Dataset diagnostics — prints per-class counts, flags classes with < 20 images, prints random-chance loss baseline |
|
||||
| Cell 4 | Class distribution bar chart |
|
||||
| Cell 5 | Augmented sample grid so you can visually verify your data loaded correctly |
|
||||
| Cell 9 | Live training curves that update every epoch via `clear_output` |
|
||||
| Cell 11 | Per-class validation accuracy with colour-coded bar chart (red < 50%, orange < 80%, green ≥ 80%) |
|
||||
| Cell 12 | Normalised confusion matrix for the full val set |
|
||||
| Cell 13 | Lists all output files with sizes |
|
||||
|
||||
**Downloading outputs:**
|
||||
|
||||
After training, download from `/kaggle/working/runs/classify/`:
|
||||
- `best.pt` — copy to `runs/classify/best.pt` locally
|
||||
- `class_names.json` — copy alongside `best.pt`
|
||||
|
||||
These two files are all you need to run `3_inference.py` and `api_server.py`.
|
||||
|
||||
#### Diagnosing a stuck val loss
|
||||
|
||||
A val loss stuck above `ln(num_classes)` (e.g. > 3.13 for 23 classes) means the model is guessing randomly or confidently wrong. Common causes:
|
||||
|
||||
| Symptom | Fix |
|
||||
|---|---|
|
||||
| Too few images per class (< 20) | Collect more data via `auto_label.py` on new images |
|
||||
| Large imbalance between classes | `WeightedRandomSampler` is already enabled — check the imbalance report printed at startup |
|
||||
| `_ignore_` class in dataset | **Remove it.** Use `--cls_conf` threshold for rejection instead |
|
||||
| Val loss high but train loss low | Classic overfitting — add `--freeze_backbone` or collect more data |
|
||||
|
||||
> **Do not use an `_ignore_` class.** Throwing unrelated products into one folder
|
||||
> creates an incoherent class that harms the entire model. Use `--cls_conf` to
|
||||
> reject low-confidence predictions at inference time instead.
|
||||
|
||||
---
|
||||
|
||||
### Step 5 — Run inference
|
||||
|
||||
Runs the full detect → classify pipeline on images, folders, videos, or webcam. Draws annotated bounding boxes and saves results.
|
||||
|
||||
```bash
|
||||
python 3_inference.py \
|
||||
--detector_weights detector/best.pt \
|
||||
--classifier_weights runs/classify/best.pt \
|
||||
--source shelf_images/ \
|
||||
--det_conf 0.30 \
|
||||
--cls_conf 0.65 \
|
||||
--output_dir inference_results \
|
||||
--show
|
||||
```
|
||||
|
||||
| Argument | Default | Description |
|
||||
|---|---|---|
|
||||
| `--source` | required | Image, folder, video, or webcam index (`0`) |
|
||||
| `--detector_weights` | required | YOLOv8n `.pt` |
|
||||
| `--classifier_weights` | required | Classifier `best.pt` |
|
||||
| `--det_conf` | `0.30` | Detection threshold |
|
||||
| `--cls_conf` | `0.50` | Min classifier confidence to draw a label |
|
||||
| `--det_iou` | `0.45` | NMS IoU threshold |
|
||||
| `--padding` | `8` | Extra pixels added around each crop before classification |
|
||||
| `--output_dir` | `inference_results` | Where annotated images/video are saved |
|
||||
| `--show` | off | Display live with `cv2.imshow` |
|
||||
| `--save_crops` | off | Also save individual crop images |
|
||||
| `--no_save` | off | Skip saving annotated output |
|
||||
|
||||
#### Dynamic batch sizing
|
||||
|
||||
The classifier automatically probes available VRAM at startup and picks the largest safe batch size (with a 20% headroom margin). If an OOM occurs at runtime (e.g. an unusually dense frame with 260+ products), the batch is halved automatically and persisted for future frames. No manual tuning needed.
|
||||
|
||||
---
|
||||
|
||||
## 5. API Server
|
||||
|
||||
A REST API with a polling pattern for integration into other applications.
|
||||
|
||||
```bash
|
||||
pip install fastapi uvicorn
|
||||
uvicorn api_server:app --host 0.0.0.0 --port 8000 --workers 1
|
||||
```
|
||||
|
||||
Configuration via environment variables:
|
||||
|
||||
```bash
|
||||
DETECTOR_WEIGHTS=detector/best.pt
|
||||
CLASSIFIER_WEIGHTS=runs/classify/best.pt
|
||||
DET_CONF=0.25
|
||||
CLS_CONF=0.50
|
||||
```
|
||||
|
||||
### Endpoints
|
||||
|
||||
#### `POST /jobs` — Submit a job
|
||||
|
||||
```json
|
||||
{
|
||||
"image": "<base64-encoded image>",
|
||||
"interested_classes": ["cola_can", "pepsi_can"],
|
||||
"cls_conf": 0.65
|
||||
}
|
||||
```
|
||||
|
||||
`interested_classes` filters the results to only those classes. Send an empty list `[]` to return all detected products.
|
||||
|
||||
Returns immediately:
|
||||
```json
|
||||
{ "job_id": "550e8400-e29b-41d4-a716-446655440000" }
|
||||
```
|
||||
|
||||
#### `GET /jobs/{job_id}` — Poll for results
|
||||
|
||||
Call every ~3 seconds until `status` is `"done"` or `"failed"`.
|
||||
|
||||
```json
|
||||
{
|
||||
"status": "processing",
|
||||
"progress": 50,
|
||||
"results": null
|
||||
}
|
||||
```
|
||||
|
||||
When done:
|
||||
```json
|
||||
{
|
||||
"status": "done",
|
||||
"progress": 100,
|
||||
"results": {
|
||||
"annotated_image": "<base64 JPEG>",
|
||||
"counts": {
|
||||
"cola_can": 3,
|
||||
"pepsi_can": 1
|
||||
},
|
||||
"detections": [
|
||||
{ "class": "cola_can", "confidence": 0.91, "bbox": [120, 45, 210, 180] },
|
||||
{ "class": "cola_can", "confidence": 0.87, "bbox": [230, 48, 318, 182] },
|
||||
{ "class": "pepsi_can", "confidence": 0.79, "bbox": [340, 51, 428, 185] }
|
||||
]
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Progress stages:
|
||||
|
||||
| Progress | Stage |
|
||||
|---|---|
|
||||
| 5 | Image decoded |
|
||||
| 10 | Detection started |
|
||||
| 30 | Detection done |
|
||||
| 50 | Crops extracted |
|
||||
| 75 | Classification done |
|
||||
| 90 | Image annotated |
|
||||
| 100 | Done |
|
||||
|
||||
#### `GET /health` — Server health check
|
||||
|
||||
```json
|
||||
{
|
||||
"status": "ok",
|
||||
"device": "cuda",
|
||||
"classes": ["cola_can", "pepsi_can", "..."],
|
||||
"active_jobs": 2
|
||||
}
|
||||
```
|
||||
|
||||
#### `GET /classes` — List available classes
|
||||
|
||||
```json
|
||||
{ "classes": ["cola_can", "pepsi_can", "lays_chips", "..."] }
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 6. Script Reference
|
||||
|
||||
| Script | Purpose | Inputs | Outputs |
|
||||
|---|---|---|---|
|
||||
| `1_generate_crop_dataset.py` | Detect & crop (no classifier) | shelf images | `data/unknown/*.jpg` |
|
||||
| `auto_label.py` | Detect + auto-classify + browser review | shelf images + both models | `data/<class>/*.jpg` |
|
||||
| `1b_split_dataset.py` | Stratified train/val/test split | `data/` | `crops_dataset/` |
|
||||
| `2_train_classifier.py` | Train EfficientNet classifier | `crops_dataset/` | `runs/classify/best.pt` |
|
||||
| `3_inference.py` | End-to-end detect + classify | images / video / webcam | annotated images/video |
|
||||
| `api_server.py` | REST polling API | — | JSON + base64 annotated image |
|
||||
| `train_classifier_kaggle.ipynb` | Kaggle training notebook | Kaggle dataset | `best.pt` + plots |
|
||||
|
||||
---
|
||||
|
||||
## 7. Troubleshooting
|
||||
|
||||
**`FileNotFoundError: No images found in: 01.jpg`**
|
||||
You passed a file path that does not exist relative to your working directory. Use the full path or `cd` into the correct folder first.
|
||||
|
||||
**`OverflowError: Python integer 276 out of bounds for uint8`**
|
||||
This was a known bug (HSV hue > 179) — already fixed in the current `3_inference.py`.
|
||||
|
||||
**Val loss stuck at ~5 during training**
|
||||
See [Diagnosing a stuck val loss](#diagnosing-a-stuck-val-loss) above.
|
||||
|
||||
**`CUDA out of memory` during inference**
|
||||
The dynamic batch sizer handles this automatically by halving the batch and retrying. If it keeps happening, your crops may be very large — reduce `--padding` or `--img_size`.
|
||||
|
||||
**`CUDA out of memory` during training**
|
||||
Reduce `--batch_size`. Start at 32 and halve until it fits.
|
||||
|
||||
**Browser does not open automatically after `auto_label.py`**
|
||||
Open `http://localhost:5000` manually.
|
||||
|
||||
**Classifier only predicts one class**
|
||||
Almost always class imbalance. Check the imbalance report printed during `build_loaders`. The `WeightedRandomSampler` should compensate, but if one class has 10× more images, collect more data for the minority classes.
|
||||
|
||||
---
|
||||
|
||||
## 8. Tips & Best Practices
|
||||
|
||||
**Data collection**
|
||||
- Aim for **≥ 50 training images per class** before expecting good results.
|
||||
- Vary lighting, angles, partial occlusion, and zoom levels in your source images.
|
||||
- Run `auto_label.py` on new batches of images regularly — the dataset grows incrementally and the model improves each iteration.
|
||||
|
||||
**Labelling**
|
||||
- Use the **confidence slider** in the review UI to start reviewing high-confidence crops first (`> 0.8`) — these are almost always correct and can be bulk-confirmed quickly.
|
||||
- Filter the sidebar to one class at a time and use **Confirm all visible** — scanning one class at a time is much faster than random-order review.
|
||||
- Never put unrecognised products into a catch-all class. Leave them in `_unreviewed_/` or reject them. Use `--cls_conf` to suppress uncertain predictions at inference time.
|
||||
|
||||
**Training**
|
||||
- Start with `--freeze_backbone` if you have fewer than 50 images per class. Once you have more data, retrain without it.
|
||||
- Use the Kaggle notebook for free T4 GPU access — it includes a diagnostic cell that checks class counts and imbalance before training starts.
|
||||
- `efficientnet_b0` is the best default — fast, accurate, small. Only upgrade to `efficientnet_b2` if B0 accuracy plateaus and you have plenty of data.
|
||||
|
||||
**Iterative improvement loop**
|
||||
|
||||
```
|
||||
Shoot new shelf photos
|
||||
↓
|
||||
auto_label.py (detect + auto-classify)
|
||||
↓
|
||||
Review UI (confirm / fix / reject)
|
||||
↓
|
||||
1b_split_dataset.py (re-split the grown dataset)
|
||||
↓
|
||||
2_train_classifier.py (retrain from scratch or from last.pt)
|
||||
↓
|
||||
3_inference.py / api_server.py (deploy updated model)
|
||||
↓
|
||||
repeat
|
||||
```
|
||||
|
||||
Each iteration your classifier gets better, which means the auto-labelling in the next iteration needs less manual correction.
|
||||
458
api_server.py
Normal file
458
api_server.py
Normal file
@ -0,0 +1,458 @@
|
||||
"""
|
||||
API Server — Detect + Classify with Polling
|
||||
=============================================
|
||||
POST /jobs — submit an image + class filter, get back a job_id
|
||||
GET /jobs/{job_id} — poll for status / progress / results
|
||||
|
||||
Polling pattern
|
||||
---------------
|
||||
1. Client POSTs image (base64) + interested_classes list
|
||||
2. Server returns { "job_id": "<uuid>" } immediately
|
||||
3. Client polls GET /jobs/{job_id} every 3 s
|
||||
4. Server returns:
|
||||
{
|
||||
"status": "queued" | "processing" | "done" | "failed",
|
||||
"progress": 0‥100,
|
||||
"results": {
|
||||
"annotated_image": "<base64 jpg>",
|
||||
"counts": { "cola_can": 3, "pepsi_can": 1, … },
|
||||
"detections": [
|
||||
{ "class": "cola_can", "confidence": 0.91,
|
||||
"bbox": [x1, y1, x2, y2] },
|
||||
…
|
||||
]
|
||||
} // null while processing
|
||||
}
|
||||
|
||||
Run
|
||||
---
|
||||
pip install fastapi uvicorn python-multipart
|
||||
|
||||
uvicorn api_server:app --host 0.0.0.0 --port 8000 --workers 1
|
||||
|
||||
Environment variables (all optional)
|
||||
-------------------------------------
|
||||
DETECTOR_WEIGHTS path to YOLOv8n .pt (default: detector/best.pt)
|
||||
CLASSIFIER_WEIGHTS path to classifier .pt (default: runs/classify/best.pt)
|
||||
DET_CONF float (default: 0.25)
|
||||
CLS_CONF float (default: 0.50)
|
||||
DET_IOU float (default: 0.45)
|
||||
PADDING int px (default: 10)
|
||||
MAX_JOBS max jobs kept in memory (default: 200)
|
||||
"""
|
||||
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import uuid
|
||||
from collections import OrderedDict
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from enum import Enum
|
||||
from typing import Optional
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from fastapi import FastAPI, HTTPException, BackgroundTasks
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from pydantic import BaseModel, Field
|
||||
from torchvision import models, transforms
|
||||
from ultralytics import YOLO
|
||||
|
||||
# ─────────────────────────── config ──────────────────────────────────────────
|
||||
|
||||
DETECTOR_WEIGHTS = os.getenv("DETECTOR_WEIGHTS", "detector/best.pt")
|
||||
CLASSIFIER_WEIGHTS = os.getenv("CLASSIFIER_WEIGHTS", "best.pt")
|
||||
DET_CONF = float(os.getenv("DET_CONF", "0.25"))
|
||||
CLS_CONF = float(os.getenv("CLS_CONF", "0.75"))
|
||||
DET_IOU = float(os.getenv("DET_IOU", "0.45"))
|
||||
PADDING = int(os.getenv("PADDING", "10"))
|
||||
MAX_JOBS = int(os.getenv("MAX_JOBS", "200"))
|
||||
|
||||
# ─────────────────────────── device ──────────────────────────────────────────
|
||||
|
||||
def get_device() -> torch.device:
|
||||
if torch.cuda.is_available():
|
||||
return torch.device("cuda")
|
||||
if torch.backends.mps.is_available():
|
||||
return torch.device("mps")
|
||||
return torch.device("cpu")
|
||||
|
||||
DEVICE = get_device()
|
||||
|
||||
# ─────────────────────────── OOM helpers (same as inference script) ──────────
|
||||
|
||||
def _is_oom(exc: BaseException) -> bool:
|
||||
if isinstance(exc, MemoryError):
|
||||
return True
|
||||
if torch.cuda.is_available() and isinstance(exc, torch.cuda.OutOfMemoryError):
|
||||
return True
|
||||
if isinstance(exc, RuntimeError):
|
||||
msg = str(exc).lower()
|
||||
return "out of memory" in msg or ("memory" in msg and "alloc" in msg)
|
||||
return False
|
||||
|
||||
def _free_device_cache(device: torch.device) -> None:
|
||||
if device.type == "cuda":
|
||||
torch.cuda.empty_cache()
|
||||
elif device.type == "mps" and hasattr(torch.mps, "empty_cache"):
|
||||
torch.mps.empty_cache()
|
||||
|
||||
def _sync_device(device: torch.device) -> None:
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize()
|
||||
elif device.type == "mps" and hasattr(torch.mps, "synchronize"):
|
||||
torch.mps.synchronize()
|
||||
|
||||
# ─────────────────────────── classifier ──────────────────────────────────────
|
||||
|
||||
class ProductClassifier:
|
||||
_BUILDERS = {
|
||||
"efficientnet_b0": models.efficientnet_b0,
|
||||
"efficientnet_b2": models.efficientnet_b2,
|
||||
"mobilenet_v3_small": models.mobilenet_v3_small,
|
||||
"resnet50": models.resnet50,
|
||||
}
|
||||
|
||||
def __init__(self, weights_path: str, device: torch.device, img_size: int = 224):
|
||||
ckpt = torch.load(weights_path, map_location=device)
|
||||
self.class_names: list[str] = ckpt["class_names"]
|
||||
num_classes = len(self.class_names)
|
||||
model_name = ckpt.get("args", {}).get("model", "efficientnet_b0")
|
||||
|
||||
base = self._BUILDERS[model_name](weights=None)
|
||||
# Attach correct head
|
||||
if "efficientnet" in model_name:
|
||||
base.classifier[1] = torch.nn.Linear(base.classifier[1].in_features, num_classes)
|
||||
elif "mobilenet" in model_name:
|
||||
base.classifier[3] = torch.nn.Linear(base.classifier[3].in_features, num_classes)
|
||||
else:
|
||||
base.fc = torch.nn.Linear(base.fc.in_features, num_classes)
|
||||
|
||||
base.load_state_dict(ckpt["model_state"])
|
||||
self.model = base.to(device).eval()
|
||||
self.device = device
|
||||
|
||||
mean = [0.485, 0.456, 0.406]
|
||||
std = [0.229, 0.224, 0.225]
|
||||
self.transform = transforms.Compose([
|
||||
transforms.ToPILImage(),
|
||||
transforms.Resize(int(img_size * 1.15)),
|
||||
transforms.CenterCrop(img_size),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean, std),
|
||||
])
|
||||
self._batch_size: int = self._calibrate(img_size)
|
||||
|
||||
def _calibrate(self, img_size: int) -> int:
|
||||
if self.device.type == "cpu":
|
||||
return 64
|
||||
dummy = torch.zeros(1, 3, img_size, img_size, device=self.device)
|
||||
safe, probe = 1, 2
|
||||
with torch.no_grad():
|
||||
while probe <= 512:
|
||||
try:
|
||||
self.model(dummy.expand(probe, -1, -1, -1))
|
||||
_sync_device(self.device)
|
||||
safe, probe = probe, probe * 2
|
||||
except Exception as exc:
|
||||
if _is_oom(exc):
|
||||
break
|
||||
raise
|
||||
finally:
|
||||
_free_device_cache(self.device)
|
||||
return max(1, int(safe * 0.8))
|
||||
|
||||
@torch.no_grad()
|
||||
def predict_batch(self, crops_bgr: list[np.ndarray]) -> list[tuple[str, float]]:
|
||||
if not crops_bgr:
|
||||
return []
|
||||
tensors = [self.transform(cv2.cvtColor(c, cv2.COLOR_BGR2RGB)) for c in crops_bgr]
|
||||
results: list[tuple[str, float]] = []
|
||||
chunk_size = self._batch_size
|
||||
i = 0
|
||||
while i < len(tensors):
|
||||
batch = torch.stack(tensors[i : i + chunk_size]).to(self.device)
|
||||
try:
|
||||
probs = F.softmax(self.model(batch), dim=1)
|
||||
confs, idxs = probs.max(dim=1)
|
||||
for idx, conf in zip(idxs, confs):
|
||||
results.append((self.class_names[idx.item()], conf.item()))
|
||||
i += chunk_size
|
||||
except Exception as exc:
|
||||
if not _is_oom(exc):
|
||||
raise
|
||||
_free_device_cache(self.device)
|
||||
chunk_size = self._batch_size = max(1, chunk_size // 2)
|
||||
finally:
|
||||
del batch
|
||||
return results
|
||||
|
||||
# ─────────────────────────── model singletons ────────────────────────────────
|
||||
|
||||
print(f"[startup] Device: {DEVICE}")
|
||||
print(f"[startup] Loading detector : {DETECTOR_WEIGHTS}")
|
||||
DETECTOR = YOLO(DETECTOR_WEIGHTS)
|
||||
print(f"[startup] Loading classifier: {CLASSIFIER_WEIGHTS}")
|
||||
CLASSIFIER = ProductClassifier(CLASSIFIER_WEIGHTS, DEVICE)
|
||||
print(f"[startup] Ready — classes: {CLASSIFIER.class_names}")
|
||||
|
||||
# Thread pool: one worker keeps GPU access serialised; increase if CPU-only
|
||||
EXECUTOR = ThreadPoolExecutor(max_workers=1)
|
||||
|
||||
# ─────────────────────────── job store ───────────────────────────────────────
|
||||
|
||||
class Status(str, Enum):
|
||||
queued = "queued"
|
||||
processing = "processing"
|
||||
done = "done"
|
||||
failed = "failed"
|
||||
|
||||
class Job:
|
||||
def __init__(self, job_id: str, image_bytes: bytes,
|
||||
interested_classes: list[str], cls_conf: float):
|
||||
self.job_id = job_id
|
||||
self.image_bytes = image_bytes
|
||||
self.interested_classes = [c.lower() for c in interested_classes]
|
||||
self.cls_conf = cls_conf
|
||||
self.status = Status.queued
|
||||
self.progress = 0
|
||||
self.results = None
|
||||
self.error = None
|
||||
|
||||
# OrderedDict so we can evict oldest jobs when MAX_JOBS is reached
|
||||
JOB_STORE: OrderedDict[str, Job] = OrderedDict()
|
||||
|
||||
def store_job(job: Job) -> None:
|
||||
if len(JOB_STORE) >= MAX_JOBS:
|
||||
JOB_STORE.popitem(last=False) # evict oldest
|
||||
JOB_STORE[job.job_id] = job
|
||||
|
||||
def get_job(job_id: str) -> Job:
|
||||
job = JOB_STORE.get(job_id)
|
||||
if job is None:
|
||||
raise HTTPException(status_code=404, detail=f"Job '{job_id}' not found")
|
||||
return job
|
||||
|
||||
# ─────────────────────────── drawing ─────────────────────────────────────────
|
||||
|
||||
_COLOUR_CACHE: dict[str, tuple] = {}
|
||||
|
||||
def class_colour(name: str) -> tuple:
|
||||
if name not in _COLOUR_CACHE:
|
||||
h = hash(name) % 180 # OpenCV hue: 0–179
|
||||
col = cv2.cvtColor(
|
||||
np.array([[[h, 220, 200]]], dtype=np.uint8), cv2.COLOR_HSV2BGR
|
||||
)[0][0]
|
||||
_COLOUR_CACHE[name] = tuple(int(c) for c in col)
|
||||
return _COLOUR_CACHE[name]
|
||||
|
||||
def draw_box(frame: np.ndarray, x1: int, y1: int, x2: int, y2: int,
|
||||
label: str, conf: float) -> None:
|
||||
colour = class_colour(label)
|
||||
cv2.rectangle(frame, (x1, y1), (x2, y2), colour, 2)
|
||||
text = f"{label} {conf:.2f}"
|
||||
(tw, th), bl = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.55, 1)
|
||||
cv2.rectangle(frame, (x1, y1 - th - bl - 4), (x1 + tw + 4, y1), colour, -1)
|
||||
cv2.putText(frame, text, (x1 + 2, y1 - bl - 2),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 255), 1, cv2.LINE_AA)
|
||||
|
||||
# ─────────────────────────── pipeline ────────────────────────────────────────
|
||||
|
||||
def _encode_image(frame: np.ndarray) -> str:
|
||||
"""Encode a BGR numpy array as a base64 JPEG string."""
|
||||
_, buf = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 90])
|
||||
return base64.b64encode(buf).decode("utf-8")
|
||||
|
||||
def run_pipeline(job: Job) -> None:
|
||||
"""Blocking function — runs in the thread-pool executor."""
|
||||
try:
|
||||
job.status = Status.processing
|
||||
job.progress = 5
|
||||
|
||||
# ── Decode image ─────────────────────────────────────────────────────
|
||||
arr = np.frombuffer(job.image_bytes, np.uint8)
|
||||
frame = cv2.imdecode(arr, cv2.IMREAD_COLOR)
|
||||
if frame is None:
|
||||
raise ValueError("Could not decode image — unsupported format or corrupt data")
|
||||
|
||||
job.progress = 10
|
||||
|
||||
# ── Detection ────────────────────────────────────────────────────────
|
||||
h, w = frame.shape[:2]
|
||||
det_results = DETECTOR(
|
||||
frame, conf=DET_CONF, iou=DET_IOU, verbose=False
|
||||
)[0]
|
||||
|
||||
job.progress = 30
|
||||
|
||||
# ── Crop each detection ───────────────────────────────────────────────
|
||||
crops, boxes = [], []
|
||||
for box in det_results.boxes:
|
||||
x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
|
||||
x1 = max(0, x1 - PADDING); y1 = max(0, y1 - PADDING)
|
||||
x2 = min(w, x2 + PADDING); y2 = min(h, y2 + PADDING)
|
||||
crop = frame[y1:y2, x1:x2]
|
||||
if crop.size == 0:
|
||||
continue
|
||||
crops.append(crop)
|
||||
boxes.append((x1, y1, x2, y2))
|
||||
|
||||
job.progress = 50
|
||||
|
||||
# ── Classify ─────────────────────────────────────────────────────────
|
||||
predictions = CLASSIFIER.predict_batch(crops)
|
||||
|
||||
job.progress = 75
|
||||
|
||||
# ── Filter by interested_classes + cls_conf ───────────────────────────
|
||||
annotated = frame.copy()
|
||||
counts: dict[str, int] = {}
|
||||
detections: list[dict] = []
|
||||
|
||||
# If the user sends an empty list → show everything above cls_conf
|
||||
filter_active = len(job.interested_classes) > 0
|
||||
|
||||
for (x1, y1, x2, y2), (cls_name, conf) in zip(boxes, predictions):
|
||||
if conf < job.cls_conf:
|
||||
continue
|
||||
if filter_active and cls_name.lower() not in job.interested_classes:
|
||||
continue
|
||||
|
||||
draw_box(annotated, x1, y1, x2, y2, cls_name, conf)
|
||||
counts[cls_name] = counts.get(cls_name, 0) + 1
|
||||
detections.append({
|
||||
"class": cls_name,
|
||||
"confidence": round(conf, 4),
|
||||
"bbox": [x1, y1, x2, y2],
|
||||
})
|
||||
|
||||
job.progress = 90
|
||||
|
||||
# ── Encode output image ───────────────────────────────────────────────
|
||||
annotated_b64 = _encode_image(annotated)
|
||||
|
||||
job.results = {
|
||||
"annotated_image": annotated_b64,
|
||||
"counts": counts,
|
||||
"detections": detections,
|
||||
}
|
||||
job.status = Status.done
|
||||
job.progress = 100
|
||||
|
||||
except Exception as exc:
|
||||
job.status = Status.failed
|
||||
job.error = str(exc)
|
||||
raise
|
||||
|
||||
# ─────────────────────────── schemas ─────────────────────────────────────────
|
||||
|
||||
class SubmitRequest(BaseModel):
|
||||
# Base64-encoded image (JPEG / PNG / BMP …)
|
||||
image: str = Field(..., description="Base64-encoded image bytes")
|
||||
# Classes to keep in results; empty list = keep all classes
|
||||
interested_classes: list[str] = Field(
|
||||
default=[],
|
||||
description="Product class names to include. Leave empty to show all.",
|
||||
examples=[["cola_can", "pepsi_can"]],
|
||||
)
|
||||
# Per-request confidence override (optional)
|
||||
cls_conf: Optional[float] = Field(
|
||||
default=None,
|
||||
ge=0.0, le=1.0,
|
||||
description="Min classifier confidence (overrides server default)",
|
||||
)
|
||||
|
||||
class SubmitResponse(BaseModel):
|
||||
job_id: str
|
||||
|
||||
class PollResponse(BaseModel):
|
||||
status: Status
|
||||
progress: int
|
||||
results: Optional[dict] = None
|
||||
error: Optional[str] = None
|
||||
|
||||
# ─────────────────────────── app ─────────────────────────────────────────────
|
||||
|
||||
app = FastAPI(
|
||||
title="Product Detection API",
|
||||
description="Detect and classify shelf products using YOLOv8 + EfficientNet",
|
||||
version="1.0.0",
|
||||
)
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
# ── POST /jobs ────────────────────────────────────────────────────────────────
|
||||
|
||||
@app.post("/jobs", response_model=SubmitResponse, status_code=202)
|
||||
async def submit_job(body: SubmitRequest, background_tasks: BackgroundTasks):
|
||||
"""
|
||||
Submit an image for detection + classification.
|
||||
|
||||
Returns a job_id to poll with GET /jobs/{job_id}.
|
||||
"""
|
||||
# Decode base64 → raw bytes
|
||||
try:
|
||||
image_bytes = base64.b64decode(body.image)
|
||||
except Exception:
|
||||
raise HTTPException(status_code=422, detail="Invalid base64 in 'image' field")
|
||||
|
||||
job_id = str(uuid.uuid4())
|
||||
job = Job(
|
||||
job_id = job_id,
|
||||
image_bytes = image_bytes,
|
||||
interested_classes = body.interested_classes,
|
||||
cls_conf = body.cls_conf if body.cls_conf is not None else CLS_CONF,
|
||||
)
|
||||
store_job(job)
|
||||
|
||||
# Submit to thread pool — does not block the event loop
|
||||
EXECUTOR.submit(run_pipeline, job)
|
||||
|
||||
return SubmitResponse(job_id=job_id)
|
||||
|
||||
|
||||
# ── GET /jobs/{job_id} ────────────────────────────────────────────────────────
|
||||
|
||||
@app.get("/jobs/{job_id}", response_model=PollResponse)
|
||||
async def poll_job(job_id: str):
|
||||
"""
|
||||
Poll the status of a submitted job.
|
||||
|
||||
Call every ~3 seconds until status is 'done' or 'failed'.
|
||||
"""
|
||||
job = get_job(job_id)
|
||||
return PollResponse(
|
||||
status = job.status,
|
||||
progress = job.progress,
|
||||
results = job.results if job.status == Status.done else None,
|
||||
error = job.error,
|
||||
)
|
||||
|
||||
|
||||
# ── GET /health ───────────────────────────────────────────────────────────────
|
||||
|
||||
@app.get("/health")
|
||||
async def health():
|
||||
return {
|
||||
"status": "ok",
|
||||
"device": str(DEVICE),
|
||||
"classes": CLASSIFIER.class_names,
|
||||
"active_jobs": sum(1 for j in JOB_STORE.values()
|
||||
if j.status in (Status.queued, Status.processing)),
|
||||
}
|
||||
|
||||
|
||||
# ── GET /classes ──────────────────────────────────────────────────────────────
|
||||
|
||||
@app.get("/classes")
|
||||
async def list_classes():
|
||||
"""Return all class names the classifier knows about."""
|
||||
return {"classes": CLASSIFIER.class_names}
|
||||
986
auto_label.py
Normal file
986
auto_label.py
Normal file
@ -0,0 +1,986 @@
|
||||
"""
|
||||
auto_label.py — Auto-classify crops + Web Review UI
|
||||
=====================================================
|
||||
Runs the full pipeline in one command:
|
||||
|
||||
Step 1 Run YOLOv8n on your source images → crop every detection
|
||||
Step 2 Run the classifier on every crop → propose a class label
|
||||
Step 3 Open a local web UI so you review, confirm, fix, or reject
|
||||
Step 4 On "Commit", write the final data/ folder ready for 1b_split_dataset.py
|
||||
|
||||
Output layout (after commit)
|
||||
-----------------------------
|
||||
data/
|
||||
cola_can/ ← confirmed crops
|
||||
pepsi_can/
|
||||
…
|
||||
_rejected_/ ← crops you marked as reject
|
||||
_unreviewed_/ ← anything you closed without reviewing
|
||||
|
||||
Usage
|
||||
-----
|
||||
python auto_label.py \
|
||||
--source shelf_images/ \
|
||||
--detector_weights detector/best.pt \
|
||||
--classifier_weights runs/classify/best.pt \
|
||||
--output_dir data \
|
||||
--det_conf 0.25 \
|
||||
--cls_conf 0.50 \
|
||||
--port 5000
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import webbrowser
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from flask import Flask, jsonify, request, send_file
|
||||
from torchvision import models, transforms
|
||||
from tqdm import tqdm
|
||||
from ultralytics import YOLO
|
||||
|
||||
# ─────────────────────────── args ────────────────────────────────────────────
|
||||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser(description="Auto-label crops and review in browser")
|
||||
p.add_argument("--source", required=True)
|
||||
p.add_argument("--detector_weights", default="detector/best.pt")
|
||||
p.add_argument("--classifier_weights", default="best.pt")
|
||||
p.add_argument("--output_dir", default="data")
|
||||
p.add_argument("--staging_dir", default=".autolabel_staging",
|
||||
help="Temp folder for staged crops before commit (hidden by default)")
|
||||
p.add_argument("--det_conf", type=float, default=0.25)
|
||||
p.add_argument("--det_iou", type=float, default=0.45)
|
||||
p.add_argument("--cls_conf", type=float, default=0.75)
|
||||
p.add_argument("--img_size", type=int, default=640)
|
||||
p.add_argument("--padding", type=int, default=10)
|
||||
p.add_argument("--min_size", type=int, default=32)
|
||||
p.add_argument("--port", type=int, default=5000)
|
||||
return p.parse_args()
|
||||
|
||||
ARGS = parse_args()
|
||||
|
||||
# ─────────────────────────── device / OOM ────────────────────────────────────
|
||||
|
||||
def get_device():
|
||||
if torch.cuda.is_available(): return torch.device("cuda")
|
||||
if torch.backends.mps.is_available(): return torch.device("mps")
|
||||
return torch.device("cpu")
|
||||
|
||||
DEVICE = get_device()
|
||||
|
||||
def _is_oom(exc):
|
||||
if isinstance(exc, MemoryError): return True
|
||||
if torch.cuda.is_available() and isinstance(exc, torch.cuda.OutOfMemoryError): return True
|
||||
if isinstance(exc, RuntimeError):
|
||||
m = str(exc).lower()
|
||||
return "out of memory" in m or ("memory" in m and "alloc" in m)
|
||||
return False
|
||||
|
||||
def _free_cache():
|
||||
if DEVICE.type == "cuda": torch.cuda.empty_cache()
|
||||
elif DEVICE.type == "mps" and hasattr(torch.mps, "empty_cache"): torch.mps.empty_cache()
|
||||
|
||||
# ─────────────────────────── classifier ──────────────────────────────────────
|
||||
|
||||
class ProductClassifier:
|
||||
def __init__(self, weights_path, device, img_size=224):
|
||||
ckpt = torch.load(weights_path, map_location=device)
|
||||
self.class_names = ckpt["class_names"]
|
||||
n = len(self.class_names)
|
||||
name = ckpt.get("args", {}).get("model", "efficientnet_b0")
|
||||
|
||||
m = {
|
||||
"efficientnet_b0": models.efficientnet_b0,
|
||||
"efficientnet_b2": models.efficientnet_b2,
|
||||
"mobilenet_v3_small": models.mobilenet_v3_small,
|
||||
"resnet50": models.resnet50,
|
||||
}[name](weights=None)
|
||||
|
||||
if "efficientnet" in name: m.classifier[1] = torch.nn.Linear(m.classifier[1].in_features, n)
|
||||
elif "mobilenet" in name: m.classifier[3] = torch.nn.Linear(m.classifier[3].in_features, n)
|
||||
else: m.fc = torch.nn.Linear(m.fc.in_features, n)
|
||||
|
||||
m.load_state_dict(ckpt["model_state"])
|
||||
self.model = m.to(device).eval()
|
||||
self.device = device
|
||||
self.tf = transforms.Compose([
|
||||
transforms.ToPILImage(),
|
||||
transforms.Resize(int(img_size * 1.15)),
|
||||
transforms.CenterCrop(img_size),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225]),
|
||||
])
|
||||
self._bs = self._probe(img_size)
|
||||
|
||||
def _probe(self, sz):
|
||||
if self.device.type == "cpu": return 64
|
||||
dummy = torch.zeros(1,3,sz,sz,device=self.device)
|
||||
safe, probe = 1, 2
|
||||
with torch.no_grad():
|
||||
while probe <= 512:
|
||||
try:
|
||||
self.model(dummy.expand(probe,-1,-1,-1))
|
||||
safe, probe = probe, probe*2
|
||||
except Exception as e:
|
||||
if _is_oom(e): break
|
||||
raise
|
||||
finally: _free_cache()
|
||||
return max(1, int(safe*0.8))
|
||||
|
||||
@torch.no_grad()
|
||||
def predict(self, crops_bgr):
|
||||
if not crops_bgr: return []
|
||||
tensors = [self.tf(cv2.cvtColor(c, cv2.COLOR_BGR2RGB)) for c in crops_bgr]
|
||||
results, i, bs = [], 0, self._bs
|
||||
while i < len(tensors):
|
||||
batch = torch.stack(tensors[i:i+bs]).to(self.device)
|
||||
try:
|
||||
probs = F.softmax(self.model(batch), dim=1)
|
||||
confs, idxs = probs.max(dim=1)
|
||||
# Return top-3 predictions per crop for the review UI
|
||||
top3_probs, top3_idxs = torch.topk(probs, min(3, probs.shape[1]), dim=1)
|
||||
for j in range(len(idxs)):
|
||||
results.append({
|
||||
"class": self.class_names[idxs[j].item()],
|
||||
"confidence": round(confs[j].item(), 4),
|
||||
"top3": [
|
||||
{"class": self.class_names[top3_idxs[j][k].item()],
|
||||
"conf": round(top3_probs[j][k].item(), 4)}
|
||||
for k in range(top3_probs.shape[1])
|
||||
]
|
||||
})
|
||||
i += bs
|
||||
except Exception as e:
|
||||
if not _is_oom(e): raise
|
||||
_free_cache()
|
||||
bs = self._bs = max(1, bs // 2)
|
||||
finally:
|
||||
del batch
|
||||
return results
|
||||
|
||||
# ─────────────────────────── pipeline ────────────────────────────────────────
|
||||
|
||||
IMAGE_EXTS = {".jpg",".jpeg",".png",".bmp",".webp",".tiff"}
|
||||
|
||||
def get_images(source):
|
||||
p = Path(source)
|
||||
if p.is_file() and p.suffix.lower() in IMAGE_EXTS: return [p]
|
||||
if p.is_dir():
|
||||
imgs = sorted(f for f in p.rglob("*") if f.suffix.lower() in IMAGE_EXTS)
|
||||
if not imgs: raise FileNotFoundError(f"No images in {p}")
|
||||
return imgs
|
||||
raise FileNotFoundError(f"Cannot open source: {source}")
|
||||
|
||||
def img_to_b64(arr):
|
||||
_, buf = cv2.imencode(".jpg", arr, [cv2.IMWRITE_JPEG_QUALITY, 88])
|
||||
return base64.b64encode(buf).decode()
|
||||
|
||||
def run_pipeline(detector, classifier, args):
|
||||
"""
|
||||
Detect → crop → classify → stage.
|
||||
Returns list of crop metadata dicts that the review UI reads.
|
||||
"""
|
||||
staging = Path(args.staging_dir)
|
||||
if staging.exists(): shutil.rmtree(staging)
|
||||
staging.mkdir(parents=True)
|
||||
|
||||
image_paths = get_images(args.source)
|
||||
print(f"\n[pipeline] {len(image_paths)} source image(s) found")
|
||||
|
||||
crops_meta = [] # master list for the review UI
|
||||
total_det = total_skip = 0
|
||||
|
||||
for img_path in tqdm(image_paths, desc="Detecting"):
|
||||
frame = cv2.imread(str(img_path))
|
||||
if frame is None:
|
||||
print(f" ⚠ Cannot read {img_path.name}")
|
||||
continue
|
||||
|
||||
h, w = frame.shape[:2]
|
||||
results = detector(frame, conf=args.det_conf, iou=args.det_iou,
|
||||
imgsz=args.img_size, verbose=False)[0]
|
||||
|
||||
boxes, crops_bgr = [], []
|
||||
for box in results.boxes:
|
||||
x1,y1,x2,y2 = map(int, box.xyxy[0].tolist())
|
||||
x1=max(0,x1-args.padding); y1=max(0,y1-args.padding)
|
||||
x2=min(w,x2+args.padding); y2=min(h,y2+args.padding)
|
||||
crop = frame[y1:y2, x1:x2]
|
||||
if crop.size==0 or (x2-x1)<args.min_size or (y2-y1)<args.min_size:
|
||||
total_skip += 1
|
||||
continue
|
||||
boxes.append((x1,y1,x2,y2))
|
||||
crops_bgr.append(crop)
|
||||
|
||||
if not crops_bgr:
|
||||
continue
|
||||
|
||||
predictions = classifier.predict(crops_bgr)
|
||||
|
||||
for i, (box, pred, crop) in enumerate(zip(boxes, predictions, crops_bgr)):
|
||||
crop_id = f"{img_path.stem}_crop{i:04d}"
|
||||
cls_name = pred["class"]
|
||||
conf = pred["confidence"]
|
||||
|
||||
# Save crop to staging/<predicted_class>/
|
||||
cls_dir = staging / cls_name
|
||||
cls_dir.mkdir(exist_ok=True)
|
||||
crop_path = cls_dir / f"{crop_id}.jpg"
|
||||
cv2.imwrite(str(crop_path), crop)
|
||||
|
||||
crops_meta.append({
|
||||
"id": crop_id,
|
||||
"file": str(crop_path),
|
||||
"source_image": img_path.name,
|
||||
"pred_class": cls_name,
|
||||
"confidence": conf,
|
||||
"top3": pred["top3"],
|
||||
"bbox": list(box),
|
||||
# review fields (mutated by the UI)
|
||||
"label": cls_name, # user's chosen label
|
||||
"status": "auto", # "auto" | "confirmed" | "rejected" | "reclassified"
|
||||
})
|
||||
total_det += 1
|
||||
|
||||
print(f"[pipeline] {total_det} crops staged | {total_skip} skipped")
|
||||
return crops_meta
|
||||
|
||||
# ─────────────────────────── Flask app ───────────────────────────────────────
|
||||
|
||||
flask_app = Flask(__name__)
|
||||
STATE = {
|
||||
"crops": [], # list of crop metadata dicts
|
||||
"classes": [], # classifier class names
|
||||
"done": False, # set to True when user commits
|
||||
}
|
||||
|
||||
@flask_app.route("/")
|
||||
def index():
|
||||
return HTML_UI
|
||||
|
||||
@flask_app.route("/api/crops")
|
||||
def api_crops():
|
||||
return jsonify(STATE["crops"])
|
||||
|
||||
@flask_app.route("/api/classes")
|
||||
def api_classes():
|
||||
return jsonify(STATE["classes"])
|
||||
|
||||
@flask_app.route("/api/crop_image/<crop_id>")
|
||||
def crop_image(crop_id):
|
||||
crop = next((c for c in STATE["crops"] if c["id"] == crop_id), None)
|
||||
if not crop:
|
||||
return "not found", 404
|
||||
return send_file(crop["file"], mimetype="image/jpeg")
|
||||
|
||||
@flask_app.route("/api/update", methods=["POST"])
|
||||
def api_update():
|
||||
"""Receive bulk label updates from the UI."""
|
||||
updates = request.json # list of {id, label, status}
|
||||
lookup = {c["id"]: c for c in STATE["crops"]}
|
||||
for u in updates:
|
||||
if u["id"] in lookup:
|
||||
lookup[u["id"]]["label"] = u["label"]
|
||||
lookup[u["id"]]["status"] = u["status"]
|
||||
return jsonify({"ok": True})
|
||||
|
||||
@flask_app.route("/api/commit", methods=["POST"])
|
||||
def api_commit():
|
||||
"""Move staged crops to final output_dir based on user decisions."""
|
||||
out_root = Path(ARGS.output_dir)
|
||||
|
||||
moved = rejected = unreviewed = 0
|
||||
for crop in STATE["crops"]:
|
||||
src = Path(crop["file"])
|
||||
if not src.exists():
|
||||
continue
|
||||
|
||||
status = crop["status"]
|
||||
label = crop["label"]
|
||||
|
||||
if status == "rejected":
|
||||
dst_dir = out_root / "_rejected_"
|
||||
elif status == "unreviewed" or status == "auto":
|
||||
dst_dir = out_root / "_unreviewed_"
|
||||
else:
|
||||
dst_dir = out_root / label
|
||||
|
||||
dst_dir.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy2(str(src), str(dst_dir / src.name))
|
||||
|
||||
if status == "rejected": rejected += 1
|
||||
elif status in ("auto","unreviewed"): unreviewed += 1
|
||||
else: moved += 1
|
||||
|
||||
summary = {
|
||||
"committed": moved,
|
||||
"rejected": rejected,
|
||||
"unreviewed": unreviewed,
|
||||
"output_dir": str(out_root.resolve()),
|
||||
}
|
||||
STATE["done"] = True
|
||||
return jsonify(summary)
|
||||
|
||||
@flask_app.route("/api/stats")
|
||||
def api_stats():
|
||||
crops = STATE["crops"]
|
||||
by_class = {}
|
||||
for c in crops:
|
||||
by_class.setdefault(c["pred_class"], 0)
|
||||
by_class[c["pred_class"]] += 1
|
||||
statuses = {}
|
||||
for c in crops:
|
||||
statuses[c["status"]] = statuses.get(c["status"], 0) + 1
|
||||
return jsonify({"by_class": by_class, "by_status": statuses, "total": len(crops)})
|
||||
|
||||
# ─────────────────────────── embedded UI ─────────────────────────────────────
|
||||
|
||||
HTML_UI = r"""<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width,initial-scale=1">
|
||||
<title>Auto-Label Review</title>
|
||||
<style>
|
||||
@import url('https://fonts.googleapis.com/css2?family=DM+Mono:wght@400;500&family=Syne:wght@700;800&display=swap');
|
||||
|
||||
:root {
|
||||
--bg: #0d0f12;
|
||||
--surface: #161a20;
|
||||
--border: #252a33;
|
||||
--accent: #00e5a0;
|
||||
--warn: #f5c542;
|
||||
--danger: #ff4c6a;
|
||||
--muted: #4a5260;
|
||||
--text: #dde3ee;
|
||||
--radius: 6px;
|
||||
}
|
||||
|
||||
*, *::before, *::after { box-sizing: border-box; margin: 0; padding: 0; }
|
||||
|
||||
body {
|
||||
background: var(--bg);
|
||||
color: var(--text);
|
||||
font-family: 'DM Mono', monospace;
|
||||
font-size: 13px;
|
||||
min-height: 100vh;
|
||||
display: grid;
|
||||
grid-template-rows: 56px 1fr;
|
||||
}
|
||||
|
||||
/* ── header ── */
|
||||
header {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 20px;
|
||||
padding: 0 24px;
|
||||
border-bottom: 1px solid var(--border);
|
||||
background: var(--surface);
|
||||
position: sticky; top: 0; z-index: 50;
|
||||
}
|
||||
header h1 {
|
||||
font-family: 'Syne', sans-serif;
|
||||
font-size: 17px;
|
||||
font-weight: 800;
|
||||
letter-spacing: -.3px;
|
||||
color: var(--accent);
|
||||
white-space: nowrap;
|
||||
}
|
||||
.stats-bar { display: flex; gap: 16px; font-size: 11px; color: var(--muted); flex: 1; }
|
||||
.stats-bar span { display: flex; align-items: center; gap: 5px; }
|
||||
.dot { width: 7px; height: 7px; border-radius: 50%; display: inline-block; }
|
||||
.dot-auto { background: var(--muted); }
|
||||
.dot-confirmed { background: var(--accent); }
|
||||
.dot-reclassify { background: var(--warn); }
|
||||
.dot-rejected { background: var(--danger); }
|
||||
|
||||
.header-actions { display: flex; gap: 10px; margin-left: auto; }
|
||||
btn { display: inline-flex; align-items: center; gap: 6px;
|
||||
padding: 7px 14px; border-radius: var(--radius);
|
||||
font-family: 'DM Mono', monospace; font-size: 12px;
|
||||
cursor: pointer; border: none; font-weight: 500; transition: .15s; }
|
||||
.btn-outline { background: transparent; border: 1px solid var(--border); color: var(--text); }
|
||||
.btn-outline:hover { border-color: var(--accent); color: var(--accent); }
|
||||
.btn-primary { background: var(--accent); color: #000; }
|
||||
.btn-primary:hover { filter: brightness(1.1); }
|
||||
.btn-danger { background: var(--danger); color: #fff; }
|
||||
.btn-danger:hover { filter: brightness(1.1); }
|
||||
|
||||
/* ── layout ── */
|
||||
.layout { display: grid; grid-template-columns: 220px 1fr; height: calc(100vh - 56px); overflow: hidden; }
|
||||
|
||||
/* ── sidebar ── */
|
||||
aside {
|
||||
border-right: 1px solid var(--border);
|
||||
overflow-y: auto;
|
||||
padding: 12px 8px;
|
||||
background: var(--surface);
|
||||
}
|
||||
.sidebar-title { font-size: 10px; color: var(--muted); padding: 4px 8px 8px; letter-spacing: 1px; text-transform: uppercase; }
|
||||
.cls-btn {
|
||||
width: 100%; text-align: left; background: none; border: none;
|
||||
color: var(--text); font-family: 'DM Mono', monospace; font-size: 12px;
|
||||
padding: 6px 10px; border-radius: var(--radius); cursor: pointer;
|
||||
display: flex; justify-content: space-between; align-items: center;
|
||||
transition: .12s;
|
||||
}
|
||||
.cls-btn:hover { background: var(--border); }
|
||||
.cls-btn.active { background: rgba(0,229,160,.12); color: var(--accent); }
|
||||
.cls-badge { font-size: 10px; background: var(--border); border-radius: 99px;
|
||||
padding: 1px 7px; color: var(--muted); }
|
||||
.cls-btn.active .cls-badge { background: rgba(0,229,160,.2); color: var(--accent); }
|
||||
|
||||
/* ── main ── */
|
||||
main { overflow-y: auto; padding: 20px 24px; }
|
||||
|
||||
.toolbar {
|
||||
display: flex; align-items: center; gap: 10px;
|
||||
margin-bottom: 16px; flex-wrap: wrap;
|
||||
}
|
||||
.toolbar-label { color: var(--muted); font-size: 11px; margin-right: 4px; }
|
||||
.filt-btn {
|
||||
padding: 4px 10px; border-radius: 99px; font-size: 11px;
|
||||
cursor: pointer; border: 1px solid var(--border); background: transparent;
|
||||
color: var(--muted); font-family: 'DM Mono', monospace; transition: .12s;
|
||||
}
|
||||
.filt-btn:hover { border-color: var(--text); color: var(--text); }
|
||||
.filt-btn.on { border-color: var(--accent); color: var(--accent); background: rgba(0,229,160,.08); }
|
||||
|
||||
.conf-filter { display: flex; align-items: center; gap: 8px; margin-left: auto; font-size: 11px; color: var(--muted); }
|
||||
.conf-filter input { accent-color: var(--accent); width: 90px; }
|
||||
|
||||
/* bulk actions */
|
||||
.bulk-bar {
|
||||
display: none; align-items: center; gap: 10px; padding: 10px 14px;
|
||||
background: rgba(0,229,160,.06); border: 1px solid rgba(0,229,160,.2);
|
||||
border-radius: var(--radius); margin-bottom: 14px; font-size: 12px;
|
||||
}
|
||||
.bulk-bar.visible { display: flex; }
|
||||
|
||||
/* grid */
|
||||
.grid { display: grid; grid-template-columns: repeat(auto-fill, minmax(148px, 1fr)); gap: 10px; }
|
||||
|
||||
.card {
|
||||
position: relative; border-radius: var(--radius); overflow: hidden;
|
||||
border: 2px solid var(--border); cursor: pointer;
|
||||
transition: border-color .15s, transform .15s;
|
||||
background: var(--surface);
|
||||
}
|
||||
.card:hover { transform: translateY(-2px); border-color: var(--muted); }
|
||||
.card.selected { border-color: var(--accent) !important; }
|
||||
.card.confirmed { border-color: var(--accent); }
|
||||
.card.reclassify { border-color: var(--warn); }
|
||||
.card.rejected { border-color: var(--danger); opacity: .45; }
|
||||
|
||||
.card img { width: 100%; aspect-ratio: 1; object-fit: cover; display: block; }
|
||||
|
||||
.card-body { padding: 6px 8px 8px; }
|
||||
.card-class { font-size: 11px; color: var(--text); white-space: nowrap;
|
||||
overflow: hidden; text-overflow: ellipsis; font-weight: 500; }
|
||||
.card-conf { font-size: 10px; color: var(--muted); margin-top: 1px; }
|
||||
.card-status { font-size: 9px; margin-top: 3px; text-transform: uppercase; letter-spacing: .5px; }
|
||||
.card-status.auto { color: var(--muted); }
|
||||
.card-status.confirmed { color: var(--accent); }
|
||||
.card-status.reclassify { color: var(--warn); }
|
||||
.card-status.rejected { color: var(--danger); }
|
||||
|
||||
.card .sel-check {
|
||||
position: absolute; top: 5px; left: 5px;
|
||||
width: 18px; height: 18px; border-radius: 4px;
|
||||
border: 1.5px solid rgba(255,255,255,.3); background: rgba(0,0,0,.4);
|
||||
display: flex; align-items: center; justify-content: center;
|
||||
font-size: 11px; transition: .12s;
|
||||
}
|
||||
.card.selected .sel-check { background: var(--accent); border-color: var(--accent); color: #000; }
|
||||
|
||||
/* ── modal ── */
|
||||
.overlay { position: fixed; inset: 0; background: rgba(0,0,0,.75); z-index: 200;
|
||||
display: none; align-items: center; justify-content: center; }
|
||||
.overlay.open { display: flex; }
|
||||
.modal {
|
||||
background: var(--surface); border: 1px solid var(--border);
|
||||
border-radius: 10px; width: 440px; max-width: 95vw;
|
||||
padding: 24px; position: relative;
|
||||
}
|
||||
.modal h2 { font-family:'Syne',sans-serif; font-size:15px; margin-bottom:14px; }
|
||||
.modal img { width: 100%; border-radius: var(--radius); margin-bottom: 14px; object-fit: contain; max-height: 220px; }
|
||||
.modal-meta { font-size: 11px; color: var(--muted); margin-bottom: 14px; line-height: 1.8; }
|
||||
.modal-meta strong { color: var(--text); }
|
||||
.top3 { display: flex; gap: 8px; margin-bottom: 16px; }
|
||||
.top3-item { flex: 1; padding: 8px; border-radius: var(--radius); border: 1px solid var(--border);
|
||||
font-size: 10px; cursor: pointer; text-align: center; transition: .12s; }
|
||||
.top3-item:hover { border-color: var(--warn); color: var(--warn); }
|
||||
.top3-item .t3c { font-weight: 500; font-size: 11px; white-space: nowrap;
|
||||
overflow: hidden; text-overflow: ellipsis; }
|
||||
.top3-item .t3p { color: var(--muted); margin-top: 2px; }
|
||||
|
||||
.cls-select { width: 100%; background: var(--bg); border: 1px solid var(--border);
|
||||
color: var(--text); padding: 7px 10px; border-radius: var(--radius);
|
||||
font-family:'DM Mono',monospace; font-size:12px; margin-bottom: 14px; }
|
||||
.cls-select:focus { outline: 1px solid var(--accent); }
|
||||
|
||||
.modal-actions { display: flex; gap: 8px; }
|
||||
.modal-close { position: absolute; top: 14px; right: 14px; background: none; border: none;
|
||||
color: var(--muted); cursor: pointer; font-size: 18px; line-height: 1; }
|
||||
|
||||
/* ── commit modal ── */
|
||||
.commit-result { font-size: 12px; line-height: 2; }
|
||||
.commit-result .k { color: var(--muted); }
|
||||
.commit-result .v { color: var(--accent); font-weight: 500; }
|
||||
.commit-result .vw { color: var(--warn); font-weight: 500; }
|
||||
.commit-result .vr { color: var(--danger); font-weight: 500; }
|
||||
|
||||
/* empty state */
|
||||
.empty { text-align: center; color: var(--muted); padding: 60px 20px; }
|
||||
.empty .big { font-size: 40px; margin-bottom: 12px; }
|
||||
|
||||
/* progress overlay */
|
||||
.loading {
|
||||
position: fixed; inset: 0; background: var(--bg);
|
||||
display: flex; flex-direction: column; align-items: center; justify-content: center;
|
||||
z-index: 999; gap: 16px;
|
||||
}
|
||||
.loading h2 { font-family:'Syne',sans-serif; font-size: 18px; color: var(--accent); }
|
||||
.loading p { color: var(--muted); font-size: 12px; }
|
||||
.spinner { width: 36px; height: 36px; border: 3px solid var(--border);
|
||||
border-top-color: var(--accent); border-radius: 50%; animation: spin .7s linear infinite; }
|
||||
@keyframes spin { to { transform: rotate(360deg); } }
|
||||
|
||||
scrollbar-width: thin;
|
||||
::-webkit-scrollbar { width: 5px; }
|
||||
::-webkit-scrollbar-track { background: transparent; }
|
||||
::-webkit-scrollbar-thumb { background: var(--border); border-radius: 99px; }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
|
||||
<!-- loading -->
|
||||
<div class="loading" id="loading">
|
||||
<div class="spinner"></div>
|
||||
<h2>Auto-Label Review</h2>
|
||||
<p id="loading-msg">Loading crops…</p>
|
||||
</div>
|
||||
|
||||
<!-- header -->
|
||||
<header>
|
||||
<h1>⬡ AutoLabel</h1>
|
||||
<div class="stats-bar" id="stats-bar"></div>
|
||||
<div class="header-actions">
|
||||
<button class="btn-outline" style="display:inline-flex;align-items:center;gap:6px;padding:7px 14px;border-radius:6px;font-family:'DM Mono',monospace;font-size:12px;cursor:pointer;border:1px solid var(--border);background:transparent;color:var(--text);" onclick="confirmAll()">✓ Confirm all visible</button>
|
||||
<button class="btn-primary" style="display:inline-flex;align-items:center;gap:6px;padding:7px 16px;border-radius:6px;font-family:'DM Mono',monospace;font-size:12px;cursor:pointer;border:none;background:var(--accent);color:#000;font-weight:500;" onclick="openCommit()">Commit dataset →</button>
|
||||
</div>
|
||||
</header>
|
||||
|
||||
<div class="layout">
|
||||
|
||||
<!-- sidebar: class list -->
|
||||
<aside id="sidebar">
|
||||
<div class="sidebar-title">Classes</div>
|
||||
<button class="cls-btn active" data-cls="__all__" onclick="filterClass(this)">
|
||||
All <span class="cls-badge" id="badge-all">0</span>
|
||||
</button>
|
||||
<div id="class-list"></div>
|
||||
</aside>
|
||||
|
||||
<!-- main content -->
|
||||
<main>
|
||||
<!-- toolbar -->
|
||||
<div class="toolbar">
|
||||
<span class="toolbar-label">Status:</span>
|
||||
<button class="filt-btn on" data-st="auto" onclick="toggleStatus(this)">Auto</button>
|
||||
<button class="filt-btn on" data-st="confirmed" onclick="toggleStatus(this)">Confirmed</button>
|
||||
<button class="filt-btn on" data-st="reclassify" onclick="toggleStatus(this)">Reclassified</button>
|
||||
<button class="filt-btn on" data-st="rejected" onclick="toggleStatus(this)">Rejected</button>
|
||||
|
||||
<div class="conf-filter">
|
||||
Min conf: <input type="range" min="0" max="1" step=".05" value="0" oninput="confThresh=+this.value;renderGrid()">
|
||||
<span id="conf-val">0.00</span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- bulk action bar -->
|
||||
<div class="bulk-bar" id="bulk-bar">
|
||||
<span id="bulk-count">0 selected</span>
|
||||
<button class="btn-outline" style="display:inline-flex;align-items:center;gap:6px;padding:5px 12px;border-radius:6px;font-family:'DM Mono',monospace;font-size:11px;cursor:pointer;border:1px solid var(--border);background:transparent;color:var(--text);" onclick="bulkAction('confirmed')">✓ Confirm</button>
|
||||
<button class="btn-outline" style="display:inline-flex;align-items:center;gap:6px;padding:5px 12px;border-radius:6px;font-family:'DM Mono',monospace;font-size:11px;cursor:pointer;border:1px solid var(--border);background:transparent;color:var(--text);" onclick="openBulkReclassify()">↷ Reclassify</button>
|
||||
<button style="display:inline-flex;align-items:center;gap:6px;padding:5px 12px;border-radius:6px;font-family:'DM Mono',monospace;font-size:11px;cursor:pointer;border:none;background:var(--danger);color:#fff;" onclick="bulkAction('rejected')">✕ Reject</button>
|
||||
<button class="btn-outline" style="margin-left:auto;display:inline-flex;align-items:center;gap:6px;padding:5px 12px;border-radius:6px;font-family:'DM Mono',monospace;font-size:11px;cursor:pointer;border:1px solid var(--border);background:transparent;color:var(--text);" onclick="clearSelection()">Clear</button>
|
||||
</div>
|
||||
|
||||
<div class="grid" id="grid"></div>
|
||||
</main>
|
||||
</div>
|
||||
|
||||
<!-- crop detail modal -->
|
||||
<div class="overlay" id="modal-overlay" onclick="if(event.target===this)closeModal()">
|
||||
<div class="modal">
|
||||
<button class="modal-close" onclick="closeModal()">✕</button>
|
||||
<h2 id="modal-title">Crop detail</h2>
|
||||
<img id="modal-img" src="" alt="">
|
||||
<div class="modal-meta" id="modal-meta"></div>
|
||||
<div style="font-size:10px;color:var(--muted);margin-bottom:6px;text-transform:uppercase;letter-spacing:.8px;">Top-3 suggestions</div>
|
||||
<div class="top3" id="modal-top3"></div>
|
||||
<div style="font-size:10px;color:var(--muted);margin-bottom:6px;text-transform:uppercase;letter-spacing:.8px;">Override class</div>
|
||||
<select class="cls-select" id="modal-cls-select"></select>
|
||||
<div class="modal-actions">
|
||||
<button class="btn-primary" style="flex:1;display:inline-flex;align-items:center;justify-content:center;gap:6px;padding:9px;border-radius:6px;font-family:'DM Mono',monospace;font-size:12px;cursor:pointer;border:none;background:var(--accent);color:#000;font-weight:500;" onclick="modalConfirm()">✓ Confirm</button>
|
||||
<button class="btn-outline" style="flex:1;display:inline-flex;align-items:center;justify-content:center;gap:6px;padding:9px;border-radius:6px;font-family:'DM Mono',monospace;font-size:12px;cursor:pointer;border:1px solid var(--border);background:transparent;color:var(--text);" onclick="modalReclassify()">↷ Apply override</button>
|
||||
<button style="flex:1;display:inline-flex;align-items:center;justify-content:center;gap:6px;padding:9px;border-radius:6px;font-family:'DM Mono',monospace;font-size:12px;cursor:pointer;border:none;background:var(--danger);color:#fff;" onclick="modalReject()">✕ Reject</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- bulk reclassify modal -->
|
||||
<div class="overlay" id="bulk-modal-overlay" onclick="if(event.target===this)closeBulkModal()">
|
||||
<div class="modal">
|
||||
<button class="modal-close" onclick="closeBulkModal()">✕</button>
|
||||
<h2>Reclassify <span id="bulk-modal-count"></span> crops</h2>
|
||||
<div style="font-size:11px;color:var(--muted);margin-bottom:12px;">All selected crops will be moved to this class.</div>
|
||||
<select class="cls-select" id="bulk-cls-select"></select>
|
||||
<div class="modal-actions">
|
||||
<button class="btn-primary" style="flex:1;display:inline-flex;align-items:center;justify-content:center;padding:9px;border-radius:6px;font-family:'DM Mono',monospace;font-size:12px;cursor:pointer;border:none;background:var(--accent);color:#000;font-weight:500;" onclick="applyBulkReclassify()">Apply</button>
|
||||
<button class="btn-outline" style="flex:1;display:inline-flex;align-items:center;justify-content:center;padding:9px;border-radius:6px;font-family:'DM Mono',monospace;font-size:12px;cursor:pointer;border:1px solid var(--border);background:transparent;color:var(--text);" onclick="closeBulkModal()">Cancel</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- commit modal -->
|
||||
<div class="overlay" id="commit-overlay" onclick="if(event.target===this)closeCommit()">
|
||||
<div class="modal">
|
||||
<button class="modal-close" onclick="closeCommit()">✕</button>
|
||||
<h2 id="commit-title">Commit dataset</h2>
|
||||
<div id="commit-body"></div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
let ALL_CROPS = [];
|
||||
let ALL_CLASSES= [];
|
||||
let activeClass= '__all__';
|
||||
let activeStatuses = new Set(['auto','confirmed','reclassify','rejected']);
|
||||
let confThresh = 0;
|
||||
let selected = new Set(); // crop ids
|
||||
let modalCropId= null;
|
||||
|
||||
// ── boot ──────────────────────────────────────────────────────────────────
|
||||
async function boot() {
|
||||
document.getElementById('loading-msg').textContent = 'Fetching crops from server…';
|
||||
const [cropsRes, clsRes] = await Promise.all([fetch('/api/crops'), fetch('/api/classes')]);
|
||||
ALL_CROPS = await cropsRes.json();
|
||||
ALL_CLASSES = await clsRes.json();
|
||||
|
||||
// Populate selects
|
||||
[document.getElementById('modal-cls-select'),
|
||||
document.getElementById('bulk-cls-select')].forEach(sel => {
|
||||
sel.innerHTML = ALL_CLASSES.map(c => `<option value="${c}">${c}</option>`).join('');
|
||||
});
|
||||
|
||||
buildSidebar();
|
||||
renderGrid();
|
||||
renderStats();
|
||||
document.getElementById('loading').style.display = 'none';
|
||||
}
|
||||
|
||||
// ── sidebar ───────────────────────────────────────────────────────────────
|
||||
function buildSidebar() {
|
||||
const counts = {};
|
||||
ALL_CROPS.forEach(c => counts[c.pred_class] = (counts[c.pred_class]||0)+1);
|
||||
const list = document.getElementById('class-list');
|
||||
list.innerHTML = ALL_CLASSES
|
||||
.filter(cls => counts[cls])
|
||||
.sort()
|
||||
.map(cls => `
|
||||
<button class="cls-btn" data-cls="${cls}" onclick="filterClass(this)">
|
||||
<span style="overflow:hidden;text-overflow:ellipsis;white-space:nowrap">${cls}</span>
|
||||
<span class="cls-badge" id="badge-${cls}">${counts[cls]||0}</span>
|
||||
</button>`).join('');
|
||||
document.getElementById('badge-all').textContent = ALL_CROPS.length;
|
||||
}
|
||||
|
||||
function filterClass(btn) {
|
||||
document.querySelectorAll('.cls-btn').forEach(b=>b.classList.remove('active'));
|
||||
btn.classList.add('active');
|
||||
activeClass = btn.dataset.cls;
|
||||
clearSelection();
|
||||
renderGrid();
|
||||
}
|
||||
|
||||
function toggleStatus(btn) {
|
||||
btn.classList.toggle('on');
|
||||
const st = btn.dataset.st;
|
||||
activeStatuses.has(st) ? activeStatuses.delete(st) : activeStatuses.add(st);
|
||||
renderGrid();
|
||||
}
|
||||
|
||||
// ── grid ──────────────────────────────────────────────────────────────────
|
||||
function visibleCrops() {
|
||||
return ALL_CROPS.filter(c => {
|
||||
if (activeClass !== '__all__' && c.pred_class !== activeClass) return false;
|
||||
if (!activeStatuses.has(c.status)) return false;
|
||||
if (c.confidence < confThresh) return false;
|
||||
return true;
|
||||
});
|
||||
}
|
||||
|
||||
function renderGrid() {
|
||||
const crops = visibleCrops();
|
||||
const grid = document.getElementById('grid');
|
||||
|
||||
// Update conf label
|
||||
document.getElementById('conf-val').textContent = confThresh.toFixed(2);
|
||||
|
||||
if (!crops.length) {
|
||||
grid.innerHTML = `<div class="empty" style="grid-column:1/-1">
|
||||
<div class="big">🔍</div>No crops match the current filters.</div>`;
|
||||
renderStats();
|
||||
return;
|
||||
}
|
||||
|
||||
grid.innerHTML = crops.map(c => `
|
||||
<div class="card ${c.status} ${selected.has(c.id)?'selected':''}"
|
||||
id="card-${c.id}"
|
||||
onclick="handleCardClick(event,'${c.id}')">
|
||||
<div class="sel-check">${selected.has(c.id)?'✓':''}</div>
|
||||
<img loading="lazy" src="/api/crop_image/${c.id}" alt="">
|
||||
<div class="card-body">
|
||||
<div class="card-class" title="${c.label}">${c.label}</div>
|
||||
<div class="card-conf">${(c.confidence*100).toFixed(1)}%</div>
|
||||
<div class="card-status ${c.status}">${c.status}</div>
|
||||
</div>
|
||||
</div>`).join('');
|
||||
renderStats();
|
||||
}
|
||||
|
||||
function handleCardClick(evt, id) {
|
||||
if (evt.shiftKey || evt.ctrlKey || evt.metaKey) {
|
||||
// multi-select
|
||||
selected.has(id) ? selected.delete(id) : selected.add(id);
|
||||
renderGrid();
|
||||
updateBulkBar();
|
||||
} else {
|
||||
openModal(id);
|
||||
}
|
||||
}
|
||||
|
||||
// ── selection / bulk ──────────────────────────────────────────────────────
|
||||
function clearSelection() { selected.clear(); renderGrid(); updateBulkBar(); }
|
||||
|
||||
function updateBulkBar() {
|
||||
const bar = document.getElementById('bulk-bar');
|
||||
document.getElementById('bulk-count').textContent = `${selected.size} selected`;
|
||||
bar.classList.toggle('visible', selected.size > 0);
|
||||
}
|
||||
|
||||
function bulkAction(status) {
|
||||
selected.forEach(id => {
|
||||
const c = ALL_CROPS.find(x=>x.id===id);
|
||||
if (c) c.status = status;
|
||||
});
|
||||
syncUpdates([...selected]);
|
||||
clearSelection();
|
||||
renderGrid();
|
||||
}
|
||||
|
||||
function openBulkReclassify() {
|
||||
document.getElementById('bulk-modal-count').textContent = selected.size;
|
||||
document.getElementById('bulk-modal-overlay').classList.add('open');
|
||||
}
|
||||
function closeBulkModal() { document.getElementById('bulk-modal-overlay').classList.remove('open'); }
|
||||
function applyBulkReclassify() {
|
||||
const cls = document.getElementById('bulk-cls-select').value;
|
||||
selected.forEach(id => {
|
||||
const c = ALL_CROPS.find(x=>x.id===id);
|
||||
if (c) { c.label = cls; c.status = 'reclassify'; }
|
||||
});
|
||||
syncUpdates([...selected]);
|
||||
closeBulkModal();
|
||||
clearSelection();
|
||||
renderGrid();
|
||||
}
|
||||
|
||||
function confirmAll() {
|
||||
const ids = visibleCrops().filter(c=>c.status==='auto').map(c=>c.id);
|
||||
ids.forEach(id => { const c = ALL_CROPS.find(x=>x.id===id); if(c) c.status='confirmed'; });
|
||||
syncUpdates(ids);
|
||||
renderGrid();
|
||||
}
|
||||
|
||||
// ── modal ─────────────────────────────────────────────────────────────────
|
||||
function openModal(id) {
|
||||
modalCropId = id;
|
||||
const c = ALL_CROPS.find(x=>x.id===id);
|
||||
document.getElementById('modal-title').textContent = c.id;
|
||||
document.getElementById('modal-img').src = `/api/crop_image/${id}`;
|
||||
document.getElementById('modal-meta').innerHTML = `
|
||||
<strong>Predicted:</strong> ${c.pred_class}<br>
|
||||
<strong>Label:</strong> ${c.label}<br>
|
||||
<strong>Confidence:</strong>${(c.confidence*100).toFixed(1)}%<br>
|
||||
<strong>Source:</strong> ${c.source_image}`;
|
||||
|
||||
document.getElementById('modal-top3').innerHTML = (c.top3||[]).map(t=>`
|
||||
<div class="top3-item" onclick="pickTop3('${t.class}')">
|
||||
<div class="t3c">${t.class}</div>
|
||||
<div class="t3p">${(t.conf*100).toFixed(1)}%</div>
|
||||
</div>`).join('');
|
||||
|
||||
// set select to current label
|
||||
document.getElementById('modal-cls-select').value = c.label;
|
||||
document.getElementById('modal-overlay').classList.add('open');
|
||||
}
|
||||
function closeModal() { document.getElementById('modal-overlay').classList.remove('open'); }
|
||||
|
||||
function pickTop3(cls) {
|
||||
document.getElementById('modal-cls-select').value = cls;
|
||||
}
|
||||
function modalConfirm() {
|
||||
const c = ALL_CROPS.find(x=>x.id===modalCropId);
|
||||
if(c) { c.status='confirmed'; syncUpdates([modalCropId]); }
|
||||
closeModal(); renderGrid();
|
||||
}
|
||||
function modalReclassify() {
|
||||
const cls = document.getElementById('modal-cls-select').value;
|
||||
const c = ALL_CROPS.find(x=>x.id===modalCropId);
|
||||
if(c) { c.label=cls; c.status='reclassify'; syncUpdates([modalCropId]); }
|
||||
closeModal(); renderGrid();
|
||||
}
|
||||
function modalReject() {
|
||||
const c = ALL_CROPS.find(x=>x.id===modalCropId);
|
||||
if(c) { c.status='rejected'; syncUpdates([modalCropId]); }
|
||||
closeModal(); renderGrid();
|
||||
}
|
||||
|
||||
// ── stats bar ─────────────────────────────────────────────────────────────
|
||||
function renderStats() {
|
||||
const st = { auto:0, confirmed:0, reclassify:0, rejected:0 };
|
||||
ALL_CROPS.forEach(c => { if(st[c.status]!==undefined) st[c.status]++; });
|
||||
document.getElementById('stats-bar').innerHTML = `
|
||||
<span><span class="dot dot-auto"></span> Auto ${st.auto}</span>
|
||||
<span><span class="dot dot-confirmed"></span> Confirmed ${st.confirmed}</span>
|
||||
<span><span class="dot dot-reclassify"></span> Reclassified ${st.reclassify}</span>
|
||||
<span><span class="dot dot-rejected"></span> Rejected ${st.rejected}</span>
|
||||
<span style="margin-left:8px"> Total ${ALL_CROPS.length}</span>`;
|
||||
}
|
||||
|
||||
// ── sync to server ─────────────────────────────────────────────────────────
|
||||
async function syncUpdates(ids) {
|
||||
const updates = ids.map(id => {
|
||||
const c = ALL_CROPS.find(x=>x.id===id);
|
||||
return c ? { id: c.id, label: c.label, status: c.status } : null;
|
||||
}).filter(Boolean);
|
||||
await fetch('/api/update', { method:'POST',
|
||||
headers:{'Content-Type':'application/json'}, body:JSON.stringify(updates) });
|
||||
}
|
||||
|
||||
// ── commit ─────────────────────────────────────────────────────────────────
|
||||
function openCommit() {
|
||||
const auto = ALL_CROPS.filter(c=>c.status==='auto').length;
|
||||
document.getElementById('commit-title').textContent = 'Commit dataset';
|
||||
document.getElementById('commit-body').innerHTML = `
|
||||
<div style="font-size:12px;color:var(--muted);margin-bottom:16px;line-height:1.7">
|
||||
${auto > 0
|
||||
? `<span style="color:var(--warn)">⚠ ${auto} crops still have status "auto" (unreviewed).<br>They will be saved to <code>_unreviewed_/</code>.</span><br><br>`
|
||||
: ''}
|
||||
This will copy all crops to <code style="color:var(--accent)">${'data/'}</code>.<br>
|
||||
Confirmed + reclassified → their class folder.<br>
|
||||
Rejected → <code>_rejected_/</code><br>
|
||||
Unreviewed → <code>_unreviewed_/</code>
|
||||
</div>
|
||||
<div style="display:flex;gap:8px">
|
||||
<button class="btn-primary" style="flex:1;display:inline-flex;align-items:center;justify-content:center;padding:10px;border-radius:6px;font-family:'DM Mono',monospace;font-size:12px;cursor:pointer;border:none;background:var(--accent);color:#000;font-weight:500;" onclick="doCommit()">Commit now</button>
|
||||
<button class="btn-outline" style="display:inline-flex;align-items:center;justify-content:center;padding:10px 16px;border-radius:6px;font-family:'DM Mono',monospace;font-size:12px;cursor:pointer;border:1px solid var(--border);background:transparent;color:var(--text);" onclick="closeCommit()">Cancel</button>
|
||||
</div>`;
|
||||
document.getElementById('commit-overlay').classList.add('open');
|
||||
}
|
||||
function closeCommit() { document.getElementById('commit-overlay').classList.remove('open'); }
|
||||
|
||||
async function doCommit() {
|
||||
document.getElementById('commit-body').innerHTML = '<div style="text-align:center;padding:20px"><div class="spinner" style="margin:0 auto"></div></div>';
|
||||
const res = await fetch('/api/commit', { method:'POST' });
|
||||
const data = await res.json();
|
||||
document.getElementById('commit-title').textContent = '✓ Done!';
|
||||
document.getElementById('commit-body').innerHTML = `
|
||||
<div class="commit-result">
|
||||
<div><span class="k">Committed </span><span class="v">${data.committed}</span></div>
|
||||
<div><span class="k">Rejected </span><span class="vr">${data.rejected}</span></div>
|
||||
<div><span class="k">Unreviewed </span><span class="vw">${data.unreviewed}</span></div>
|
||||
<div style="margin-top:12px"><span class="k">Saved to </span><code style="color:var(--accent)">${data.output_dir}</code></div>
|
||||
</div>
|
||||
<div style="margin-top:18px;font-size:11px;color:var(--muted)">
|
||||
Next step:<br>
|
||||
<code style="color:var(--text)">python 1b_split_dataset.py --data_dir data --output_dir crops_dataset</code>
|
||||
</div>`;
|
||||
}
|
||||
|
||||
// keyboard shortcuts
|
||||
document.addEventListener('keydown', e => {
|
||||
if (!modalCropId) return;
|
||||
if (e.key==='ArrowRight'||e.key==='ArrowLeft'||e.key==='d'||e.key==='a'||e.key==='q') {
|
||||
const visible = visibleCrops();
|
||||
const idx = visible.findIndex(c=>c.id===modalCropId);
|
||||
const next = (e.key==='ArrowRight' || e.key==='d') ? idx+1 : idx-1;
|
||||
if (next>=0 && next<visible.length) openModal(visible[next].id);
|
||||
}
|
||||
if (e.key==='c') modalConfirm();
|
||||
if (e.key==='r') modalReject();
|
||||
if (e.key==='Escape') closeModal();
|
||||
});
|
||||
|
||||
boot();
|
||||
</script>
|
||||
</body>
|
||||
</html>"""
|
||||
|
||||
# ─────────────────────────── entry point ─────────────────────────────────────
|
||||
|
||||
def main():
|
||||
args = ARGS
|
||||
|
||||
print(f"[startup] Device : {DEVICE}")
|
||||
print(f"[startup] Detector : {args.detector_weights}")
|
||||
print(f"[startup] Classifier : {args.classifier_weights}")
|
||||
|
||||
detector = YOLO(args.detector_weights)
|
||||
classifier = ProductClassifier(args.classifier_weights, DEVICE)
|
||||
|
||||
print("\n[step 1/2] Running detection + auto-classification …")
|
||||
crops = run_pipeline(detector, classifier, args)
|
||||
|
||||
if not crops:
|
||||
print("⚠ No crops were generated. Check --source and --det_conf.")
|
||||
sys.exit(1)
|
||||
|
||||
STATE["crops"] = crops
|
||||
STATE["classes"] = classifier.class_names
|
||||
|
||||
print(f"\n[step 2/2] Launching review UI on http://localhost:{args.port}")
|
||||
print(" Keyboard shortcuts inside a crop:")
|
||||
print(" C → confirm")
|
||||
print(" R → reject")
|
||||
print(" ← → → previous / next crop")
|
||||
print(" Esc → close modal")
|
||||
print(" Ctrl/Shift+click → multi-select for bulk actions")
|
||||
print("\n Close the browser tab or press Ctrl+C to stop.\n")
|
||||
|
||||
# Open browser after a short delay so Flask is ready
|
||||
def _open():
|
||||
time.sleep(1.2)
|
||||
webbrowser.open(f"http://localhost:{args.port}")
|
||||
threading.Thread(target=_open, daemon=True).start()
|
||||
|
||||
# Run Flask (single-threaded to keep GPU access safe)
|
||||
flask_app.run(host="0.0.0.0", port=args.port, threaded=False, use_reloader=False)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
32
class_names.json
Normal file
32
class_names.json
Normal file
@ -0,0 +1,32 @@
|
||||
{
|
||||
"0": "1001",
|
||||
"1": "1001 Gold",
|
||||
"2": "A100",
|
||||
"3": "Aroma",
|
||||
"4": "Aroma Espresso",
|
||||
"5": "Aroma Familial",
|
||||
"6": "Aroma Gold",
|
||||
"7": "Bonal",
|
||||
"8": "Bonal Excellence",
|
||||
"9": "Cafe D Or Caps",
|
||||
"10": "Canastel",
|
||||
"11": "Caps",
|
||||
"12": "CityOne Instant Coffee",
|
||||
"13": "Dozia",
|
||||
"14": "El Kabir",
|
||||
"15": "Facto",
|
||||
"16": "Facto no",
|
||||
"17": "Famico",
|
||||
"18": "Famico Caps",
|
||||
"19": "Famico Exclusive",
|
||||
"20": "Molino",
|
||||
"21": "Mundo",
|
||||
"22": "Mundo Caps",
|
||||
"23": "Nizier",
|
||||
"24": "Nouara",
|
||||
"25": "Oscar",
|
||||
"26": "Primo",
|
||||
"27": "Ricamar Thon",
|
||||
"28": "Siglo",
|
||||
"29": "Skor"
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 5.3 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 4.4 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.9 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 4.0 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 4.9 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.8 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 3.6 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.8 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.6 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.8 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 3.0 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.5 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.7 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.9 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.4 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 3.0 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.6 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.7 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 7.7 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 8.4 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 6.4 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 3.6 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.8 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.6 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 5.3 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 4.4 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.9 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 4.0 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 4.9 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.8 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.8 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 3.0 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.5 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.7 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.9 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.4 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 3.0 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.6 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.7 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 4.4 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 3.8 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.7 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.9 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.4 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.4 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 4.2 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.0 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.1 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.0 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 9.7 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 8.5 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 11 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 7.5 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 12 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 13 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 3.3 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.8 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.8 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 3.4 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 3.5 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 3.4 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.8 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.8 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 3.2 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 3.0 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.6 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.8 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.5 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.3 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 6.0 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 7.5 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 6.4 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 7.3 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 8.9 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 5.6 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 4.3 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 6.9 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 9.2 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.4 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.6 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.4 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.9 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.8 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.1 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.0 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.9 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.2 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.3 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.9 KiB |
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in New Issue
Block a user