# ui/workflow.py from datetime import datetime import streamlit as st from PIL import Image, ImageOps from streamlit_drawable_canvas import st_canvas from core.drawing import ( parse_boxes, # from canvas JSON -> list[xyxy] boxes_to_canvas_json, # (kept for compatibility; not used here) rects_only_json, get_rect_conf, inject_index_labels, seed_canvas_from_boxes, # list[xyxy] -> canvas JSON ) from core.state import set_defaults_from_preds, sync_samples_with_state from core.exports import export_session # not used here but retained from core.detect_infer import DetConfig import torch # ---------------- helpers ---------------- def _auto_canvas_width(): try: from streamlit_js_eval import get_page_info info = get_page_info() or {} w = int(info.get("clientWidth") or info.get("windowWidth") or 0) return 1200 if w >= 1800 else (1000 if w >= 1350 else 820) except Exception: return 820 def _det_cfg_from_state(): c = st.session_state.get("det_cfg", {}) return DetConfig( conf=float(c.get("conf", 0.25)), iou=float(c.get("iou", 0.50)), imgsz=int(c.get("imgsz", 640)), max_det=int(c.get("max_det", 300)), half=bool(c.get("half", True)), ) def _iou_xyxy(a, b, eps=1e-6): ax1, ay1, ax2, ay2 = map(float, a) bx1, by1, bx2, by2 = map(float, b) ix1, iy1 = max(ax1, bx1), max(ay1, by1) ix2, iy2 = min(ax2, bx2), min(ay2, by2) iw, ih = max(0.0, ix2 - ix1), max(0.0, iy2 - iy1) inter = iw * ih aa = max(0.0, (ax2 - ax1)) * max(0.0, (ay2 - ay1)) ba = max(0.0, (bx2 - bx1)) * max(0.0, (by2 - by1)) union = aa + ba - inter + eps return inter / union def _match_det_info_to_live_boxes(item, live_boxes, iou_thresh=0.90): """Align per-box detector info (conf/cls/name) to current live boxes.""" base_boxes = item.get("boxes") or [] info = item.get("detector_info") or [] out = [] for b in live_boxes: best_iou, best_idx = 0.0, -1 for j, bb in enumerate(base_boxes): i = _iou_xyxy(b, bb) if i > best_iou: best_iou, best_idx = i, j out.append(info[best_idx] if (best_iou >= iou_thresh and 0 <= best_idx < len(info)) else None) return out def _reindex_det_info(old_boxes, old_info, new_boxes, iou_thresh=0.90): """Rebuild detector_info aligned to new_boxes by matching from (old_boxes, old_info).""" info_new = [] for b in new_boxes: best_iou, best_idx = 0.0, -1 for j, ob in enumerate(old_boxes): i = _iou_xyxy(b, ob) if i > best_iou: best_iou, best_idx = i, j info_new.append(old_info[best_idx] if (best_iou >= iou_thresh and 0 <= best_idx < len(old_info)) else None) return info_new def _run_detection(item): det = st.session_state.get("detector") if det is None: return [] cfg = _det_cfg_from_state() return det.predict_one(item["pil"], cfg) or [] # list of dicts {'bbox':[x1,y1,x2,y2], 'score','cls','name',...} def _parse_canvas_into_boxes(item, scale, W, H): """Read current canvas JSON back into xyxy boxes on the original image grid.""" if item.get("canvas_json"): item["boxes"] = parse_boxes(item["canvas_json"], scale, W, H) def _classify_now(item, predict_fn): """Run classifier on current boxes, then sync to session samples.""" if not item.get("boxes"): item["preds"] = []; item["user_labels"] = []; item["actions"] = [] sync_samples_with_state(); return crops = [item["pil"].crop(b) for b in item["boxes"]] with torch.inference_mode(): item["preds"] = predict_fn(crops, topk=3) # set_defaults fills current_label/confidence/margin/badges and seeds samples set_defaults_from_preds(item) item["actions"] = ["pending"] * len(item["boxes"]) sync_samples_with_state() def _bump_on_image_switch(ak: str, item: dict): """Force a canvas remount when the active image changes.""" if st.session_state.get("_last_active_key") != ak: item["canvas_rev"] = (item.get("canvas_rev", 0) + 1) st.session_state["_last_active_key"] = ak # ---------------- main UI ---------------- def render_workflow_tab(predict_fn, class_names, class_to_id, model_meta, export_root): CANVAS_W = _auto_canvas_width() if "uploader_rev" not in st.session_state: st.session_state.uploader_rev = 0 # =============== (1) Load images =============== st.subheader("1) Load image(s)") ups = st.file_uploader( "Upload one or more images", type=["jpg", "jpeg", "png", "bmp"], accept_multiple_files=True, key=f"uploader_{st.session_state.uploader_rev}" # resettable ) for up in (ups or []): key = f"{up.name}-{up.size}" if key not in st.session_state.images: pil = Image.open(up).convert("RGB") st.session_state.images[key] = { "key": key, "name": up.name, "pil": pil, "boxes": [], "preds": [], "user_labels": [], "actions": [], "detector_info": [], # <— conf/cls/name aligned with boxes "canvas_json": {"objects": []}, "canvas_rev": 0, } if st.session_state.active_key is None: st.session_state.active_key = key keys = list(st.session_state.images.keys()) if not keys: st.info("Upload images or load from S3/Drive in the sidebar.") return names = [st.session_state.images[k]["name"] for k in keys] idx = keys.index(st.session_state.active_key) if st.session_state.active_key in keys else 0 chosen = st.selectbox("Active image", names, index=idx) ak = keys[names.index(chosen)] st.session_state.active_key = ak # =============== (2) Canvas with Assist+ autorun =============== item = st.session_state.images[ak] base = item["pil"] disp_w = CANVAS_W scale = float(disp_w) / float(base.width) disp_h = int(round(base.height * scale)) _bump_on_image_switch(ak, item) # Assist+: auto-detect → seed canvas → auto-classify (only if no boxes yet) if st.session_state.assist_plus and not item.get("boxes") and not item.get("_skip_autodetect_once"): det_out = _run_detection(item) if det_out: boxes_xyxy = [tuple(int(v) for v in o["bbox"]) for o in det_out] det_info = [] for o in det_out: conf = o.get("score", o.get("conf", o.get("confidence"))) clsid = o.get("cls", o.get("class")) name = o.get("name") det_info.append({ "conf": float(conf) if conf is not None else None, "cls": int(clsid) if clsid is not None else None, "name": name if name is not None else None, }) item["boxes"] = boxes_xyxy item["detector_info"] = det_info item["canvas_json"] = seed_canvas_from_boxes( boxes_xyxy, scale, det_scores=[d["conf"] for d in det_info] ) item["canvas_mode"] = "edit" _parse_canvas_into_boxes(item, scale, base.width, base.height) _classify_now(item, predict_fn) st.rerun() # Build display background bg_pil = ImageOps.exif_transpose(base).resize((disp_w, disp_h), Image.BILINEAR).convert("RGB").copy() # 2) Draw / edit rectangles st.subheader("2) Draw / edit rectangles") # Per-image canvas mode (so you can move/resize in Edit mode) if "canvas_mode" not in item: item["canvas_mode"] = "edit" if item.get("boxes") else "draw" mode_idx = 1 if item["canvas_mode"] == "edit" else 0 canvas_mode_choice = st.radio( "Canvas mode", ["Draw", "Edit"], index=mode_idx, horizontal=True, key=f"canvas_mode_{ak}_{item.get('canvas_rev', 0)}" ) item["canvas_mode"] = "edit" if canvas_mode_choice == "Edit" else "draw" drawing_mode = "transform" if item["canvas_mode"] == "edit" else "rect" # Hints if item["canvas_mode"] == "edit": st.caption("Tip: Click a box to move/resize. Use the live list below to remove a box.") else: st.caption("Tip: Drag on the image to draw new boxes.") # ---------- initial drawing ---------- base_json = item.get("canvas_json") or {"objects": []} show_labels = st.checkbox("Show box labels on canvas", value=True, key=f"show_labels_{ak}") display_json = inject_index_labels(base_json, show_conf=False) if show_labels else base_json # ---------- Canvas ---------- canvas = st_canvas( fill_color="rgba(0,0,0,0)", stroke_width=3, stroke_color="#FF9900", background_image=bg_pil, update_streamlit=True, height=disp_h, width=disp_w, drawing_mode=drawing_mode, display_toolbar=True, initial_drawing=display_json, key=f"canvas_{ak}_{item.get('canvas_rev', 0)}_{disp_w}x{disp_h}" ) # ---------- Live canvas snapshot ---------- live_json = canvas.json_data or base_json objs = list(live_json.get("objects") or []) # Rect indices & objects (ignore overlay textboxes) rect_indices = [k for k, o in enumerate(objs) if o.get("type") == "rect"] rect_objs = [objs[k] for k in rect_indices] # Live boxes in original image coords live_boxes = parse_boxes(live_json, scale, base.width, base.height) # Detector summary (use detector_info matched to live boxes) infos = _match_det_info_to_live_boxes(item, live_boxes) confs = [d["conf"] for d in infos if d and d.get("conf") is not None] if confs: n = len(confs) cmin, cmax = min(confs), max(confs) cmean = sum(confs)/n th = float(st.session_state.get("det_cfg", {}).get("conf", 0.25)) below = sum(1 for c in confs if c < th) st.markdown( f"**Active image (detector)** — boxes: {n} · mean **{cmean:.3f}** · " f"min **{cmin:.3f}** · max **{cmax:.3f}** · below {th:.2f}: **{below}/{n}**" ) # Commit canvas into boxes (Clear / Remove) col_clear, col_remove = st.columns([1, 1]) if col_clear.button("Clear boxes (THIS image)", key=f"clear_{ak}"): item["canvas_json"] = {"objects": []} item["boxes"] = []; item["preds"] = []; item["user_labels"] = []; item["actions"] = [] item["detector_info"] = [] # also clear detector stats item["canvas_mode"] = "draw" item["canvas_rev"] = (item.get("canvas_rev", 0) + 1) item["_skip_autodetect_once"] = True st.session_state.all_samples = [s for s in st.session_state.all_samples if s.image_key != ak] st.rerun() if col_remove.button("Remove image", key=f"remove_{ak}"): remaining_keys = [k for k in st.session_state.images.keys() if k != ak] next_key = remaining_keys[0] if remaining_keys else None del st.session_state.images[ak] st.session_state.all_samples = [s for s in st.session_state.all_samples if s.image_key != ak] st.session_state.uploader_rev += 1 st.session_state.active_key = next_key if next_key: nxt = st.session_state.images[next_key] nxt["_skip_autodetect_once"] = True nxt["canvas_mode"] = "edit" if nxt.get("boxes") else "draw" nxt["canvas_rev"] = nxt.get("canvas_rev", 0) + 1 st.rerun() # =============== (3) Detect / Classify controls =============== st.subheader("3) Detect & Classify") assist_on = bool(st.session_state.get("assist_plus")) has_boxes = bool(item.get("boxes")) canvas_objects = (canvas.json_data or {}).get("objects") or [] has_canvas_rects = any(obj.get("type") == "rect" for obj in canvas_objects) can_classify = has_boxes or has_canvas_rects # ---------- DETECT ---------- if assist_on: st.caption("Detect: find objects and write boxes onto the canvas.") det_mode = st.radio( "Detection mode", ["Replace current boxes", "Append to existing"], index=0, horizontal=True, key=f"det_mode_{ak}", help="Replace overwrites all boxes; Append merges new detections." ) det_label = "Detect objects (active image)" if not has_boxes else "Run detection (active image)" if st.button( det_label, help="Run detection with the mode selected above.", key=f"detect_{ak}" ): det_out = _run_detection(item) new_boxes = [tuple(int(v) for v in o["bbox"]) for o in det_out] new_info = [] for o in det_out: conf = o.get("score", o.get("conf", o.get("confidence"))) clsid = o.get("cls", o.get("class")) name = o.get("name") new_info.append({ "conf": float(conf) if conf is not None else None, "cls": int(clsid) if clsid is not None else None, "name": name if name is not None else None, }) if new_boxes: old_boxes = list(item.get("boxes") or []) old_info = list(item.get("detector_info") or []) if det_mode.startswith("Replace"): merged_boxes = new_boxes merged_info = new_info else: merged_boxes = old_boxes + new_boxes merged_info = old_info + new_info item["boxes"] = merged_boxes item["detector_info"] = merged_info item["canvas_json"] = seed_canvas_from_boxes( merged_boxes, scale, det_scores=[d["conf"] for d in merged_info] ) item["canvas_mode"] = "edit" item["canvas_rev"] = (item.get("canvas_rev", 0) + 1) # auto-run classification so users don't need a second click after detection _parse_canvas_into_boxes(item, scale, base.width, base.height) _classify_now(item, predict_fn) st.rerun() else: st.info("No objects detected. Try adjusting thresholds or draw boxes.") # ---------- CLASSIFY ---------- st.caption("Classify: run the classifier on the current canvas boxes.") has_preds = bool(item.get("preds")) boxes_changed = live_boxes != (item.get("boxes") or []) classify_label = "Run classification" if not has_preds else ("Update after edits" if boxes_changed else "Re-run classification") classify_help = "Add or detect a box first." if not can_classify else "Run the classifier on the current canvas boxes." if st.button( classify_label, disabled=not can_classify, help=classify_help, key=f"classify_{ak}" ): item["canvas_json"] = canvas.json_data or item.get("canvas_json") or {} _parse_canvas_into_boxes(item, scale, base.width, base.height) _classify_now(item, predict_fn) st.rerun() # ---------- Crops (single list for remove + relabel + stats) ---------- st.subheader("Crops") st.caption(f"Crops in this image: **{len(live_boxes)}**") samples = [s for s in st.session_state.all_samples if s.image_key == ak] samples_by_idx = {s.crop_idx: s for s in samples} if not live_boxes: st.info("Draw boxes or run detection to see crops.") else: for i, b in enumerate(live_boxes): info = infos[i] if i < len(infos) else None det_conf = (info or {}).get("conf", None) det_label = None if info: if info.get("name") is not None: det_label = str(info["name"]) elif info.get("cls") is not None: names = st.session_state.get("detector_names") or [] cid = info["cls"] if isinstance(names, (list, tuple)) and isinstance(cid, int) and 0 <= cid < len(names): det_label = str(names[cid]) else: det_label = str(cid) sample = samples_by_idx.get(i) col_img, col_meta, col_actions = st.columns([1.2, 2.6, 1.3]) with col_img: try: crop_image = base.crop(b) st.image(crop_image, caption=f"Crop {i+1}", width=180) except Exception: st.caption(f"Crop {i+1}") with col_meta: st.write(f"Box #{i+1} [x0={b[0]}, y0={b[1]}, x1={b[2]}, y1={b[3]}]") if det_conf is None and not det_label: st.caption("det: —") else: th = float(st.session_state.get("det_cfg", {}).get("conf", 0.25)) status = " ⚠️" if (det_conf is not None and det_conf < th) else "" parts = [] if det_conf is not None: parts.append(f"{det_conf:.2f}") if det_label: parts.append(det_label) st.caption("det: " + " · ".join(parts) + status) if sample: clf_top_name, clf_top_prob = (sample.top3[0] if sample.top3 else (sample.current_label, sample.confidence)) if det_conf is not None and clf_top_prob is not None: labels_match = bool(det_label and clf_top_name and str(det_label) == str(clf_top_name)) det_flag = ":red[low det]" if det_conf < th else ":green[det ok]" lbl_flag = ":green[label match]" if labels_match else (":orange[label diff]" if det_label else ":gray[det unlabeled]") st.markdown( f"**Det vs Clf:** det **{det_conf:.2f}**{f' ({det_label})' if det_label else ''} " f"vs clf **{clf_top_prob:.3f}** ({clf_top_name}) · {det_flag} · {lbl_flag}" ) top3_str = ", ".join([f"{n} ({p:.3f})" for n, p in sample.top3[:3]]) st.write(f"**Current Label:** {sample.current_label}") st.write(f"**Top-3:** {top3_str}") st.write(f"**Confidence:** {sample.confidence:.3f}") st.write(f"**Margin:** {sample.margin:.3f}") badge_str = " ".join([f":red[{b}]" for b in sample.badges]) if sample.badges else ":green[OK]" st.write(f"**Badges:** {badge_str}") st.write(f"**Action Status:** {(':green[' if sample.action=='accepted' else ':red[') + sample.action.upper() + ']'}") else: st.caption("Not classified yet.") with col_actions: if sample: try: idx = class_names.index(sample.current_label) if sample.current_label in class_names else 0 except Exception: idx = 0 new_label = st.selectbox( f"Relabel Crop {sample.crop_idx+1}", class_names, index=idx, key=f"wf_relabel_{ak}_{i}", ) if new_label != sample.current_label: sample.user_label = new_label sample.action = "relabel" from core.state import update_image_state_from_samples update_image_state_from_samples() st.rerun() else: st.caption("Classify before relabeling.") if st.button("Remove", key=f"rm_crop_{ak}_{i}"): true_idx = rect_indices[i] if i < len(rect_indices) else None if true_idx is not None and 0 <= true_idx < len(objs): old_boxes = item.get("boxes") or [] old_info = item.get("detector_info") or [] objs.pop(true_idx) if true_idx < len(objs) and objs[true_idx].get("type") == "textbox": objs.pop(true_idx) live_json["objects"] = objs item["canvas_json"] = live_json _parse_canvas_into_boxes(item, scale, base.width, base.height) new_boxes = item.get("boxes") or [] item["detector_info"] = _reindex_det_info(old_boxes, old_info, new_boxes) item["canvas_mode"] = "edit" item["canvas_rev"] = (item.get("canvas_rev", 0) + 1) item["preds"] = []; item["user_labels"] = []; item["actions"] = [] sync_samples_with_state() st.rerun() st.divider()