1045 lines
38 KiB
Python
1045 lines
38 KiB
Python
"""
|
||
class_labeller.py — Class Labelling Webapp
|
||
============================================
|
||
A high-performance web app for assigning class labels to crop images.
|
||
Handles 14,000+ images via:
|
||
• Parallel thumbnail generation (multiprocessing)
|
||
• On-disk thumbnail + metadata cache (skips re-processing)
|
||
• Virtual scrolling (only visible DOM nodes rendered)
|
||
• Color-grouped layout so visually similar products sit together
|
||
|
||
Workflow
|
||
--------
|
||
1. Run this script pointing at your data/unknown/ (or any flat image folder)
|
||
2. Browser opens — images are sorted by dominant color into groups
|
||
3. Click / Shift-click / Ctrl-click to select images
|
||
4. Type a class name (or pick existing) → Assign
|
||
5. All assigned images are MOVED to data/<class_name>/
|
||
6. Remaining unassigned images stay in the source folder
|
||
|
||
Usage
|
||
-----
|
||
python class_labeller.py \
|
||
--source data/unknown \
|
||
--data_dir data \
|
||
--port 5001 \
|
||
--workers 8 \
|
||
--thumb_size 112
|
||
"""
|
||
|
||
import argparse
|
||
import hashlib
|
||
import json
|
||
import math
|
||
import multiprocessing
|
||
import os
|
||
import shutil
|
||
import time
|
||
from functools import partial
|
||
from pathlib import Path
|
||
|
||
import cv2
|
||
import numpy as np
|
||
from flask import Flask, Response, jsonify, request, send_file
|
||
|
||
|
||
# ─────────────────────────── args ────────────────────────────────────────────
|
||
|
||
def parse_args():
|
||
p = argparse.ArgumentParser()
|
||
p.add_argument("--source", default="data/unknown",
|
||
help="Folder of unlabelled images")
|
||
p.add_argument("--data_dir", default="data",
|
||
help="Root data folder — assigned images go to data/<class>/")
|
||
p.add_argument("--port", type=int, default=5001)
|
||
p.add_argument("--workers", type=int,
|
||
default=max(1, multiprocessing.cpu_count() - 1),
|
||
help="Parallel workers for thumbnail generation")
|
||
p.add_argument("--thumb_size", type=int, default=112,
|
||
help="Thumbnail size in pixels (square)")
|
||
p.add_argument("--cache_dir", default=".label_cache",
|
||
help="Where thumbnails and metadata are cached")
|
||
p.add_argument("--n_colors", type=int, default=16,
|
||
help="Number of color buckets for grouping")
|
||
return p.parse_args()
|
||
|
||
ARGS = parse_args()
|
||
|
||
IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tiff"}
|
||
|
||
# ─────────────────────────── color helpers ────────────────────────────────────
|
||
|
||
# Named color buckets (hue ranges in OpenCV HSV, hue 0-179)
|
||
COLOR_BUCKETS = [
|
||
( 0, 10, "Red" ),
|
||
( 10, 20, "Orange"),
|
||
( 20, 35, "Yellow"),
|
||
( 35, 80, "Green" ),
|
||
( 80, 95, "Cyan" ),
|
||
( 95, 130, "Blue" ),
|
||
(130, 150, "Purple"),
|
||
(150, 170, "Pink" ),
|
||
(170, 180, "Red" ), # red wraps around
|
||
]
|
||
|
||
def dominant_color_hsv(img_bgr: np.ndarray) -> tuple[int, int, int]:
|
||
"""
|
||
Return the dominant color of an image as (H, S, V) in OpenCV ranges.
|
||
Strategy: resize to 32×32, convert to HSV, take the median of each channel.
|
||
Fast enough to run on thousands of images without a GPU.
|
||
"""
|
||
small = cv2.resize(img_bgr, (32, 32), interpolation=cv2.INTER_AREA)
|
||
hsv = cv2.cvtColor(small, cv2.COLOR_BGR2HSV)
|
||
h = int(np.median(hsv[:, :, 0]))
|
||
s = int(np.median(hsv[:, :, 1]))
|
||
v = int(np.median(hsv[:, :, 2]))
|
||
return h, s, v
|
||
|
||
def color_bucket(h: int, s: int, v: int) -> str:
|
||
"""Map (H, S, V) to a human-readable color name."""
|
||
if v < 50: return "Black"
|
||
if s < 40:
|
||
return "White" if v > 180 else "Gray"
|
||
for lo, hi, name in COLOR_BUCKETS:
|
||
if lo <= h < hi:
|
||
return name
|
||
return "Red"
|
||
|
||
BUCKET_ORDER = [
|
||
"Red","Orange","Yellow","Green","Cyan",
|
||
"Blue","Purple","Pink","White","Gray","Black"
|
||
]
|
||
|
||
def sort_key(meta: dict) -> tuple:
|
||
"""Sort key: bucket index → hue → saturation descending → value."""
|
||
b = meta["bucket"]
|
||
idx = BUCKET_ORDER.index(b) if b in BUCKET_ORDER else 99
|
||
return (idx, meta["h"], -meta["s"], -meta["v"])
|
||
|
||
# ─────────────────────────── thumbnail generation ────────────────────────────
|
||
|
||
def process_one(img_path_str: str, thumb_dir: str, size: int) -> dict | None:
|
||
"""
|
||
Worker function (runs in subprocess):
|
||
Read one image, make a thumbnail, extract dominant color.
|
||
Returns metadata dict or None on failure.
|
||
"""
|
||
try:
|
||
img_path = Path(img_path_str)
|
||
img = cv2.imread(img_path_str)
|
||
if img is None:
|
||
return None
|
||
|
||
# Thumbnail
|
||
h, w = img.shape[:2]
|
||
scale = size / max(h, w)
|
||
th, tw = int(h * scale), int(w * scale)
|
||
thumb = cv2.resize(img, (tw, th), interpolation=cv2.INTER_AREA)
|
||
|
||
# Center-pad to square
|
||
canvas = np.zeros((size, size, 3), dtype=np.uint8)
|
||
y0 = (size - th) // 2
|
||
x0 = (size - tw) // 2
|
||
canvas[y0:y0+th, x0:x0+tw] = thumb
|
||
|
||
# Save thumbnail — use a hash of the original path as filename
|
||
# so re-runs are skipped cleanly
|
||
key = hashlib.md5(img_path_str.encode()).hexdigest()
|
||
out_path = Path(thumb_dir) / f"{key}.jpg"
|
||
if not out_path.exists():
|
||
cv2.imwrite(str(out_path), canvas,
|
||
[cv2.IMWRITE_JPEG_QUALITY, 82])
|
||
|
||
dh, ds, dv = dominant_color_hsv(img)
|
||
bucket = color_bucket(dh, ds, dv)
|
||
|
||
return {
|
||
"id": key,
|
||
"path": img_path_str,
|
||
"name": img_path.name,
|
||
"bucket": bucket,
|
||
"h": dh, "s": ds, "v": dv,
|
||
"w": w, "height": h,
|
||
}
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def build_metadata(source_dir: Path, cache_dir: Path,
|
||
thumb_size: int, n_workers: int) -> list[dict]:
|
||
"""
|
||
Generate thumbnails for every image in source_dir (parallel),
|
||
cache results to disk, return sorted metadata list.
|
||
Re-running only processes new images.
|
||
"""
|
||
thumb_dir = cache_dir / "thumbs"
|
||
thumb_dir.mkdir(parents=True, exist_ok=True)
|
||
meta_file = cache_dir / "metadata.json"
|
||
|
||
# Load existing cache
|
||
cached: dict[str, dict] = {}
|
||
if meta_file.exists():
|
||
try:
|
||
for m in json.loads(meta_file.read_text()):
|
||
cached[m["path"]] = m
|
||
except Exception:
|
||
pass
|
||
|
||
# Find all images in source_dir
|
||
all_paths = sorted(
|
||
str(p) for p in source_dir.rglob("*")
|
||
if p.suffix.lower() in IMAGE_EXTS
|
||
)
|
||
|
||
# Only process images not in cache
|
||
new_paths = [p for p in all_paths if p not in cached]
|
||
|
||
if new_paths:
|
||
print(f"[cache] Processing {len(new_paths)} new images with {n_workers} workers…")
|
||
worker = partial(process_one, thumb_dir=str(thumb_dir), size=thumb_size)
|
||
|
||
with multiprocessing.Pool(processes=n_workers) as pool:
|
||
chunk = max(1, len(new_paths) // (n_workers * 4))
|
||
results = []
|
||
done = 0
|
||
for batch in [new_paths[i:i+chunk] for i in range(0, len(new_paths), chunk)]:
|
||
batch_results = pool.map(worker, batch)
|
||
for r in batch_results:
|
||
if r:
|
||
cached[r["path"]] = r
|
||
done += len(batch)
|
||
pct = done / len(new_paths) * 100
|
||
print(f" {done}/{len(new_paths)} ({pct:.0f}%)", end="\r", flush=True)
|
||
print()
|
||
else:
|
||
print(f"[cache] All {len(all_paths)} images already cached.")
|
||
|
||
# Keep only images that still exist in source_dir
|
||
meta = [cached[p] for p in all_paths if p in cached]
|
||
|
||
# Sort by color
|
||
meta.sort(key=sort_key)
|
||
|
||
# Persist cache
|
||
meta_file.write_text(json.dumps(meta))
|
||
print(f"[cache] {len(meta)} images ready.")
|
||
return meta
|
||
|
||
# ─────────────────────────── Flask app ───────────────────────────────────────
|
||
|
||
flask_app = Flask(__name__)
|
||
STATE: dict = {} # populated in main()
|
||
|
||
@flask_app.route("/")
|
||
def index():
|
||
return HTML
|
||
|
||
@flask_app.route("/api/meta")
|
||
def api_meta():
|
||
"""
|
||
Return lightweight metadata (no pixel data).
|
||
The client paginates via ?offset=&limit=
|
||
"""
|
||
offset = int(request.args.get("offset", 0))
|
||
limit = int(request.args.get("limit", 500))
|
||
bucket = request.args.get("bucket", "")
|
||
|
||
items = STATE["meta"]
|
||
if bucket and bucket != "__all__":
|
||
items = [m for m in items if m["bucket"] == bucket]
|
||
|
||
total = len(items)
|
||
page = items[offset : offset + limit]
|
||
# strip heavy fields before sending
|
||
slim = [{"id":m["id"],"name":m["name"],"bucket":m["bucket"],
|
||
"h":m["h"],"s":m["s"],"v":m["v"]} for m in page]
|
||
return jsonify({"total": total, "offset": offset, "items": slim})
|
||
|
||
@flask_app.route("/api/buckets")
|
||
def api_buckets():
|
||
counts: dict[str, int] = {}
|
||
for m in STATE["meta"]:
|
||
counts[m["bucket"]] = counts.get(m["bucket"], 0) + 1
|
||
ordered = [{"name": b, "count": counts[b]}
|
||
for b in BUCKET_ORDER if b in counts]
|
||
return jsonify(ordered)
|
||
|
||
@flask_app.route("/api/classes")
|
||
def api_classes():
|
||
data_dir = Path(ARGS.data_dir)
|
||
classes = sorted(
|
||
d.name for d in data_dir.iterdir()
|
||
if d.is_dir() and not d.name.startswith("_") and d.name != "unknown"
|
||
) if data_dir.exists() else []
|
||
return jsonify(classes)
|
||
|
||
@flask_app.route("/thumb/<img_id>")
|
||
def thumb(img_id: str):
|
||
path = Path(ARGS.cache_dir) / "thumbs" / f"{img_id}.jpg"
|
||
if not path.exists():
|
||
return "", 404
|
||
return send_file(str(path), mimetype="image/jpeg",
|
||
max_age=86400) # cache 24 h in browser
|
||
|
||
@flask_app.route("/api/assign", methods=["POST"])
|
||
def api_assign():
|
||
"""
|
||
Move selected images into data/<class_name>/.
|
||
Body: { "ids": [...], "class_name": "cola_can" }
|
||
"""
|
||
body = request.json
|
||
ids = set(body.get("ids", []))
|
||
class_name = body["class_name"].strip().replace(" ", "_")
|
||
|
||
if not class_name:
|
||
return jsonify({"error": "class_name is required"}), 400
|
||
|
||
dst_dir = Path(ARGS.data_dir) / class_name
|
||
dst_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
moved = 0
|
||
errors = []
|
||
id_map = {m["id"]: m for m in STATE["meta"]}
|
||
|
||
for img_id in ids:
|
||
meta = id_map.get(img_id)
|
||
if not meta:
|
||
continue
|
||
src = Path(meta["path"])
|
||
if not src.exists():
|
||
continue
|
||
dst = dst_dir / src.name
|
||
# If name collision, append a short hash suffix
|
||
if dst.exists():
|
||
stem = src.stem
|
||
ext = src.suffix
|
||
dst = dst_dir / f"{stem}_{img_id[:6]}{ext}"
|
||
try:
|
||
shutil.move(str(src), str(dst))
|
||
moved += 1
|
||
except Exception as e:
|
||
errors.append(str(e))
|
||
|
||
# Remove moved images from in-memory metadata
|
||
STATE["meta"] = [m for m in STATE["meta"] if m["id"] not in ids]
|
||
|
||
return jsonify({"moved": moved, "errors": errors,
|
||
"remaining": len(STATE["meta"])})
|
||
|
||
@flask_app.route("/api/stats")
|
||
def api_stats():
|
||
return jsonify({
|
||
"total": len(STATE["meta"]),
|
||
"buckets": len(set(m["bucket"] for m in STATE["meta"])),
|
||
})
|
||
|
||
# ─────────────────────────── HTML / JS UI ────────────────────────────────────
|
||
|
||
HTML = r"""<!DOCTYPE html>
|
||
<html lang="en">
|
||
<head>
|
||
<meta charset="UTF-8">
|
||
<meta name="viewport" content="width=device-width,initial-scale=1">
|
||
<title>Class Labeller</title>
|
||
<style>
|
||
@import url('https://fonts.googleapis.com/css2?family=IBM+Plex+Mono:wght@400;500;600&family=Epilogue:wght@800;900&display=swap');
|
||
|
||
:root{
|
||
--bg:#0a0b0e; --s1:#111318; --s2:#1a1d24; --bd:#24282f;
|
||
--text:#cdd5e0; --muted:#4e5566; --accent:#00ffb2;
|
||
--sel:#1a3d2e; --sel-bd:#00ffb2;
|
||
--red:#ff4558; --amber:#ffb347;
|
||
--radius:6px; --thumb:112px;
|
||
}
|
||
*{box-sizing:border-box;margin:0;padding:0}
|
||
html,body{height:100%;overflow:hidden}
|
||
body{background:var(--bg);color:var(--text);font-family:'IBM Plex Mono',monospace;
|
||
font-size:13px;display:grid;grid-template-rows:52px 1fr;grid-template-columns:240px 1fr}
|
||
|
||
/* ── header ── */
|
||
header{grid-column:1/-1;display:flex;align-items:center;gap:16px;
|
||
padding:0 20px;border-bottom:1px solid var(--bd);background:var(--s1);z-index:10}
|
||
header h1{font-family:'Epilogue',sans-serif;font-size:16px;font-weight:900;
|
||
color:var(--accent);letter-spacing:-.5px;white-space:nowrap}
|
||
.hstats{font-size:11px;color:var(--muted);flex:1}
|
||
.hstats b{color:var(--text)}
|
||
|
||
/* assign bar */
|
||
.assign-bar{display:flex;gap:8px;align-items:center;margin-left:auto}
|
||
#sel-count{font-size:11px;color:var(--muted);white-space:nowrap;min-width:80px;text-align:right}
|
||
.cls-wrap{position:relative}
|
||
#cls-input{background:var(--s2);border:1px solid var(--bd);color:var(--text);
|
||
padding:6px 12px;border-radius:var(--radius);font-family:inherit;
|
||
font-size:12px;width:200px;outline:none}
|
||
#cls-input:focus{border-color:var(--accent)}
|
||
#cls-drop{position:absolute;top:calc(100% + 4px);left:0;right:0;
|
||
background:var(--s2);border:1px solid var(--bd);border-radius:var(--radius);
|
||
max-height:200px;overflow-y:auto;display:none;z-index:100}
|
||
#cls-drop.open{display:block}
|
||
.cls-opt{padding:7px 12px;cursor:pointer;font-size:12px}
|
||
.cls-opt:hover{background:var(--sel);color:var(--accent)}
|
||
.btn{display:inline-flex;align-items:center;gap:6px;padding:6px 16px;
|
||
border-radius:var(--radius);font-family:inherit;font-size:12px;
|
||
cursor:pointer;border:none;font-weight:600;transition:.12s}
|
||
.btn-primary{background:var(--accent);color:#000}
|
||
.btn-primary:hover{filter:brightness(1.1)}
|
||
.btn-primary:disabled{background:var(--bd);color:var(--muted);cursor:not-allowed;filter:none}
|
||
.btn-ghost{background:transparent;border:1px solid var(--bd);color:var(--muted)}
|
||
.btn-ghost:hover{border-color:var(--text);color:var(--text)}
|
||
|
||
/* ── sidebar ── */
|
||
aside{background:var(--s1);border-right:1px solid var(--bd);
|
||
overflow-y:auto;padding:12px 8px}
|
||
.side-title{font-size:10px;letter-spacing:1px;text-transform:uppercase;
|
||
color:var(--muted);padding:4px 8px 8px}
|
||
.bkt-btn{width:100%;background:none;border:none;color:var(--text);
|
||
font-family:inherit;font-size:12px;padding:7px 10px;
|
||
border-radius:var(--radius);cursor:pointer;
|
||
display:flex;justify-content:space-between;align-items:center;transition:.1s}
|
||
.bkt-btn:hover{background:var(--s2)}
|
||
.bkt-btn.active{background:rgba(0,255,178,.08);color:var(--accent)}
|
||
.bkt-swatch{width:10px;height:10px;border-radius:2px;flex-shrink:0;margin-right:6px}
|
||
.bkt-left{display:flex;align-items:center;gap:0}
|
||
.bkt-count{font-size:10px;color:var(--muted);background:var(--bd);
|
||
border-radius:99px;padding:1px 7px}
|
||
.bkt-btn.active .bkt-count{background:rgba(0,255,178,.15);color:var(--accent)}
|
||
|
||
/* ── main ── */
|
||
main{overflow:hidden;display:flex;flex-direction:column;position:relative}
|
||
|
||
/* toolbar */
|
||
.toolbar{display:flex;align-items:center;gap:10px;padding:10px 16px;
|
||
border-bottom:1px solid var(--bd);flex-shrink:0;background:var(--s1)}
|
||
.sel-bar{display:none;align-items:center;gap:8px;
|
||
background:rgba(0,255,178,.06);border:1px solid rgba(0,255,178,.2);
|
||
border-radius:var(--radius);padding:6px 12px;font-size:12px}
|
||
.sel-bar.on{display:flex}
|
||
.filter-input{background:var(--s2);border:1px solid var(--bd);color:var(--text);
|
||
padding:5px 10px;border-radius:var(--radius);
|
||
font-family:inherit;font-size:12px;outline:none;width:180px}
|
||
.filter-input:focus{border-color:var(--accent)}
|
||
.toolbar-right{display:flex;align-items:center;gap:10px;margin-left:auto;font-size:11px;color:var(--muted)}
|
||
|
||
/* scroller */
|
||
.scroll-host{flex:1;overflow-y:auto;overflow-x:hidden;position:relative}
|
||
#spacer{width:1px} /* height set by JS */
|
||
#grid-rows{position:absolute;top:0;left:0;right:0;padding:12px 14px}
|
||
|
||
/* row */
|
||
.grid-row{display:flex;flex-wrap:wrap;gap:6px;margin-bottom:6px;contain:layout style paint}
|
||
|
||
/* card */
|
||
.card{width:var(--thumb);flex-shrink:0;border-radius:var(--radius);
|
||
overflow:hidden;cursor:pointer;border:2px solid transparent;
|
||
background:var(--s2);transition:border-color .1s,transform .1s;
|
||
position:relative;user-select:none;contain:layout style paint}
|
||
.card:hover{border-color:var(--muted);transform:translateY(-2px)}
|
||
.card.sel{border-color:var(--sel-bd);background:var(--sel)}
|
||
.card img{width:100%;height:var(--thumb);object-fit:cover;display:block;
|
||
background:var(--bd)}
|
||
.card .cname{font-size:10px;padding:3px 6px;white-space:nowrap;
|
||
overflow:hidden;text-overflow:ellipsis;color:var(--muted)}
|
||
.card.sel .cname{color:var(--accent)}
|
||
.card .chk{position:absolute;top:4px;left:4px;width:16px;height:16px;
|
||
border-radius:3px;border:1.5px solid rgba(255,255,255,.25);
|
||
background:rgba(0,0,0,.35);display:flex;align-items:center;
|
||
justify-content:center;font-size:9px;transition:.1s}
|
||
.card.sel .chk{background:var(--accent);border-color:var(--accent);color:#000}
|
||
|
||
/* bucket label row */
|
||
.bucket-label{width:100%;font-size:10px;letter-spacing:.8px;
|
||
text-transform:uppercase;color:var(--muted);
|
||
padding:10px 2px 4px;display:flex;align-items:center;gap:8px}
|
||
.bucket-label::after{content:'';flex:1;height:1px;background:var(--bd)}
|
||
|
||
/* loading overlay */
|
||
.loading{position:absolute;inset:0;background:var(--bg);
|
||
display:flex;flex-direction:column;align-items:center;
|
||
justify-content:center;gap:14px;z-index:50}
|
||
.loading h2{font-family:'Epilogue',sans-serif;font-size:18px;color:var(--accent)}
|
||
.loading p{font-size:12px;color:var(--muted)}
|
||
.bar-wrap{width:260px;height:4px;background:var(--bd);border-radius:99px;overflow:hidden}
|
||
.bar-fill{height:100%;background:var(--accent);border-radius:99px;
|
||
transition:width .3s;width:0}
|
||
.spinner{width:28px;height:28px;border:2px solid var(--bd);
|
||
border-top-color:var(--accent);border-radius:50%;
|
||
animation:spin .6s linear infinite}
|
||
@keyframes spin{to{transform:rotate(360deg)}}
|
||
|
||
/* toast */
|
||
.toast{position:fixed;bottom:24px;right:24px;background:var(--s2);
|
||
border:1px solid var(--bd);border-radius:var(--radius);
|
||
padding:10px 18px;font-size:12px;opacity:0;
|
||
transform:translateY(8px);transition:.25s;pointer-events:none;z-index:200}
|
||
.toast.show{opacity:1;transform:none}
|
||
.toast.ok{border-color:var(--accent);color:var(--accent)}
|
||
.toast.err{border-color:var(--red);color:var(--red)}
|
||
|
||
/* empty */
|
||
.empty{text-align:center;padding:60px 20px;color:var(--muted);
|
||
font-size:13px;line-height:2}
|
||
|
||
/* scrollbar */
|
||
::-webkit-scrollbar{width:4px}
|
||
::-webkit-scrollbar-track{background:transparent}
|
||
::-webkit-scrollbar-thumb{background:var(--bd);border-radius:99px}
|
||
</style>
|
||
</head>
|
||
<body>
|
||
|
||
<!-- header -->
|
||
<header>
|
||
<h1>⬡ ClassLabeller</h1>
|
||
<div class="hstats" id="hstats">Loading…</div>
|
||
<div class="assign-bar">
|
||
<span id="sel-count"></span>
|
||
<div class="cls-wrap">
|
||
<input id="cls-input" type="text" placeholder="class name…" autocomplete="off"
|
||
oninput="filterDrop()" onkeydown="clsKeydown(event)">
|
||
<div id="cls-drop"></div>
|
||
</div>
|
||
<button class="btn btn-primary" id="assign-btn" disabled onclick="doAssign()">Assign →</button>
|
||
<button class="btn btn-ghost" onclick="clearSel()">Clear</button>
|
||
</div>
|
||
</header>
|
||
|
||
<!-- sidebar -->
|
||
<aside id="sidebar">
|
||
<div class="side-title">Color groups</div>
|
||
<div id="bucket-list"></div>
|
||
</aside>
|
||
|
||
<!-- main -->
|
||
<main>
|
||
<div class="toolbar">
|
||
<input class="filter-input" type="text" id="name-filter"
|
||
placeholder="filter by name…" oninput="applyFilter()">
|
||
<div class="sel-bar" id="sel-bar">
|
||
<span id="sel-bar-txt"></span>
|
||
<button class="btn btn-ghost" style="padding:3px 10px;font-size:11px"
|
||
onclick="selectAll()">Select all visible</button>
|
||
</div>
|
||
<div class="toolbar-right">
|
||
<span id="visible-count"></span>
|
||
</div>
|
||
</div>
|
||
|
||
<!-- virtual scroll host -->
|
||
<div class="scroll-host" id="scroll-host">
|
||
<div id="spacer"></div>
|
||
<div id="grid-rows"></div>
|
||
|
||
<!-- loading state -->
|
||
<div class="loading" id="loading">
|
||
<div class="spinner"></div>
|
||
<h2>Class Labeller</h2>
|
||
<p id="load-msg">Building thumbnail cache…</p>
|
||
<div class="bar-wrap"><div class="bar-fill" id="bar-fill"></div></div>
|
||
</div>
|
||
</div>
|
||
</main>
|
||
|
||
<div class="toast" id="toast"></div>
|
||
|
||
<script>
|
||
// ─── state ────────────────────────────────────────────────────────────────
|
||
const THUMB = parseInt(getComputedStyle(document.documentElement)
|
||
.getPropertyValue('--thumb')) || 112;
|
||
const GAP = 6;
|
||
const PAD = 28;
|
||
|
||
let ALL = [];
|
||
let VISIBLE = [];
|
||
let SEL = new Set();
|
||
let CLASSES = [];
|
||
let BUCKETS = [];
|
||
let activeBucket = '__all__';
|
||
let nameFilter = '';
|
||
let cols = 1;
|
||
|
||
const CARD_W = THUMB + 4;
|
||
const CARD_H = THUMB + 22;
|
||
const ROW_H = CARD_H + GAP;
|
||
const OVERSCAN = 8;
|
||
|
||
let lastFirstRow = -1;
|
||
let lastLastRow = -1;
|
||
const loadedThumbs = new Set();
|
||
|
||
// ─── boot ─────────────────────────────────────────────────────────────────
|
||
async function boot() {
|
||
try {
|
||
setLoadMsg('Fetching metadata...', 10);
|
||
const [buckRes, clsRes] = await Promise.all([
|
||
fetch('/api/buckets'), fetch('/api/classes')
|
||
]);
|
||
BUCKETS = await buckRes.json();
|
||
CLASSES = await clsRes.json();
|
||
buildSidebar();
|
||
populateClsDrop();
|
||
|
||
setLoadMsg('Loading image list...', 30);
|
||
await loadAll();
|
||
|
||
document.getElementById('loading').style.display = 'none';
|
||
applyFilter();
|
||
updateStats();
|
||
} catch(e) {
|
||
document.getElementById('load-msg').textContent = 'Error: ' + e.message;
|
||
console.error(e);
|
||
}
|
||
}
|
||
|
||
function setLoadMsg(msg, pct) {
|
||
document.getElementById('load-msg').textContent = msg;
|
||
document.getElementById('bar-fill').style.width = pct + '%';
|
||
}
|
||
|
||
async function loadAll() {
|
||
let offset = 0;
|
||
const LIMIT = 2000;
|
||
let total = Infinity;
|
||
while (offset < total) {
|
||
const r = await fetch('/api/meta?offset=' + offset + '&limit=' + LIMIT);
|
||
const d = await r.json();
|
||
total = d.total;
|
||
ALL.push(...d.items);
|
||
offset += LIMIT;
|
||
const pct = Math.round(30 + (Math.min(offset, total) / total) * 60);
|
||
setLoadMsg('Loaded ' + Math.min(offset, total).toLocaleString() + ' / ' + total.toLocaleString() + ' images...', pct);
|
||
}
|
||
}
|
||
|
||
// ─── sidebar ──────────────────────────────────────────────────────────────
|
||
const SWATCH = {
|
||
Red:'#e05050',Orange:'#e09050',Yellow:'#d4c040',Green:'#50a050',
|
||
Cyan:'#30b0b0',Blue:'#4060d0',Purple:'#8060c0',Pink:'#c060a0',
|
||
White:'#e0e0e0',Gray:'#808080',Black:'#303030'
|
||
};
|
||
|
||
function buildSidebar() {
|
||
const list = document.getElementById('bucket-list');
|
||
const total = BUCKETS.reduce((a,b)=>a+b.count,0);
|
||
let html = `<button class="bkt-btn active" data-bkt="__all__" onclick="setBucket(this)">
|
||
<span class="bkt-left"><span class="bkt-swatch" style="background:linear-gradient(135deg,#e05050,#50a050,#4060d0)"></span>All</span>
|
||
<span class="bkt-count">${total.toLocaleString()}</span></button>`;
|
||
for (const b of BUCKETS) {
|
||
html += `<button class="bkt-btn" data-bkt="${b.name}" onclick="setBucket(this)">
|
||
<span class="bkt-left">
|
||
<span class="bkt-swatch" style="background:${SWATCH[b.name]||'#666'}"></span>
|
||
${b.name}
|
||
</span>
|
||
<span class="bkt-count">${b.count.toLocaleString()}</span>
|
||
</button>`;
|
||
}
|
||
list.innerHTML = html;
|
||
}
|
||
|
||
function setBucket(btn) {
|
||
document.querySelectorAll('.bkt-btn').forEach(b=>b.classList.remove('active'));
|
||
btn.classList.add('active');
|
||
activeBucket = btn.dataset.bkt;
|
||
SEL.clear();
|
||
applyFilter();
|
||
}
|
||
|
||
// ─── filter ───────────────────────────────────────────────────────────────
|
||
function applyFilter() {
|
||
const nf = nameFilter.toLowerCase();
|
||
VISIBLE = ALL.filter(m => {
|
||
if (activeBucket !== '__all__' && m.bucket !== activeBucket) return false;
|
||
if (nf && !m.name.toLowerCase().includes(nf)) return false;
|
||
return true;
|
||
});
|
||
document.getElementById('visible-count').textContent =
|
||
VISIBLE.length.toLocaleString() + ' images';
|
||
computeLayout();
|
||
renderViewport(true);
|
||
updateSelBar();
|
||
}
|
||
|
||
document.getElementById('name-filter').addEventListener('input', e => {
|
||
nameFilter = e.target.value;
|
||
applyFilter();
|
||
});
|
||
|
||
// ─── layout ───────────────────────────────────────────────────────────────
|
||
function computeLayout() {
|
||
const w = document.getElementById('scroll-host').clientWidth - PAD;
|
||
cols = Math.max(1, Math.floor((w + GAP) / (CARD_W + GAP)));
|
||
const rows = Math.ceil(VISIBLE.length / cols);
|
||
const totalH = rows * ROW_H;
|
||
document.getElementById('spacer').style.height = totalH + 'px';
|
||
}
|
||
|
||
window.addEventListener('resize', () => { computeLayout(); renderViewport(true); });
|
||
|
||
// ─── virtual scroll ───────────────────────────────────────────────────────
|
||
const scrollHost = document.getElementById('scroll-host');
|
||
let rafPending = false;
|
||
|
||
//scrollHost.addEventListener('scroll', () => {
|
||
// if (!rafPending) {
|
||
// rafPending = true;
|
||
// requestAnimationFrame(() => { renderViewport(true); rafPending = false; });
|
||
// }
|
||
//});
|
||
|
||
scrollHost.addEventListener('scroll', () => {
|
||
if (!rafPending) {
|
||
rafPending = true;
|
||
requestAnimationFrame(() => { renderViewport(true); rafPending = false; });
|
||
}
|
||
}, { passive: true });
|
||
|
||
// Image loading queue — limits concurrent fetches
|
||
const IMG_CONCURRENCY = 12;
|
||
let loadingNow = 0;
|
||
const loadQueue = [];
|
||
|
||
//function enqueueImg(img, src) {
|
||
// const doLoad = () => {
|
||
// loadingNow++;
|
||
// img.src = src;
|
||
// img.onload = img.onerror = () => {
|
||
// loadingNow--;
|
||
// if (loadQueue.length) loadQueue.shift()();
|
||
// };
|
||
// };
|
||
// if (loadingNow < IMG_CONCURRENCY) doLoad();
|
||
// else loadQueue.push(doLoad);
|
||
//}
|
||
|
||
function enqueueImg(img, src) {
|
||
const id = src.split('/').pop();
|
||
if (loadedThumbs.has(id)) {
|
||
img.src = src;
|
||
return;
|
||
}
|
||
const doLoad = () => {
|
||
loadingNow++;
|
||
img.src = src;
|
||
img.onload = img.onerror = () => {
|
||
loadedThumbs.add(id);
|
||
loadingNow--;
|
||
if (loadQueue.length) loadQueue.shift()();
|
||
};
|
||
};
|
||
if (loadingNow < IMG_CONCURRENCY) doLoad();
|
||
else loadQueue.push(doLoad);
|
||
}
|
||
|
||
//function renderViewport() {
|
||
// const host = scrollHost;
|
||
// const scrollY = host.scrollTop;
|
||
// const vpH = host.clientHeight;
|
||
// const firstRow = Math.max(0, Math.floor(scrollY / ROW_H) - OVERSCAN);
|
||
// const lastRow = Math.min(
|
||
// Math.ceil(VISIBLE.length / cols),
|
||
// Math.ceil((scrollY + vpH) / ROW_H) + OVERSCAN
|
||
// );
|
||
//
|
||
// const container = document.getElementById('grid-rows');
|
||
// container.style.top = (firstRow * ROW_H) + 'px';
|
||
//
|
||
// // Build HTML for visible rows
|
||
// let html = '';
|
||
// let prevBucket = null;
|
||
// for (let row = firstRow; row < lastRow; row++) {
|
||
// const startIdx = row * cols;
|
||
// if (startIdx >= VISIBLE.length) break;
|
||
//
|
||
// // Bucket separator
|
||
// const firstInRow = VISIBLE[startIdx];
|
||
// if (firstInRow && firstInRow.bucket !== prevBucket && activeBucket === '__all__') {
|
||
// html += `<div class="bucket-label">${firstInRow.bucket}</div>`;
|
||
// prevBucket = firstInRow.bucket;
|
||
// }
|
||
//
|
||
// html += '<div class="grid-row">';
|
||
// for (let c = 0; c < cols; c++) {
|
||
// const idx = startIdx + c;
|
||
// if (idx >= VISIBLE.length) break;
|
||
// const m = VISIBLE[idx];
|
||
// const sel = SEL.has(m.id) ? ' sel' : '';
|
||
// html += `<div class="card${sel}" id="c-${m.id}"
|
||
// onclick="cardClick(event,'${m.id}')"
|
||
// data-id="${m.id}">
|
||
// <div class="chk">${SEL.has(m.id)?'✓':''}</div>
|
||
// <img data-src="/thumb/${m.id}" alt="${m.name}">
|
||
// <div class="cname" title="${m.name}">${m.name}</div>
|
||
// </div>`;
|
||
// }
|
||
// html += '</div>';
|
||
// }
|
||
//
|
||
// if (!html && VISIBLE.length === 0) {
|
||
// html = `<div class="empty">No images in this group.<br>
|
||
// Try a different filter or colour group.</div>`;
|
||
// }
|
||
//
|
||
// container.innerHTML = html;
|
||
//
|
||
// // Lazy-load images via queue
|
||
// container.querySelectorAll('img[data-src]').forEach(img => {
|
||
// enqueueImg(img, img.dataset.src);
|
||
// img.removeAttribute('data-src');
|
||
// });
|
||
//}
|
||
|
||
function renderViewport(force = false) {
|
||
const host = scrollHost;
|
||
const scrollY = host.scrollTop;
|
||
const vpH = host.clientHeight;
|
||
const firstRow = Math.max(0, Math.floor(scrollY / ROW_H) - OVERSCAN);
|
||
const lastRow = Math.min(
|
||
Math.ceil(VISIBLE.length / cols),
|
||
Math.ceil((scrollY + vpH) / ROW_H) + OVERSCAN
|
||
);
|
||
|
||
if (!force && firstRow === lastFirstRow && lastRow === lastLastRow) return;
|
||
lastFirstRow = firstRow;
|
||
lastLastRow = lastRow;
|
||
|
||
const container = document.getElementById('grid-rows');
|
||
container.style.top = (firstRow * ROW_H) + 'px';
|
||
|
||
let html = '';
|
||
let prevBucket = null;
|
||
for (let row = firstRow; row < lastRow; row++) {
|
||
const startIdx = row * cols;
|
||
if (startIdx >= VISIBLE.length) break;
|
||
|
||
const firstInRow = VISIBLE[startIdx];
|
||
if (firstInRow && firstInRow.bucket !== prevBucket && activeBucket === '__all__') {
|
||
html += `<div class="bucket-label">${firstInRow.bucket}</div>`;
|
||
prevBucket = firstInRow.bucket;
|
||
}
|
||
|
||
html += '<div class="grid-row">';
|
||
for (let c = 0; c < cols; c++) {
|
||
const idx = startIdx + c;
|
||
if (idx >= VISIBLE.length) break;
|
||
const m = VISIBLE[idx];
|
||
const sel = SEL.has(m.id) ? ' sel' : '';
|
||
const url = `/thumb/${m.id}`;
|
||
const img = loadedThumbs.has(m.id)
|
||
? `<img src="${url}" alt="${m.name}">`
|
||
: `<img data-src="${url}" alt="${m.name}">`;
|
||
html += `<div class="card${sel}" id="c-${m.id}"
|
||
onclick="cardClick(event,'${m.id}')"
|
||
data-id="${m.id}">
|
||
<div class="chk">${SEL.has(m.id)?'✓':''}</div>
|
||
${img}
|
||
<div class="cname" title="${m.name}">${m.name}</div>
|
||
</div>`;
|
||
}
|
||
html += '</div>';
|
||
}
|
||
|
||
if (!html && VISIBLE.length === 0) {
|
||
html = `<div class="empty">No images in this group.<br>
|
||
Try a different filter or colour group.</div>`;
|
||
}
|
||
|
||
container.innerHTML = html;
|
||
|
||
container.querySelectorAll('img[data-src]').forEach(img => {
|
||
enqueueImg(img, img.dataset.src);
|
||
img.removeAttribute('data-src');
|
||
});
|
||
}
|
||
|
||
// ─── selection ────────────────────────────────────────────────────────────
|
||
let lastClickIdx = null;
|
||
|
||
function cardClick(evt, id) {
|
||
const idx = VISIBLE.findIndex(m => m.id === id);
|
||
const card = document.getElementById('c-' + id);
|
||
|
||
if (evt.shiftKey && lastClickIdx !== null) {
|
||
// Range select
|
||
const lo = Math.min(lastClickIdx, idx);
|
||
const hi = Math.max(lastClickIdx, idx);
|
||
for (let i = lo; i <= hi; i++) SEL.add(VISIBLE[i].id);
|
||
} else if (evt.ctrlKey || evt.metaKey) {
|
||
SEL.has(id) ? SEL.delete(id) : SEL.add(id);
|
||
} else {
|
||
SEL.has(id) ? SEL.delete(id) : SEL.add(id);
|
||
}
|
||
lastClickIdx = idx;
|
||
updateSelBar();
|
||
renderViewport(true);
|
||
}
|
||
|
||
function clearSel() { SEL.clear(); updateSelBar(); renderViewport(true); }
|
||
|
||
function selectAll() {
|
||
VISIBLE.forEach(m => SEL.add(m.id));
|
||
updateSelBar();
|
||
renderViewport(true);
|
||
}
|
||
|
||
function updateSelBar() {
|
||
const n = SEL.size;
|
||
const bar = document.getElementById('sel-bar');
|
||
bar.classList.toggle('on', n > 0);
|
||
document.getElementById('sel-bar-txt').textContent = `${n.toLocaleString()} selected`;
|
||
document.getElementById('sel-count').textContent = n > 0 ? `${n} selected` : '';
|
||
document.getElementById('assign-btn').disabled = n === 0;
|
||
}
|
||
|
||
// ─── class dropdown ───────────────────────────────────────────────────────
|
||
function populateClsDrop() {
|
||
const drop = document.getElementById('cls-drop');
|
||
drop.innerHTML = CLASSES.map(c =>
|
||
`<div class="cls-opt" onclick="pickCls('${c}')">${c}</div>`
|
||
).join('');
|
||
}
|
||
|
||
function filterDrop() {
|
||
const val = document.getElementById('cls-input').value.toLowerCase();
|
||
const drop = document.getElementById('cls-drop');
|
||
const opts = CLASSES.filter(c => c.toLowerCase().includes(val));
|
||
drop.innerHTML = opts.map(c =>
|
||
`<div class="cls-opt" onclick="pickCls('${c}')">${c}</div>`
|
||
).join('');
|
||
drop.classList.toggle('open', opts.length > 0 && val.length > 0);
|
||
document.getElementById('assign-btn').disabled = SEL.size === 0;
|
||
}
|
||
|
||
function pickCls(name) {
|
||
document.getElementById('cls-input').value = name;
|
||
document.getElementById('cls-drop').classList.remove('open');
|
||
}
|
||
|
||
function clsKeydown(e) {
|
||
if (e.key === 'Enter') doAssign();
|
||
if (e.key === 'Escape') document.getElementById('cls-drop').classList.remove('open');
|
||
}
|
||
|
||
document.addEventListener('click', e => {
|
||
if (!(e.target instanceof Element) || !e.target.closest('.cls-wrap'))
|
||
document.getElementById('cls-drop').classList.remove('open');
|
||
});
|
||
|
||
// ─── assign ───────────────────────────────────────────────────────────────
|
||
async function doAssign() {
|
||
const cls = document.getElementById('cls-input').value.trim();
|
||
if (!cls) { toast('Enter a class name', 'err'); return; }
|
||
if (SEL.size === 0) { toast('Nothing selected', 'err'); return; }
|
||
|
||
const ids = [...SEL];
|
||
const res = await fetch('/api/assign', {
|
||
method: 'POST',
|
||
headers: {'Content-Type':'application/json'},
|
||
body: JSON.stringify({ids, class_name: cls})
|
||
});
|
||
const d = await res.json();
|
||
|
||
// Remove from local state
|
||
ALL = ALL.filter(m => !SEL.has(m.id));
|
||
SEL.clear();
|
||
document.getElementById('cls-input').value = '';
|
||
|
||
// Add class to dropdown if new
|
||
if (!CLASSES.includes(cls)) {
|
||
CLASSES.push(cls);
|
||
CLASSES.sort();
|
||
populateClsDrop();
|
||
}
|
||
|
||
applyFilter();
|
||
updateStats();
|
||
toast(`✓ Moved ${d.moved} images to ${cls}/`, 'ok');
|
||
}
|
||
|
||
// ─── stats ────────────────────────────────────────────────────────────────
|
||
function updateStats() {
|
||
document.getElementById('hstats').innerHTML =
|
||
`<b>${ALL.length.toLocaleString()}</b> remaining · <b>${CLASSES.length}</b> classes`;
|
||
}
|
||
|
||
// ─── toast ────────────────────────────────────────────────────────────────
|
||
let toastTimer = null;
|
||
function toast(msg, type) {
|
||
if (!type) type = 'ok';
|
||
const el = document.getElementById('toast');
|
||
el.textContent = msg;
|
||
el.className = `toast ${type} show`;
|
||
if (toastTimer) clearTimeout(toastTimer);
|
||
toastTimer = setTimeout(() => el.classList.remove('show'), 3000);
|
||
}
|
||
|
||
// ─── keyboard shortcuts ───────────────────────────────────────────────────
|
||
document.addEventListener('keydown', function(e) {
|
||
const tag = document.activeElement && document.activeElement.id;
|
||
if (tag === 'cls-input' || tag === 'name-filter') return;
|
||
if (e.key === 'a' && (e.ctrlKey || e.metaKey)) { e.preventDefault(); selectAll(); }
|
||
if (e.key === 'Escape') clearSel();
|
||
if (e.key === 'Enter' && SEL.size > 0) doAssign();
|
||
});
|
||
|
||
boot();
|
||
</script>
|
||
</body>
|
||
</html>"""
|
||
|
||
# ─────────────────────────── main ────────────────────────────────────────────
|
||
|
||
def main():
|
||
args = ARGS
|
||
|
||
source_dir = Path(args.source)
|
||
cache_dir = Path(args.cache_dir)
|
||
data_dir = Path(args.data_dir)
|
||
|
||
if not source_dir.exists():
|
||
print(f"[error] Source directory not found: {source_dir}")
|
||
print(" Create it and add your unlabelled crops, then re-run.")
|
||
return
|
||
|
||
data_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
print(f"""
|
||
╔══════════════════════════════════════════╗
|
||
║ Class Labeller ║
|
||
╠══════════════════════════════════════════╣
|
||
Source : {source_dir}
|
||
Data dir : {data_dir}
|
||
Cache : {cache_dir}
|
||
Workers : {args.workers}
|
||
Thumb : {args.thumb_size}px
|
||
╚══════════════════════════════════════════╝
|
||
""")
|
||
|
||
# Build / update thumbnail cache (blocking, before server starts)
|
||
meta = build_metadata(source_dir, cache_dir, args.thumb_size, args.workers)
|
||
|
||
STATE["meta"] = meta
|
||
|
||
import threading
|
||
import webbrowser
|
||
url = f"http://localhost:{args.port}"
|
||
|
||
def _open():
|
||
time.sleep(1.0)
|
||
webbrowser.open(url)
|
||
threading.Thread(target=_open, daemon=True).start()
|
||
|
||
print(f"\n[server] Starting at {url}")
|
||
print(" Keyboard shortcuts:")
|
||
print(" Ctrl+A → select all visible")
|
||
print(" Shift+click → range select")
|
||
print(" Ctrl+click → toggle single")
|
||
print(" Enter → assign selected")
|
||
print(" Esc → clear selection")
|
||
print("\n Press Ctrl+C to stop.\n")
|
||
|
||
flask_app.run(host="0.0.0.0", port=args.port,
|
||
threaded=True, use_reloader=False)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|