feat: sync live-count, exemplar annotation modules, and update .gitignore
This commit is contained in:
1 parent
b6624eeff9
commit
ac95674c07
39 files changed
+3979
-444
No files matched your search
@@ -259,125 +259,3 @@ def _reset_reviewed(batch_id: int) -> None:
|
||||
"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
|
||||
Reference in new issue
Block a user