asus
5c7c122105
This commit includes major additions and updates to the frontend and backend architectures, introducing new dataset management, live counting features, batch processing, and triage logic. Includes new UI pages, components, and API routes.
384 lines
17 KiB
Python
384 lines
17 KiB
Python
"""The auto-annotation job: SAM3 over every frame of a batch (REQ-030…034).
|
|
|
|
Detection itself is `labeling.label_image`, unchanged — one `set_image` per
|
|
frame with the class prompts looped over that cached state, then greedy IoU NMS
|
|
across prompts. This module's job is only to turn its output into rows and to
|
|
keep the user's own corrections out of the way.
|
|
"""
|
|
|
|
import os
|
|
from typing import List, Optional
|
|
|
|
from PIL import Image
|
|
from backend import batches, db, jobs, labeling, projects, review
|
|
from backend.batches import BatchError
|
|
|
|
DEFAULT_THRESHOLD = 0.35
|
|
DEFAULT_IOU = 0.0
|
|
|
|
|
|
|
|
|
|
def start(batch_id: int, threshold: float = DEFAULT_THRESHOLD,
|
|
iou_threshold: float = DEFAULT_IOU, min_box_frac: float = 0.0,
|
|
resume: bool = False, append: bool = False, engine: str = "base_model",
|
|
engines: Optional[List[str]] = None,
|
|
class_ids: Optional[List[int]] = None,
|
|
engine_classes: Optional[dict[str, List[str]]] = None,
|
|
custom_model_path: Optional[str] = None,
|
|
target_class_names: Optional[List[str]] = None) -> dict:
|
|
batch = batches.get(batch_id)
|
|
if batch is None:
|
|
raise batches.BatchError("No such batch")
|
|
if batch["frame_count"] == 0:
|
|
raise batches.BatchError("This batch has no frames yet")
|
|
|
|
active_engines = engines if (engines and len(engines) > 0) else [engine]
|
|
|
|
job = jobs.create(
|
|
"autolabel",
|
|
params={"batch_id": batch_id, "threshold": threshold,
|
|
"iou_threshold": iou_threshold, "min_box_frac": min_box_frac,
|
|
"resume": resume, "append": append, "engine": active_engines[0], "engines": active_engines,
|
|
"class_ids": class_ids, "engine_classes": engine_classes,
|
|
"custom_model_path": custom_model_path, "target_class_names": target_class_names},
|
|
project_id=batch["project_id"],
|
|
batch_id=batch_id,
|
|
message=f"{batch['date_label']}/{batch['batch_label']} ({'+'.join(e.upper() for e in active_engines)})",
|
|
)
|
|
return job.to_dict()
|
|
|
|
|
|
def _geometries(detection, width: int, height: int, label_type: str) -> List[dict]:
|
|
if label_type == "bbox":
|
|
x0, y0, x1, y1 = detection.box
|
|
return [review.bbox(x0 / width, y0 / height, x1 / width, y1 / height)]
|
|
|
|
shapes = []
|
|
for points in review.mask_to_polygons(detection.mask):
|
|
if len(points) >= 3:
|
|
shapes.append(review.polygon([(x / width, y / height) for x, y in points]))
|
|
return shapes
|
|
|
|
|
|
@jobs.handler("autolabel")
|
|
def _run_autolabel(job) -> None:
|
|
batch = batches.get(job.params["batch_id"])
|
|
if batch is None:
|
|
raise batches.BatchError("The batch disappeared before labeling started")
|
|
project = projects.get(batch["project_id"])
|
|
|
|
selected_engine = job.params.get("engine", "base_model")
|
|
|
|
frames = batches.frames(batch["id"])
|
|
batches.set_status(batch["id"], "labeling")
|
|
job.progress(0, len(frames))
|
|
|
|
skip = review.frames_with_auto(batch["id"]) if job.params.get("resume") else set()
|
|
if skip:
|
|
job.log(f"Resuming: skipping {len(skip)} frame(s) that already have automatic shapes")
|
|
|
|
directory = batches.frames_dir(batch["project_slug"], batch["id"])
|
|
written = 0
|
|
attempted = 0
|
|
failures = []
|
|
|
|
conf = job.params.get("threshold", DEFAULT_THRESHOLD)
|
|
iou_thresh = job.params.get("iou_threshold", DEFAULT_IOU)
|
|
|
|
yolo_model = None
|
|
sam3_target_classes = []
|
|
|
|
custom_path = job.params.get("custom_model_path")
|
|
target_class_names = job.params.get("target_class_names")
|
|
engine_classes = job.params.get("engine_classes")
|
|
if not target_class_names and isinstance(engine_classes, dict):
|
|
c_names = engine_classes.get(selected_engine) or engine_classes.get("sam3") or []
|
|
if isinstance(c_names, list) and len(c_names) > 0:
|
|
target_class_names = c_names
|
|
|
|
if selected_engine == "sam3" and not custom_path:
|
|
allowed_classes_set = {c.strip().lower() for c in target_class_names} if target_class_names else None
|
|
if allowed_classes_set:
|
|
# Add any new target class names that aren't in project classes yet
|
|
existing_names = {c["name"].strip().lower() for c in project["classes"]}
|
|
for name in target_class_names:
|
|
if name.strip().lower() not in existing_names:
|
|
try:
|
|
updated_proj = projects.add_class(project["id"], name=name.strip(), prompt=name.strip())
|
|
project["classes"] = updated_proj["classes"]
|
|
except Exception as exc:
|
|
job.log(f"Warning adding class '{name}': {exc}")
|
|
sam3_target_classes = [c for c in project["classes"] if c["name"].strip().lower() in allowed_classes_set or c["prompt"].strip().lower() in allowed_classes_set]
|
|
else:
|
|
sam3_target_classes = [c for c in project["classes"]]
|
|
|
|
prompts = [c["prompt"] for c in sam3_target_classes]
|
|
if prompts:
|
|
from backend.sam3_engine import engine_is_loaded, get_engine
|
|
if not engine_is_loaded():
|
|
job.log("Loading SAM3 (the first run downloads ~3.4 GB from HuggingFace)…")
|
|
engine = get_engine()
|
|
job.log(f"SAM3 ready on {engine.device}; prompts: {', '.join(prompts)}")
|
|
else:
|
|
job.log("SAM3 selected but 0 prompts match project classes.")
|
|
else:
|
|
from ultralytics import YOLO
|
|
if custom_path and os.path.isfile(custom_path):
|
|
m_path = custom_path
|
|
job.log(f"Loading Custom Model: {os.path.basename(m_path)}...")
|
|
else:
|
|
m_path = projects.training_start_point(project)
|
|
with db.cursor() as cur:
|
|
cur.execute("SELECT weights_path FROM model_versions WHERE project_id = ? ORDER BY version DESC LIMIT 1", (project["id"],))
|
|
row = cur.fetchone()
|
|
if row and os.path.isfile(row[0]):
|
|
m_path = row[0]
|
|
job.log(f"Loading Base Model: {os.path.basename(m_path)}...")
|
|
|
|
yolo_model = YOLO(m_path)
|
|
|
|
name_to_class_id = {item["name"].strip().lower(): item["class_id"] for item in project["classes"]}
|
|
allowed_classes_set = {c.strip().lower() for c in target_class_names} if target_class_names else None
|
|
|
|
job.log(f"Starting auto-labeling with {selected_engine}...")
|
|
|
|
for index, frame in enumerate(frames):
|
|
if job.cancelled:
|
|
job.log(f"Cancelled after {index} frame(s)")
|
|
break
|
|
if frame["id"] in skip:
|
|
job.progress(index + 1, len(frames))
|
|
continue
|
|
attempted += 1
|
|
|
|
try:
|
|
frame_file = os.path.join(directory, frame["filename"])
|
|
fw = max(1, frame.get("width") or 1)
|
|
fh = max(1, frame.get("height") or 1)
|
|
all_raw_detections = []
|
|
|
|
if yolo_model is not None:
|
|
results = yolo_model.predict(frame_file, conf=conf, verbose=False)
|
|
if results and len(results) > 0:
|
|
model_names = results[0].names
|
|
for box in results[0].boxes:
|
|
cls_idx = int(box.cls[0].item())
|
|
raw_cls_name = str(model_names.get(cls_idx, cls_idx)).strip().lower()
|
|
|
|
target_class_id = name_to_class_id.get(raw_cls_name)
|
|
if target_class_id is None:
|
|
for item in project["classes"]:
|
|
if item["class_id"] == cls_idx:
|
|
target_class_id = item["class_id"]
|
|
break
|
|
if target_class_id is None and 0 <= cls_idx < len(project["classes"]):
|
|
target_class_id = project["classes"][cls_idx]["class_id"]
|
|
|
|
if target_class_id is None:
|
|
continue
|
|
|
|
target_cls_obj = next((c for c in project["classes"] if c["class_id"] == target_class_id), None)
|
|
proj_cls_name = target_cls_obj["name"].strip().lower() if target_cls_obj else ""
|
|
|
|
if allowed_classes_set is not None:
|
|
if (raw_cls_name not in allowed_classes_set and
|
|
proj_cls_name not in allowed_classes_set and
|
|
str(target_class_id) not in allowed_classes_set):
|
|
continue
|
|
|
|
score = float(box.conf[0].item())
|
|
xyxyn = box.xyxyn[0].tolist()
|
|
all_raw_detections.append(labeling.Detection(
|
|
class_id=target_class_id,
|
|
class_name=proj_cls_name or raw_cls_name,
|
|
box=[xyxyn[0]*fw, xyxyn[1]*fh, xyxyn[2]*fw, xyxyn[3]*fh],
|
|
score=score,
|
|
mask=None
|
|
))
|
|
|
|
if selected_engine == "sam3" and sam3_target_classes:
|
|
prompts = [(c.get("prompt") or c["name"]).strip() for c in sam3_target_classes]
|
|
res = labeling.label_image(
|
|
frame_file, frame["filename"], prompts, conf,
|
|
iou_threshold=iou_thresh, min_box_frac=job.params.get("min_box_frac", 0.0)
|
|
)
|
|
if not res.error and res.detections:
|
|
for det in res.detections:
|
|
if 0 <= det.class_id < len(sam3_target_classes):
|
|
real_cls = sam3_target_classes[det.class_id]
|
|
det.class_id = real_cls["class_id"]
|
|
det.class_name = real_cls["name"]
|
|
all_raw_detections.append(det)
|
|
elif res.error:
|
|
job.log(f"[SAM3 ERROR] {frame['filename']}: {res.error}")
|
|
|
|
kept = labeling.deduplicate(all_raw_detections, iou_threshold=iou_thresh)
|
|
items = []
|
|
for det in kept:
|
|
if project["label_type"] == "bbox" or det.mask is None:
|
|
geom = review.bbox(det.box[0]/fw, det.box[1]/fh, det.box[2]/fw, det.box[3]/fh)
|
|
items.append({"class_id": det.class_id, "geometry": geom, "score": det.score})
|
|
else:
|
|
for geometry in _geometries(det, fw, fh, project["label_type"]):
|
|
items.append({"class_id": det.class_id, "geometry": geometry, "score": det.score})
|
|
|
|
if job.params.get("append"):
|
|
review.append_auto(frame["id"], items)
|
|
else:
|
|
review.replace_auto(frame["id"], items)
|
|
written += len(items)
|
|
job.progress(index + 1, len(frames), f"{frame['filename']}: {len(items)} shape(s)")
|
|
except Exception as exc:
|
|
failures.append(str(exc))
|
|
job.log(f"[ERROR] {frame['filename']}: {exc}")
|
|
job.progress(index + 1, len(frames))
|
|
|
|
|
|
|
|
# "Every frame failed" is not a finished job with no findings — it is a
|
|
# broken run, and reporting `done` for it would be the system lying about
|
|
# its own state. An empty frame is fine (REQ-033); an errored one is not.
|
|
if attempted and len(failures) == attempted:
|
|
batches.set_status(batch["id"], "failed")
|
|
raise BatchError(f"All {attempted} frame(s) failed. First error: {failures[0]}")
|
|
|
|
batches.set_status(batch["id"], "reviewing")
|
|
_reset_reviewed(batch["id"])
|
|
if failures:
|
|
job.log(f"{len(failures)} of {attempted} frame(s) failed — see the errors above")
|
|
job.log(f"Wrote {written} shape(s) across {attempted - len(failures)} frame(s)")
|
|
|
|
|
|
def _reset_reviewed(batch_id: int) -> None:
|
|
"""Approvals were given against the previous labels, so a re-run puts those
|
|
frames back in the queue. Manual shapes stay; the sign-off does not."""
|
|
with db.cursor() as cur:
|
|
cur.execute(
|
|
"UPDATE frames SET review_status = 'pending' WHERE batch_id = ? "
|
|
"AND review_status = 'approved'",
|
|
(batch_id,),
|
|
)
|
|
|
|
|
|
def preview_frame(
|
|
batch_id: int,
|
|
frame_id: int,
|
|
engine: str,
|
|
threshold: float = DEFAULT_THRESHOLD,
|
|
iou_threshold: float = DEFAULT_IOU,
|
|
min_box_frac: float = 0.0,
|
|
target_class_names: Optional[List[str]] = None,
|
|
custom_model_path: Optional[str] = None
|
|
) -> List[dict]:
|
|
batch = batches.get(batch_id)
|
|
if not batch:
|
|
raise ValueError("No such batch")
|
|
project = projects.get(batch["project_id"])
|
|
|
|
frame = next((f for f in batches.frames(batch_id) if f["id"] == frame_id), None)
|
|
if not frame:
|
|
raise ValueError("Frame not found")
|
|
|
|
directory = batches.frames_dir(batch["project_slug"], batch_id)
|
|
frame_file = os.path.join(directory, frame["filename"])
|
|
fw = max(1, frame.get("width") or 1)
|
|
fh = max(1, frame.get("height") or 1)
|
|
|
|
yolo_model = None
|
|
sam3_target_classes = []
|
|
|
|
if engine == "sam3" and not custom_model_path:
|
|
allowed_classes_set = {c.strip().lower() for c in target_class_names} if target_class_names else None
|
|
if allowed_classes_set:
|
|
sam3_target_classes = [c for c in project["classes"] if c["name"].strip().lower() in allowed_classes_set or c["prompt"].strip().lower() in allowed_classes_set]
|
|
else:
|
|
sam3_target_classes = [c for c in project["classes"]]
|
|
|
|
prompts = [c["prompt"] for c in sam3_target_classes]
|
|
if prompts:
|
|
from backend.sam3_engine import get_engine
|
|
get_engine()
|
|
else:
|
|
from ultralytics import YOLO
|
|
if custom_model_path and os.path.isfile(custom_model_path):
|
|
m_path = custom_model_path
|
|
else:
|
|
m_path = projects.training_start_point(project)
|
|
with db.cursor() as cur:
|
|
cur.execute("SELECT weights_path FROM model_versions WHERE project_id = ? ORDER BY version DESC LIMIT 1", (project["id"],))
|
|
row = cur.fetchone()
|
|
if row and os.path.isfile(row[0]):
|
|
m_path = row[0]
|
|
yolo_model = YOLO(m_path)
|
|
|
|
name_to_class_id = {item["name"].strip().lower(): item["class_id"] for item in project["classes"]}
|
|
allowed_classes_set = {c.strip().lower() for c in target_class_names} if target_class_names else None
|
|
|
|
all_raw_detections = []
|
|
|
|
if yolo_model is not None:
|
|
results = yolo_model.predict(frame_file, conf=threshold, verbose=False)
|
|
if results and len(results) > 0:
|
|
model_names = results[0].names
|
|
for box in results[0].boxes:
|
|
cls_idx = int(box.cls[0].item())
|
|
raw_cls_name = str(model_names.get(cls_idx, cls_idx)).strip().lower()
|
|
|
|
target_class_id = name_to_class_id.get(raw_cls_name)
|
|
if target_class_id is None:
|
|
for item in project["classes"]:
|
|
if item["class_id"] == cls_idx:
|
|
target_class_id = item["class_id"]
|
|
break
|
|
if target_class_id is None and 0 <= cls_idx < len(project["classes"]):
|
|
target_class_id = project["classes"][cls_idx]["class_id"]
|
|
|
|
if target_class_id is None:
|
|
continue
|
|
|
|
target_cls_obj = next((c for c in project["classes"] if c["class_id"] == target_class_id), None)
|
|
proj_cls_name = target_cls_obj["name"].strip().lower() if target_cls_obj else ""
|
|
|
|
if allowed_classes_set is not None:
|
|
if (raw_cls_name not in allowed_classes_set and
|
|
proj_cls_name not in allowed_classes_set and
|
|
str(target_class_id) not in allowed_classes_set):
|
|
continue
|
|
|
|
score = float(box.conf[0].item())
|
|
xyxyn = box.xyxyn[0].tolist()
|
|
all_raw_detections.append(labeling.Detection(
|
|
class_id=target_class_id,
|
|
class_name=proj_cls_name or raw_cls_name,
|
|
box=[xyxyn[0]*fw, xyxyn[1]*fh, xyxyn[2]*fw, xyxyn[3]*fh],
|
|
score=score,
|
|
mask=None
|
|
))
|
|
|
|
if engine == "sam3" and sam3_target_classes:
|
|
prompts = [(c.get("prompt") or c["name"]).strip() for c in sam3_target_classes]
|
|
res = labeling.label_image(
|
|
frame_file, frame["filename"], prompts, threshold,
|
|
iou_threshold=iou_threshold, min_box_frac=min_box_frac
|
|
)
|
|
if not res.error and res.detections:
|
|
for det in res.detections:
|
|
if 0 <= det.class_id < len(sam3_target_classes):
|
|
real_cls = sam3_target_classes[det.class_id]
|
|
det.class_id = real_cls["class_id"]
|
|
det.class_name = real_cls["name"]
|
|
all_raw_detections.append(det)
|
|
|
|
kept = labeling.deduplicate(all_raw_detections, iou_threshold=iou_threshold)
|
|
items = []
|
|
for det in kept:
|
|
if project["label_type"] == "bbox" or det.mask is None:
|
|
geom = review.bbox(det.box[0]/fw, det.box[1]/fh, det.box[2]/fw, det.box[3]/fh)
|
|
items.append({"class_id": det.class_id, "geometry": geom, "score": det.score})
|
|
else:
|
|
for geometry in _geometries(det, fw, fh, project["label_type"]):
|
|
items.append({"class_id": det.class_id, "geometry": geometry, "score": det.score})
|
|
|
|
return items
|