987 lines
42 KiB
Python
987 lines
42 KiB
Python
"""
|
|
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="runs/classify/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()
|