{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Material Hunters POC Notebook\n", "\n", "This notebook runs a minimal end-to-end baseline:\n", "- box/brush guided segmentation with SAM 3.1 (`facebook/sam3.1`)\n", "- optional brush-hint guidance (converted to SAM box + positive points)\n", "- material classification with SigCLIP\n", "- optional user edits\n", "- passport JSON export\n", "- evaluation row append + go/no-go metrics\n", "\n", "Requires access to the gated SAM 3.1 checkpoint on Hugging Face." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Run once in Colab if needed.\n", "%pip -q install transformers accelerate pillow matplotlib pandas huggingface_hub\n", "%pip -q install git+https://github.com/facebookresearch/sam3.git" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from pathlib import Path\n", "from contextlib import nullcontext\n", "import os\n", "import base64\n", "import io\n", "from datetime import datetime, timezone\n", "import json\n", "\n", "import numpy as np\n", "import pandas as pd\n", "import torch\n", "import matplotlib.pyplot as plt\n", "import matplotlib.patches as patches\n", "from PIL import Image\n", "\n", "from transformers import AutoImageProcessor, AutoModel, AutoProcessor\n", "\n", "from huggingface_hub import whoami\n", "from sam3.model_builder import build_sam3_image_model, download_ckpt_from_hf\n", "from sam3.model.sam3_image_processor import Sam3Processor\n", "\n", "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", "SAM3_REPO_ID = \"facebook/sam3.1\"\n", "SAM3_CHECKPOINT_PATH = None # set a local path like /content/checkpoints/sam3.1_multiplex.pt to skip HF download\n", "SAM3_CONFIDENCE_THRESHOLD = 0.10\n", "SAM3_POINT_HINT_BOX_RADIUS_PX = 12\n", "SAM3_MAX_POINT_HINTS = 4\n", "SAM3_TORCH_DTYPE = torch.bfloat16 if DEVICE == \"cuda\" else torch.float32\n", "DINO_MODEL_ID = \"facebook/dinov3-vitb16-pretrain-lvd1689m\"\n", "SIGCLIP_FUSION_WEIGHT = 0.60\n", "DINO_FUSION_WEIGHT = 0.40\n", "DINO_SIM_TEMPERATURE = 0.10\n", "LOW_CONF_THRESHOLD = 0.50\n", "MIXED_MARGIN_THRESHOLD = 0.15\n", "MAX_MIXED_ENTRIES = 3\n", "MIN_MIXED_PERCENT = 10\n", "MATERIAL_PROTOTYPE_IMAGE_PATHS = {\n", " # \"wood\": [\"/content/prototypes/wood_01.jpg\", \"/content/prototypes/wood_02.jpg\"],\n", " # \"steel\": [\"/content/prototypes/steel_01.jpg\"],\n", "}\n", "SIGCLIP_MODEL_ID = \"google/siglip-base-patch16-224\"\n", "\n", "CSV_PATH = Path(\"/content/poc_eval_sheet.csv\")\n", "TEMPLATE_PATH = Path(\"/content/poc_eval_sheet_template.csv\")\n", "PASSPORT_DIR = Path(\"/content/passports\")\n", "PASSPORT_DIR.mkdir(parents=True, exist_ok=True)\n", "\n", "THRESHOLDS = {\n", " \"segmentation_acceptance_rate\": 0.85,\n", " \"top1_accuracy\": 0.70,\n", " \"top3_hit_rate\": 0.90,\n", " \"manual_edit_rate_max\": 0.35,\n", " \"median_e2e_time_sec_max\": 12.0,\n", "}\n", "\n", "MATERIAL_TAXONOMY = {\n", " \"concrete\": [\"concrete\", \"cement\"],\n", " \"brick\": [\"brick\", \"masonry brick\"],\n", " \"wood\": [\"wood\", \"timber\", \"plywood\"],\n", " \"steel\": [\"steel\", \"iron metal\"],\n", " \"aluminum\": [\"aluminum\", \"aluminium\"],\n", " \"glass\": [\"glass\", \"window glass\"],\n", " \"plastic\": [\"plastic\", \"polymer\", \"pvc\"],\n", " \"ceramic\": [\"ceramic\", \"tile\", \"porcelain\"],\n", " \"gypsum\": [\"gypsum\", \"drywall\", \"plasterboard\"],\n", " \"insulation wool\": [\"insulation wool\", \"mineral wool\", \"fiberglass insulation\"],\n", "}\n", "\n", "PROMPT_TEMPLATES = [\n", " \"a close-up photo of {material}\",\n", " \"a construction material made of {material}\",\n", " \"the surface texture is {material}\",\n", "]\n", "\n", "OBJECT_MATERIAL_PRIORS = {\n", " \"window\": [\"glass\", \"aluminum\", \"wood\", \"plastic\", \"steel\"],\n", " \"door\": [\"wood\", \"steel\", \"aluminum\", \"glass\", \"plastic\"],\n", " \"wall\": [\"concrete\", \"brick\", \"gypsum\", \"wood\"],\n", " \"floor\": [\"concrete\", \"wood\", \"ceramic\", \"plastic\"],\n", " \"roof\": [\"steel\", \"wood\", \"ceramic\", \"insulation wool\"],\n", " \"pipe\": [\"steel\", \"plastic\", \"aluminum\"],\n", " \"beam\": [\"steel\", \"wood\", \"concrete\"],\n", " \"column\": [\"concrete\", \"steel\", \"brick\", \"wood\"],\n", "}" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def ensure_eval_csv(csv_path: Path, template_path: Path):\n", " if csv_path.exists():\n", " return\n", " if not template_path.exists():\n", " raise FileNotFoundError(\n", " f\"Missing {template_path}. Upload poc_eval_sheet_template.csv first.\"\n", " )\n", " pd.read_csv(template_path).iloc[0:0].to_csv(csv_path, index=False)\n", "\n", "\n", "def load_image(image_path: str) -> Image.Image:\n", " return Image.open(image_path).convert(\"RGB\")\n", "\n", "\n", "def select_bbox_colab(image: Image.Image, prompt: str = \"Draw a box around the object\"):\n", " try:\n", " from google.colab.output import eval_js\n", " from IPython.display import Javascript, display\n", " except Exception as exc:\n", " raise RuntimeError(\"Box selection requires Google Colab.\") from exc\n", "\n", " buf = io.BytesIO()\n", " image.save(buf, format=\"PNG\")\n", " data_url = \"data:image/png;base64,\" + base64.b64encode(buf.getvalue()).decode(\"utf-8\")\n", "\n", " js = \"\"\"\n", " async function materialHuntersPickBox(dataUrl, promptText) {\n", " const overlay = document.createElement('div');\n", " overlay.style.position = 'fixed';\n", " overlay.style.inset = '0';\n", " overlay.style.background = 'rgba(0,0,0,0.75)';\n", " overlay.style.zIndex = '9999';\n", " overlay.style.display = 'flex';\n", " overlay.style.flexDirection = 'column';\n", " overlay.style.alignItems = 'center';\n", " overlay.style.justifyContent = 'center';\n", " overlay.style.gap = '8px';\n", "\n", " const label = document.createElement('div');\n", " label.textContent = `${promptText} (click-drag to draw, Done to confirm, Esc to cancel)`;\n", " label.style.color = 'white';\n", " label.style.font = '16px sans-serif';\n", "\n", " const container = document.createElement('div');\n", " container.style.position = 'relative';\n", "\n", " const img = document.createElement('img');\n", " img.src = dataUrl;\n", " img.style.width = '95vw';\n", " img.style.maxWidth = '1800px';\n", " img.style.maxHeight = '90vh';\n", " img.style.border = '2px solid white';\n", " img.style.display = 'block';\n", " img.style.userSelect = 'none';\n", "\n", " const canvas = document.createElement('canvas');\n", " canvas.style.position = 'absolute';\n", " canvas.style.left = '0';\n", " canvas.style.top = '0';\n", " canvas.style.cursor = 'crosshair';\n", " canvas.style.touchAction = 'none';\n", "\n", " container.appendChild(img);\n", " container.appendChild(canvas);\n", "\n", " const controls = document.createElement('div');\n", " controls.style.display = 'flex';\n", " controls.style.gap = '8px';\n", " const doneBtn = document.createElement('button');\n", " doneBtn.textContent = 'Done';\n", " const resetBtn = document.createElement('button');\n", " resetBtn.textContent = 'Reset';\n", " controls.appendChild(doneBtn);\n", " controls.appendChild(resetBtn);\n", "\n", " overlay.appendChild(label);\n", " overlay.appendChild(container);\n", " overlay.appendChild(controls);\n", " document.body.appendChild(overlay);\n", "\n", " await new Promise((resolve) => {\n", " if (img.complete && img.naturalWidth > 0) resolve();\n", " else img.onload = () => resolve();\n", " });\n", "\n", " const resizeCanvas = () => {\n", " const r = img.getBoundingClientRect();\n", " canvas.width = Math.max(1, Math.round(r.width));\n", " canvas.height = Math.max(1, Math.round(r.height));\n", " canvas.style.width = `${r.width}px`;\n", " canvas.style.height = `${r.height}px`;\n", " };\n", " resizeCanvas();\n", "\n", " const ctx = canvas.getContext('2d');\n", " const clamp = (v, lo, hi) => Math.max(lo, Math.min(hi, v));\n", " const toCanvas = (e) => {\n", " const r = canvas.getBoundingClientRect();\n", " return [\n", " clamp(e.clientX - r.left, 0, r.width - 1),\n", " clamp(e.clientY - r.top, 0, r.height - 1),\n", " ];\n", " };\n", " const toNatural = (x, y) => [\n", " Math.round(x * (img.naturalWidth / canvas.width)),\n", " Math.round(y * (img.naturalHeight / canvas.height)),\n", " ];\n", "\n", " let dragging = false;\n", " let sx = 0;\n", " let sy = 0;\n", " let latest = null;\n", "\n", " const redraw = () => {\n", " ctx.clearRect(0, 0, canvas.width, canvas.height);\n", " if (!latest) return;\n", " const [x0, y0, x1, y1] = latest;\n", " ctx.fillStyle = 'rgba(0,229,255,0.18)';\n", " ctx.strokeStyle = 'rgba(0,229,255,1)';\n", " ctx.lineWidth = 2;\n", " ctx.fillRect(x0, y0, x1 - x0, y1 - y0);\n", " ctx.strokeRect(x0, y0, x1 - x0, y1 - y0);\n", " };\n", "\n", " canvas.onpointerdown = (e) => {\n", " canvas.setPointerCapture(e.pointerId);\n", " dragging = true;\n", " [sx, sy] = toCanvas(e);\n", " latest = [sx, sy, sx + 1, sy + 1];\n", " redraw();\n", " };\n", "\n", " canvas.onpointermove = (e) => {\n", " if (!dragging) return;\n", " const [cx, cy] = toCanvas(e);\n", " latest = [Math.min(sx, cx), Math.min(sy, cy), Math.max(sx, cx), Math.max(sy, cy)];\n", " redraw();\n", " };\n", "\n", " const stopDrag = () => { dragging = false; };\n", " canvas.onpointerup = stopDrag;\n", " canvas.onpointercancel = stopDrag;\n", "\n", " resetBtn.onclick = () => { latest = null; redraw(); };\n", "\n", " let resolvePromise;\n", " const promise = new Promise((resolve) => { resolvePromise = resolve; });\n", "\n", " const cleanup = () => {\n", " document.removeEventListener('keydown', onKey);\n", " window.removeEventListener('resize', resizeCanvas);\n", " overlay.remove();\n", " };\n", "\n", " const onKey = (e) => {\n", " if (e.key === 'Escape') {\n", " cleanup();\n", " resolvePromise([-1, -1, -1, -1]);\n", " }\n", " };\n", " document.addEventListener('keydown', onKey);\n", " window.addEventListener('resize', resizeCanvas);\n", "\n", " doneBtn.onclick = () => {\n", " if (!latest) {\n", " label.textContent = 'Draw a box first.';\n", " return;\n", " }\n", " const [cx0, cy0, cx1, cy1] = latest;\n", " if ((cx1 - cx0) < 3 || (cy1 - cy0) < 3) {\n", " label.textContent = 'Box too small. Drag a larger box.';\n", " return;\n", " }\n", " const [nx0, ny0] = toNatural(cx0, cy0);\n", " const [nx1, ny1] = toNatural(cx1, cy1);\n", " cleanup();\n", " resolvePromise([Math.min(nx0,nx1), Math.min(ny0,ny1), Math.max(nx0,nx1), Math.max(ny0,ny1)]);\n", " };\n", "\n", " return await promise;\n", " }\n", " \"\"\"\n", "\n", " display(Javascript(js))\n", " x0, y0, x1, y1 = eval_js(\n", " f\"materialHuntersPickBox({json.dumps(data_url)}, {json.dumps(prompt)})\"\n", " )\n", " if x0 < 0 or y0 < 0:\n", " raise RuntimeError(\"Box selection cancelled.\")\n", " if x1 <= x0 or y1 <= y0:\n", " raise RuntimeError(\"Invalid box. Draw a box with non-zero area.\")\n", " return (int(x0), int(y0), int(x1), int(y1))\n", "\n", "\n", "def select_multiple_bboxes_colab(image: Image.Image, expected_boxes: int):\n", " if expected_boxes < 1:\n", " raise ValueError(\"expected_boxes must be >= 1\")\n", " bboxes = []\n", " for i in range(expected_boxes):\n", " bbox = select_bbox_colab(image, prompt=f\"Draw box {i + 1}/{expected_boxes}\")\n", " bboxes.append(bbox)\n", " return bboxes\n", "\n", "\n", "def review_bbox_acceptance(image: Image.Image, bbox_xyxy, title: str = \"Review selected box\"):\n", " x0, y0, x1, y1 = [int(v) for v in bbox_xyxy]\n", " plt.figure(figsize=(8, 6))\n", " ax = plt.gca()\n", " ax.imshow(image)\n", " rect = patches.Rectangle((x0, y0), x1 - x0, y1 - y0, linewidth=2, edgecolor=\"cyan\", facecolor=\"none\")\n", " ax.add_patch(rect)\n", " ax.set_title(title)\n", " ax.axis(\"off\")\n", " plt.show()\n", "\n", " answer = input(\"Accept this box? [y/n]: \" ).strip().lower()\n", " return answer in {\"y\", \"yes\"}\n", "\n", "\n", "def select_bbox_with_redraw_colab(image: Image.Image, prompt: str, max_attempts: int = 5):\n", " for attempt in range(1, max_attempts + 1):\n", " bbox = select_bbox_colab(image, prompt=f\"{prompt} (attempt {attempt}/{max_attempts})\")\n", " accepted = review_bbox_acceptance(image, bbox, title=f\"{prompt} - attempt {attempt}\")\n", " if accepted:\n", " return bbox\n", " print(\"Re-draw requested.\")\n", "\n", " raise RuntimeError(\"Box not accepted within max attempts.\")\n", "\n", "\n", "def select_multiple_bboxes_with_redraw_colab(image: Image.Image, expected_boxes: int, max_attempts_per_box: int = 5):\n", " if expected_boxes < 1:\n", " raise ValueError(\"expected_boxes must be >= 1\")\n", " bboxes = []\n", " for i in range(expected_boxes):\n", " bbox = select_bbox_with_redraw_colab(\n", " image=image,\n", " prompt=f\"Draw box {i + 1}/{expected_boxes}\",\n", " max_attempts=max_attempts_per_box,\n", " )\n", " bboxes.append(bbox)\n", " return bboxes\n", "\n", "\n", "def select_brush_hint_colab(image: Image.Image, prompt: str = \"Brush over target object\"):\n", " try:\n", " from google.colab.output import eval_js\n", " from IPython.display import Javascript, display\n", " except Exception as exc:\n", " raise RuntimeError(\"Brush selection requires Google Colab.\") from exc\n", "\n", " buf = io.BytesIO()\n", " image.save(buf, format=\"PNG\")\n", " data_url = \"data:image/png;base64,\" + base64.b64encode(buf.getvalue()).decode(\"utf-8\")\n", "\n", " js = \"\"\"\n", " async function materialHuntersBrushHint(dataUrl, promptText) {\n", " const overlay = document.createElement('div');\n", " overlay.style.position = 'fixed';\n", " overlay.style.inset = '0';\n", " overlay.style.background = 'rgba(0,0,0,0.75)';\n", " overlay.style.zIndex = '9999';\n", " overlay.style.display = 'flex';\n", " overlay.style.flexDirection = 'column';\n", " overlay.style.alignItems = 'center';\n", " overlay.style.justifyContent = 'center';\n", " overlay.style.gap = '8px';\n", "\n", " const label = document.createElement('div');\n", " label.textContent = `${promptText} (paint object area, Done to confirm, Esc to cancel)`;\n", " label.style.color = 'white';\n", " label.style.font = '16px sans-serif';\n", "\n", " const container = document.createElement('div');\n", " container.style.position = 'relative';\n", "\n", " const img = document.createElement('img');\n", " img.src = dataUrl;\n", " img.style.width = '95vw';\n", " img.style.maxWidth = '1800px';\n", " img.style.maxHeight = '90vh';\n", " img.style.display = 'block';\n", " img.style.userSelect = 'none';\n", " img.style.border = '2px solid white';\n", "\n", " const canvas = document.createElement('canvas');\n", " canvas.style.position = 'absolute';\n", " canvas.style.left = '0';\n", " canvas.style.top = '0';\n", " canvas.style.cursor = 'crosshair';\n", " canvas.style.touchAction = 'none';\n", "\n", " container.appendChild(img);\n", " container.appendChild(canvas);\n", "\n", " const controls = document.createElement('div');\n", " controls.style.display = 'flex';\n", " controls.style.gap = '8px';\n", " const doneBtn = document.createElement('button');\n", " doneBtn.textContent = 'Done';\n", " const clearBtn = document.createElement('button');\n", " clearBtn.textContent = 'Clear';\n", " const brush = document.createElement('input');\n", " brush.type = 'range';\n", " brush.min = '4';\n", " brush.max = '48';\n", " brush.value = '18';\n", "\n", " controls.appendChild(doneBtn);\n", " controls.appendChild(clearBtn);\n", " controls.appendChild(brush);\n", "\n", " overlay.appendChild(label);\n", " overlay.appendChild(container);\n", " overlay.appendChild(controls);\n", " document.body.appendChild(overlay);\n", "\n", " await new Promise((resolve) => {\n", " if (img.complete && img.naturalWidth > 0) resolve();\n", " else img.onload = () => resolve();\n", " });\n", "\n", " const resizeCanvas = () => {\n", " const r = img.getBoundingClientRect();\n", " canvas.width = Math.max(1, Math.round(r.width));\n", " canvas.height = Math.max(1, Math.round(r.height));\n", " canvas.style.width = `${r.width}px`;\n", " canvas.style.height = `${r.height}px`;\n", " };\n", " resizeCanvas();\n", "\n", " const ctx = canvas.getContext('2d');\n", " const clamp = (v, lo, hi) => Math.max(lo, Math.min(hi, v));\n", " const toCanvas = (e) => {\n", " const r = canvas.getBoundingClientRect();\n", " return [\n", " clamp(e.clientX - r.left, 0, r.width - 1),\n", " clamp(e.clientY - r.top, 0, r.height - 1),\n", " ];\n", " };\n", " const toNatural = (x, y) => [\n", " Math.round(x * (img.naturalWidth / canvas.width)),\n", " Math.round(y * (img.naturalHeight / canvas.height)),\n", " ];\n", "\n", " const drawn = [];\n", " let drawing = false;\n", " let lx = 0;\n", " let ly = 0;\n", "\n", " const paintSegment = (x0, y0, x1, y1) => {\n", " ctx.lineCap = 'round';\n", " ctx.lineJoin = 'round';\n", " ctx.strokeStyle = 'rgba(0,255,255,0.85)';\n", " ctx.lineWidth = Number(brush.value);\n", " ctx.beginPath();\n", " ctx.moveTo(x0, y0);\n", " ctx.lineTo(x1, y1);\n", " ctx.stroke();\n", " };\n", "\n", " const maybeRecord = (x, y) => {\n", " if (drawn.length === 0) { drawn.push([x, y]); return; }\n", " const [px, py] = drawn[drawn.length - 1];\n", " const dx = x - px;\n", " const dy = y - py;\n", " if ((dx*dx + dy*dy) >= 9) drawn.push([x, y]);\n", " };\n", "\n", " canvas.onpointerdown = (e) => {\n", " canvas.setPointerCapture(e.pointerId);\n", " drawing = true;\n", " [lx, ly] = toCanvas(e);\n", " paintSegment(lx, ly, lx, ly);\n", " maybeRecord(lx, ly);\n", " };\n", "\n", " canvas.onpointermove = (e) => {\n", " if (!drawing) return;\n", " const [cx, cy] = toCanvas(e);\n", " paintSegment(lx, ly, cx, cy);\n", " lx = cx;\n", " ly = cy;\n", " maybeRecord(cx, cy);\n", " };\n", "\n", " const stopDraw = () => { drawing = false; };\n", " canvas.onpointerup = stopDraw;\n", " canvas.onpointercancel = stopDraw;\n", "\n", " clearBtn.onclick = () => {\n", " ctx.clearRect(0, 0, canvas.width, canvas.height);\n", " drawn.length = 0;\n", " };\n", "\n", " let resolvePromise;\n", " const promise = new Promise((resolve) => { resolvePromise = resolve; });\n", "\n", " const cleanup = () => {\n", " document.removeEventListener('keydown', onKey);\n", " window.removeEventListener('resize', resizeCanvas);\n", " overlay.remove();\n", " };\n", "\n", " const onKey = (e) => {\n", " if (e.key === 'Escape') {\n", " cleanup();\n", " resolvePromise({bbox:[-1,-1,-1,-1], points:[]});\n", " }\n", " };\n", " document.addEventListener('keydown', onKey);\n", " window.addEventListener('resize', resizeCanvas);\n", "\n", " doneBtn.onclick = () => {\n", " if (drawn.length < 1) {\n", " label.textContent = 'Paint over the object before pressing Done.';\n", " return;\n", " }\n", "\n", " let minX = Infinity, minY = Infinity, maxX = -Infinity, maxY = -Infinity;\n", " for (const [x, y] of drawn) {\n", " if (x < minX) minX = x;\n", " if (y < minY) minY = y;\n", " if (x > maxX) maxX = x;\n", " if (y > maxY) maxY = y;\n", " }\n", "\n", " const pad = Math.max(4, Math.round(Number(brush.value) * 0.8));\n", " minX = Math.max(0, minX - pad);\n", " minY = Math.max(0, minY - pad);\n", " maxX = Math.min(canvas.width - 1, maxX + pad);\n", " maxY = Math.min(canvas.height - 1, maxY + pad);\n", "\n", " const [nx0, ny0] = toNatural(minX, minY);\n", " const [nx1, ny1] = toNatural(maxX, maxY);\n", "\n", " const sampled = [];\n", " const limit = 16;\n", " const step = Math.max(1, Math.floor(drawn.length / limit));\n", " for (let i = 0; i < drawn.length && sampled.length < limit; i += step) {\n", " sampled.push(toNatural(drawn[i][0], drawn[i][1]));\n", " }\n", "\n", " cleanup();\n", " resolvePromise({bbox:[nx0, ny0, nx1, ny1], points: sampled});\n", " };\n", "\n", " return await promise;\n", " }\n", " \"\"\"\n", "\n", " display(Javascript(js))\n", " result = eval_js(\n", " f\"materialHuntersBrushHint({json.dumps(data_url)}, {json.dumps(prompt)})\"\n", " )\n", " bbox = result.get(\"bbox\", [-1, -1, -1, -1])\n", " points = result.get(\"points\", [])\n", " if bbox[0] < 0 or bbox[1] < 0:\n", " raise RuntimeError(\"Brush selection cancelled.\")\n", " if bbox[2] <= bbox[0] or bbox[3] <= bbox[1]:\n", " raise RuntimeError(\"Invalid brush hint bbox.\")\n", " points_xy = [(int(p[0]), int(p[1])) for p in points]\n", " if not points_xy:\n", " cx = (int(bbox[0]) + int(bbox[2])) // 2\n", " cy = (int(bbox[1]) + int(bbox[3])) // 2\n", " points_xy = [(cx, cy)]\n", " return {\n", " \"prompt_bbox_xyxy\": [int(v) for v in bbox],\n", " \"prompt_points_xy\": points_xy,\n", " }\n", "\n", "\n", "def select_brush_hint_with_redraw_colab(image: Image.Image, prompt: str, max_attempts: int = 5):\n", " for attempt in range(1, max_attempts + 1):\n", " hint = select_brush_hint_colab(image, prompt=f\"{prompt} (attempt {attempt}/{max_attempts})\")\n", " accepted = review_bbox_acceptance(\n", " image=image,\n", " bbox_xyxy=hint[\"prompt_bbox_xyxy\"],\n", " title=f\"{prompt} - attempt {attempt}\",\n", " )\n", " if accepted:\n", " return hint\n", " print(\"Re-draw requested.\")\n", "\n", " raise RuntimeError(\"Brush hint not accepted within max attempts.\")\n", "\n", "\n", "def mask_bbox_xywh(mask: np.ndarray):\n", " ys, xs = np.where(mask)\n", " if len(xs) == 0 or len(ys) == 0:\n", " return (0, 0, 0, 0)\n", " x0, x1 = int(xs.min()), int(xs.max())\n", " y0, y1 = int(ys.min()), int(ys.max())\n", " return (x0, y0, x1 - x0 + 1, y1 - y0 + 1)\n", "\n", "\n", "def crop_masked_rgb(image: Image.Image, mask: np.ndarray, bbox_xywh, pad: int = 8) -> Image.Image:\n", " arr = np.array(image)\n", " h, w = arr.shape[:2]\n", " x, y, bw, bh = bbox_xywh\n", " x0 = max(0, x - pad)\n", " y0 = max(0, y - pad)\n", " x1 = min(w, x + bw + pad)\n", " y1 = min(h, y + bh + pad)\n", "\n", " crop = arr[y0:y1, x0:x1].copy()\n", " crop_mask = mask[y0:y1, x0:x1]\n", "\n", " bg = np.full_like(crop, 255)\n", " bg[crop_mask] = crop[crop_mask]\n", " return Image.fromarray(bg)\n", "\n", "\n", "def crop_bbox_rgb(image: Image.Image, bbox_xywh, pad: int = 8) -> Image.Image:\n", " arr = np.array(image)\n", " h, w = arr.shape[:2]\n", " x, y, bw, bh = bbox_xywh\n", " x0 = max(0, x - pad)\n", " y0 = max(0, y - pad)\n", " x1 = min(w, x + bw + pad)\n", " y1 = min(h, y + bh + pad)\n", " return Image.fromarray(arr[y0:y1, x0:x1])\n", "\n", "\n", "def build_material_prompt_table(material_taxonomy: dict, prompt_templates: list):\n", " prompts = []\n", " prompt_materials = []\n", " for material, aliases in material_taxonomy.items():\n", " for alias in aliases:\n", " for template in prompt_templates:\n", " prompts.append(template.format(material=alias))\n", " prompt_materials.append(material)\n", " return prompts, prompt_materials\n", "\n", "\n", "ALL_PROMPTS, ALL_PROMPT_MATERIALS = build_material_prompt_table(MATERIAL_TAXONOMY, PROMPT_TEMPLATES)\n", "ALL_CANONICAL_MATERIALS = list(MATERIAL_TAXONOMY.keys())\n", "\n", "\n", "def resolve_candidate_materials(object_label: str | None):\n", " if not object_label:\n", " return ALL_CANONICAL_MATERIALS\n", " key = object_label.strip().lower()\n", " candidates = OBJECT_MATERIAL_PRIORS.get(key, [])\n", " if len(candidates) < 2:\n", " return ALL_CANONICAL_MATERIALS\n", " return candidates\n", "\n", "\n", "def bbox_xyxy_to_center_xy(bbox_xyxy):\n", " x0, y0, x1, y1 = [int(v) for v in bbox_xyxy]\n", " return ((x0 + x1) // 2, (y0 + y1) // 2)\n", "\n", "\n", "def sam3_autocast_context():\n", " if DEVICE == \"cuda\" and SAM3_TORCH_DTYPE in (torch.float16, torch.bfloat16):\n", " return torch.autocast(device_type=\"cuda\", dtype=SAM3_TORCH_DTYPE)\n", " return nullcontext()\n", "\n", "\n", "def bbox_xyxy_to_norm_cxcywh(bbox_xyxy, image_width: int, image_height: int):\n", " x0, y0, x1, y1 = [float(v) for v in bbox_xyxy]\n", " x0 = max(0.0, min(float(image_width - 1), x0))\n", " y0 = max(0.0, min(float(image_height - 1), y0))\n", " x1 = max(0.0, min(float(image_width - 1), x1))\n", " y1 = max(0.0, min(float(image_height - 1), y1))\n", "\n", " cx = ((x0 + x1) / 2.0) / max(float(image_width), 1.0)\n", " cy = ((y0 + y1) / 2.0) / max(float(image_height), 1.0)\n", " w = max(1.0, abs(x1 - x0)) / max(float(image_width), 1.0)\n", " h = max(1.0, abs(y1 - y0)) / max(float(image_height), 1.0)\n", "\n", " return [float(cx), float(cy), float(w), float(h)]\n", "\n", "\n", "def build_brush_hint_mask(image_shape_hw, prompt_bbox_xyxy, prompt_points_xy, point_radius_px: int = 16):\n", " h, w = image_shape_hw\n", " hint = np.zeros((h, w), dtype=bool)\n", "\n", " x0, y0, x1, y1 = [int(v) for v in prompt_bbox_xyxy]\n", " x0 = max(0, min(w - 1, x0))\n", " y0 = max(0, min(h - 1, y0))\n", " x1 = max(0, min(w - 1, x1))\n", " y1 = max(0, min(h - 1, y1))\n", "\n", " if not prompt_points_xy:\n", " if x1 > x0 and y1 > y0:\n", " hint[y0:y1, x0:x1] = True\n", " return hint\n", "\n", " yy, xx = np.ogrid[:h, :w]\n", " for px, py in prompt_points_xy:\n", " px = int(max(0, min(w - 1, px)))\n", " py = int(max(0, min(h - 1, py)))\n", " disk = (xx - px) ** 2 + (yy - py) ** 2 <= point_radius_px ** 2\n", " hint |= disk\n", "\n", " return hint\n", "\n", "\n", "def score_sam_candidate(mask, sam_score, prompt_bbox_xyxy, prompt_points_xy, hint_mask):\n", " mask = mask.astype(bool)\n", " mask_area = int(mask.sum())\n", "\n", " if mask_area == 0:\n", " return {\n", " \"total\": -1.0,\n", " \"sam_score\": float(sam_score),\n", " \"point_coverage\": 0.0,\n", " \"hint_iou\": 0.0,\n", " \"inside_box_ratio\": 0.0,\n", " \"mask_area\": 0,\n", " }\n", "\n", " h, w = mask.shape\n", " x0, y0, x1, y1 = [int(v) for v in prompt_bbox_xyxy]\n", " x0 = max(0, min(w - 1, x0))\n", " y0 = max(0, min(h - 1, y0))\n", " x1 = max(0, min(w, x1))\n", " y1 = max(0, min(h, y1))\n", "\n", " box_mask = np.zeros_like(mask)\n", " if x1 > x0 and y1 > y0:\n", " box_mask[y0:y1, x0:x1] = True\n", "\n", " inside_box_ratio = float((mask & box_mask).sum() / max(mask_area, 1))\n", "\n", " if prompt_points_xy:\n", " hits = 0\n", " for px, py in prompt_points_xy:\n", " if 0 <= int(py) < h and 0 <= int(px) < w and mask[int(py), int(px)]:\n", " hits += 1\n", " point_coverage = float(hits / len(prompt_points_xy))\n", " else:\n", " point_coverage = 0.0\n", "\n", " if hint_mask is not None:\n", " inter = float((mask & hint_mask).sum())\n", " union = float((mask | hint_mask).sum())\n", " hint_iou = inter / max(union, 1.0)\n", " else:\n", " hint_iou = 0.0\n", "\n", " if prompt_points_xy:\n", " total = 0.50 * hint_iou + 0.30 * point_coverage + 0.10 * inside_box_ratio + 0.10 * float(sam_score)\n", " else:\n", " total = 0.80 * float(sam_score) + 0.20 * inside_box_ratio\n", "\n", " return {\n", " \"total\": float(total),\n", " \"sam_score\": float(sam_score),\n", " \"point_coverage\": float(point_coverage),\n", " \"hint_iou\": float(hint_iou),\n", " \"inside_box_ratio\": float(inside_box_ratio),\n", " \"mask_area\": int(mask_area),\n", " }\n", "\n", "\n", "def run_sam_guided_segmentation(image: Image.Image, prompt_bbox_xyxy, prompt_points_xy, sam_model, sam_processor):\n", " img_w, img_h = image.size\n", "\n", " with torch.inference_mode(), sam3_autocast_context():\n", " state = sam_processor.set_image(image, state={})\n", "\n", " box_norm = bbox_xyxy_to_norm_cxcywh(prompt_bbox_xyxy, image_width=img_w, image_height=img_h)\n", " state = sam_processor.add_geometric_prompt(box=box_norm, label=True, state=state)\n", "\n", " if prompt_points_xy:\n", " hint_points = prompt_points_xy[:SAM3_MAX_POINT_HINTS]\n", " for px, py in hint_points:\n", " r = int(SAM3_POINT_HINT_BOX_RADIUS_PX)\n", " pbox = [\n", " int(max(0, px - r)),\n", " int(max(0, py - r)),\n", " int(min(img_w - 1, px + r)),\n", " int(min(img_h - 1, py + r)),\n", " ]\n", " pbox_norm = bbox_xyxy_to_norm_cxcywh(pbox, image_width=img_w, image_height=img_h)\n", " state = sam_processor.add_geometric_prompt(box=pbox_norm, label=True, state=state)\n", "\n", " masks_t = state.get(\"masks\", None)\n", " scores_t = state.get(\"scores\", None)\n", "\n", " if masks_t is None or scores_t is None:\n", " raise RuntimeError(\"SAM3 inference did not return masks/scores.\")\n", "\n", " masks_np = masks_t.detach().cpu().numpy()\n", " if masks_np.ndim == 4 and masks_np.shape[1] == 1:\n", " candidate_masks = masks_np[:, 0].astype(bool)\n", " elif masks_np.ndim == 3:\n", " candidate_masks = masks_np.astype(bool)\n", " else:\n", " raise RuntimeError(f\"Unexpected SAM3 mask shape: {masks_np.shape}\")\n", "\n", " sam_scores = scores_t.detach().cpu().numpy().astype(float)\n", " if sam_scores.ndim > 1:\n", " sam_scores = sam_scores.reshape(-1)\n", "\n", " if len(candidate_masks) == 0:\n", " raise RuntimeError(\"SAM3 returned zero candidate masks. Try redrawing the prompt.\")\n", "\n", " img_arr = np.array(image)\n", " hint_mask = None\n", " if prompt_points_xy:\n", " hint_mask = build_brush_hint_mask(img_arr.shape[:2], prompt_bbox_xyxy, prompt_points_xy)\n", "\n", " candidate_debug = []\n", " for idx in range(candidate_masks.shape[0]):\n", " metrics = score_sam_candidate(\n", " mask=candidate_masks[idx],\n", " sam_score=float(sam_scores[min(idx, len(sam_scores) - 1)]),\n", " prompt_bbox_xyxy=prompt_bbox_xyxy,\n", " prompt_points_xy=prompt_points_xy,\n", " hint_mask=hint_mask,\n", " )\n", " metrics[\"candidate_idx\"] = int(idx)\n", " candidate_debug.append(metrics)\n", "\n", " best = max(candidate_debug, key=lambda d: d[\"total\"])\n", " best_idx = int(best[\"candidate_idx\"])\n", " mask = candidate_masks[best_idx]\n", " bbox = mask_bbox_xywh(mask)\n", "\n", " return {\n", " \"mask\": mask,\n", " \"mask_area_px\": int(mask.sum()),\n", " \"bbox_xywh\": bbox,\n", " \"prompt_bbox_xyxy\": [int(v) for v in prompt_bbox_xyxy],\n", " \"prompt_points_xy\": [[int(x), int(y)] for x, y in prompt_points_xy],\n", " \"sam_score\": float(sam_scores[best_idx]),\n", " \"selection_score\": float(best[\"total\"]),\n", " \"selection_metrics\": {\n", " \"point_coverage\": float(best[\"point_coverage\"]),\n", " \"hint_iou\": float(best[\"hint_iou\"]),\n", " \"inside_box_ratio\": float(best[\"inside_box_ratio\"]),\n", " },\n", " \"sam_candidates\": candidate_debug,\n", " }\n", "\n", "\n", "def classify_material_sigclip(\n", " image_views: list,\n", " object_label: str | None,\n", " sigclip_model,\n", " sigclip_processor,\n", " top_k: int = 3,\n", " view_weights: list | None = None,\n", "):\n", " candidate_materials = resolve_candidate_materials(object_label)\n", "\n", " selected_idx = [\n", " i for i, m in enumerate(ALL_PROMPT_MATERIALS) if m in candidate_materials\n", " ]\n", " selected_prompts = [ALL_PROMPTS[i] for i in selected_idx]\n", " selected_prompt_materials = [ALL_PROMPT_MATERIALS[i] for i in selected_idx]\n", "\n", " if not selected_prompts:\n", " raise ValueError(\"No text prompts available for candidate materials.\")\n", "\n", " if view_weights is None:\n", " view_weights = [1.0 / len(image_views)] * len(image_views)\n", " if len(view_weights) != len(image_views):\n", " raise ValueError(\"view_weights length must match image_views length\")\n", "\n", " agg_scores = {m: 0.0 for m in candidate_materials}\n", "\n", " for image, weight in zip(image_views, view_weights):\n", " inputs = sigclip_processor(\n", " text=selected_prompts,\n", " images=image,\n", " return_tensors=\"pt\",\n", " padding=\"max_length\",\n", " )\n", " inputs = {k: v.to(DEVICE) if hasattr(v, \"to\") else v for k, v in inputs.items()}\n", "\n", " with torch.no_grad():\n", " outputs = sigclip_model(**inputs)\n", "\n", " probs = torch.softmax(outputs.logits_per_image[0], dim=0).detach().cpu().numpy()\n", "\n", " per_material_max = {m: 0.0 for m in candidate_materials}\n", " for p, m in zip(probs, selected_prompt_materials):\n", " per_material_max[m] = max(per_material_max[m], float(p))\n", "\n", " for m in candidate_materials:\n", " agg_scores[m] += weight * per_material_max[m]\n", "\n", " total = sum(agg_scores.values())\n", " if total > 0:\n", " agg_scores = {k: v / total for k, v in agg_scores.items()}\n", "\n", " ranked = sorted(agg_scores.items(), key=lambda kv: kv[1], reverse=True)\n", " top_k = min(top_k, len(ranked))\n", " top_items = ranked[:top_k]\n", "\n", " top_labels = [x[0] for x in top_items]\n", " top_scores = [x[1] for x in top_items]\n", "\n", " return {\n", " \"pred_material_top1\": top_labels[0],\n", " \"pred_material_top3\": \"|\".join(top_labels),\n", " \"pred_material_conf_top1\": float(top_scores[0]),\n", " \"topk_labels\": top_labels,\n", " \"topk_scores\": [float(s) for s in top_scores],\n", " \"per_material_probs\": {k: float(v) for k, v in agg_scores.items()},\n", " \"candidate_materials\": candidate_materials,\n", " }\n", "\n", "\n", "def l2_normalize_np(vec: np.ndarray, eps: float = 1e-12):\n", " n = float(np.linalg.norm(vec))\n", " if n < eps:\n", " return vec\n", " return vec / n\n", "\n", "\n", "def extract_dinov3_embedding(image: Image.Image, dino_model, dino_processor):\n", " inputs = dino_processor(images=image, return_tensors=\"pt\")\n", " inputs = {k: v.to(DEVICE) if hasattr(v, \"to\") else v for k, v in inputs.items()}\n", " with torch.inference_mode():\n", " outputs = dino_model(**inputs)\n", " emb = outputs.pooler_output[0].detach().float().cpu().numpy()\n", " return l2_normalize_np(emb)\n", "\n", "\n", "def build_dino_prototype_bank(material_image_paths: dict, dino_model, dino_processor):\n", " bank = {}\n", " missing = []\n", " for material in ALL_CANONICAL_MATERIALS:\n", " paths = material_image_paths.get(material, [])\n", " embs = []\n", " for p in paths:\n", " if not Path(p).exists():\n", " missing.append(p)\n", " continue\n", " img = load_image(str(p))\n", " embs.append(extract_dinov3_embedding(img, dino_model, dino_processor))\n", " if embs:\n", " bank[material] = l2_normalize_np(np.mean(np.stack(embs, axis=0), axis=0))\n", " return bank, missing\n", "\n", "\n", "def normalize_prob_dict(score_dict: dict, keys: list):\n", " vals = np.array([float(score_dict.get(k, 0.0)) for k in keys], dtype=np.float64)\n", " vals = np.clip(vals, a_min=0.0, a_max=None)\n", " s = float(vals.sum())\n", " if s <= 0:\n", " vals = np.ones_like(vals) / max(len(vals), 1)\n", " else:\n", " vals = vals / s\n", " return {k: float(v) for k, v in zip(keys, vals)}\n", "\n", "\n", "def classify_material_dinov3(\n", " image_views: list,\n", " object_label: str | None,\n", " dino_model,\n", " dino_processor,\n", " prototype_bank: dict,\n", " top_k: int = 3,\n", " view_weights: list | None = None,\n", "):\n", " candidate_materials = resolve_candidate_materials(object_label)\n", "\n", " if view_weights is None:\n", " view_weights = [1.0 / len(image_views)] * len(image_views)\n", " if len(view_weights) != len(image_views):\n", " raise ValueError(\"view_weights length must match image_views length\")\n", "\n", " proto_materials = [m for m in candidate_materials if m in prototype_bank]\n", " prototype_coverage = float(len(proto_materials) / max(len(candidate_materials), 1))\n", "\n", " if len(proto_materials) < 2:\n", " probs = {m: 1.0 / len(candidate_materials) for m in candidate_materials}\n", " ranked = sorted(probs.items(), key=lambda kv: kv[1], reverse=True)\n", " top_items = ranked[: min(top_k, len(ranked))]\n", " return {\n", " \"pred_material_top1\": top_items[0][0],\n", " \"pred_material_top3\": \"|\".join([x[0] for x in top_items]),\n", " \"pred_material_conf_top1\": float(top_items[0][1]),\n", " \"topk_labels\": [x[0] for x in top_items],\n", " \"topk_scores\": [float(x[1]) for x in top_items],\n", " \"per_material_probs\": probs,\n", " \"candidate_materials\": candidate_materials,\n", " \"prototype_coverage\": prototype_coverage,\n", " \"prototypes_available\": False,\n", " }\n", "\n", " agg = {m: 0.0 for m in candidate_materials}\n", " for image, weight in zip(image_views, view_weights):\n", " emb = extract_dinov3_embedding(image, dino_model, dino_processor)\n", " sims = []\n", " for m in candidate_materials:\n", " if m in prototype_bank:\n", " sims.append(float(np.dot(emb, prototype_bank[m])))\n", " else:\n", " sims.append(-1e3)\n", " sims = np.array(sims, dtype=np.float64)\n", " probs = np.exp((sims - np.max(sims)) / max(DINO_SIM_TEMPERATURE, 1e-6))\n", " probs = probs / max(probs.sum(), 1e-12)\n", " for m, p in zip(candidate_materials, probs.tolist()):\n", " agg[m] += float(weight) * float(p)\n", "\n", " agg = normalize_prob_dict(agg, candidate_materials)\n", " ranked = sorted(agg.items(), key=lambda kv: kv[1], reverse=True)\n", " top_items = ranked[: min(top_k, len(ranked))]\n", "\n", " return {\n", " \"pred_material_top1\": top_items[0][0],\n", " \"pred_material_top3\": \"|\".join([x[0] for x in top_items]),\n", " \"pred_material_conf_top1\": float(top_items[0][1]),\n", " \"topk_labels\": [x[0] for x in top_items],\n", " \"topk_scores\": [float(x[1]) for x in top_items],\n", " \"per_material_probs\": agg,\n", " \"candidate_materials\": candidate_materials,\n", " \"prototype_coverage\": prototype_coverage,\n", " \"prototypes_available\": True,\n", " }\n", "\n", "\n", "def fuse_material_predictions(sigclip_out: dict, dino_out: dict):\n", " mats = sorted(set(sigclip_out[\"candidate_materials\"]) | set(dino_out[\"candidate_materials\"]))\n", " p_sig = normalize_prob_dict(sigclip_out.get(\"per_material_probs\", {}), mats)\n", " p_dino = normalize_prob_dict(dino_out.get(\"per_material_probs\", {}), mats)\n", "\n", " w_sig = float(SIGCLIP_FUSION_WEIGHT)\n", " w_dino = float(DINO_FUSION_WEIGHT) if dino_out.get(\"prototypes_available\", False) else 0.0\n", " if w_sig + w_dino <= 0:\n", " w_sig = 1.0\n", " w_sum = w_sig + w_dino\n", " w_sig, w_dino = w_sig / w_sum, w_dino / w_sum\n", "\n", " fused = {m: w_sig * p_sig[m] + w_dino * p_dino[m] for m in mats}\n", " fused = normalize_prob_dict(fused, mats)\n", "\n", " ranked = sorted(fused.items(), key=lambda kv: kv[1], reverse=True)\n", " top3 = ranked[: min(3, len(ranked))]\n", "\n", " sig_top1 = max(p_sig.items(), key=lambda kv: kv[1])[0]\n", " dino_top1 = max(p_dino.items(), key=lambda kv: kv[1])[0]\n", "\n", " return {\n", " \"pred_material_top1\": ranked[0][0],\n", " \"pred_material_top3\": \"|\".join([x[0] for x in top3]),\n", " \"pred_material_conf_top1\": float(ranked[0][1]),\n", " \"topk_labels\": [x[0] for x in top3],\n", " \"topk_scores\": [float(x[1]) for x in top3],\n", " \"per_material_probs\": fused,\n", " \"candidate_materials\": mats,\n", " \"weights\": {\"sigclip\": w_sig, \"dinov3\": w_dino},\n", " \"disagreement\": {\n", " \"top1_match\": bool(sig_top1 == dino_top1),\n", " \"sigclip_top1\": sig_top1,\n", " \"dinov3_top1\": dino_top1,\n", " },\n", " }\n", "\n", "\n", "def build_mixed_materials(prob_dict: dict):\n", " ranked = sorted(prob_dict.items(), key=lambda kv: kv[1], reverse=True)\n", " selected = ranked[: max(2, MAX_MIXED_ENTRIES)]\n", " vals = np.array([max(0.0, float(v)) for _, v in selected], dtype=np.float64)\n", " vals = vals / max(vals.sum(), 1e-12)\n", "\n", " pairs = []\n", " for (m, _), p in zip(selected, vals.tolist()):\n", " pairs.append({\"material\": m, \"percentage\": int(round(100 * p))})\n", "\n", " pairs = [x for x in pairs if x[\"percentage\"] >= MIN_MIXED_PERCENT]\n", " if len(pairs) < 2:\n", " return []\n", "\n", " total = sum(x[\"percentage\"] for x in pairs)\n", " if total <= 0:\n", " return []\n", " for x in pairs:\n", " x[\"percentage\"] = int(round(100 * x[\"percentage\"] / total))\n", " diff = 100 - sum(x[\"percentage\"] for x in pairs)\n", " pairs[-1][\"percentage\"] += diff\n", "\n", " pairs = sorted(pairs, key=lambda x: x[\"percentage\"], reverse=True)[:MAX_MIXED_ENTRIES]\n", " total2 = sum(x[\"percentage\"] for x in pairs)\n", " if total2 != 100:\n", " pairs[-1][\"percentage\"] += 100 - total2\n", " return pairs\n", "\n", "\n", "def suggest_material_decision(fused_out: dict):\n", " probs = fused_out[\"per_material_probs\"]\n", " ranked = sorted(probs.items(), key=lambda kv: kv[1], reverse=True)\n", " top1_mat, top1_conf = ranked[0]\n", " top2_conf = ranked[1][1] if len(ranked) > 1 else 0.0\n", " margin12 = float(top1_conf - top2_conf)\n", "\n", " low_conf = float(top1_conf) < float(LOW_CONF_THRESHOLD)\n", " mixed_suggestion = build_mixed_materials(probs)\n", "\n", " if low_conf:\n", " suggested_mode = \"unknown\"\n", " elif margin12 < float(MIXED_MARGIN_THRESHOLD) and len(mixed_suggestion) >= 2:\n", " suggested_mode = \"mixed\"\n", " else:\n", " suggested_mode = \"single\"\n", "\n", " return {\n", " \"suggested_mode\": suggested_mode,\n", " \"suggested_single_material\": top1_mat if suggested_mode == \"single\" else None,\n", " \"suggested_mixed_materials\": mixed_suggestion if suggested_mode == \"mixed\" else [],\n", " \"low_confidence_flag\": bool(low_conf),\n", " \"requires_user_confirmation\": bool(low_conf),\n", " \"margin12\": margin12,\n", " }\n", "\n", "\n", "def validate_mixed_materials(mixed_materials):\n", " if not isinstance(mixed_materials, list):\n", " return False, \"mixed_materials must be a list\"\n", " if len(mixed_materials) < 2 or len(mixed_materials) > MAX_MIXED_ENTRIES:\n", " return False, f\"mixed_materials must have 2..{MAX_MIXED_ENTRIES} entries\"\n", "\n", " total = 0\n", " for item in mixed_materials:\n", " if not isinstance(item, dict) or \"material\" not in item or \"percentage\" not in item:\n", " return False, \"each mixed entry must be {material, percentage}\"\n", " pct = int(item[\"percentage\"])\n", " if pct < MIN_MIXED_PERCENT:\n", " return False, f\"each mixed percentage must be >= {MIN_MIXED_PERCENT}\"\n", " total += pct\n", " if total != 100:\n", " return False, \"mixed percentages must sum to 100\"\n", " return True, \"\"\n", "\n", "\n", "def encode_final_material_for_csv(final_mode, final_single, final_mixed):\n", " if final_mode == \"single\":\n", " return str(final_single)\n", " if final_mode == \"mixed\":\n", " parts = [f\"{x['material']}={x['percentage']}\" for x in final_mixed]\n", " return \"mixed:\" + \"|\".join(parts)\n", " return \"unknown\"\n", "\n", "\n", "def save_passport(passport: dict, out_dir: Path) -> Path:\n", " out_dir.mkdir(parents=True, exist_ok=True)\n", " sample_id = passport[\"sample_id\"]\n", " out_path = out_dir / f\"{sample_id}.json\"\n", " out_path.write_text(json.dumps(passport, indent=2), encoding=\"utf-8\")\n", " return out_path\n", "\n", "\n", "def append_eval_row(csv_path: Path, row: dict):\n", " df_new = pd.DataFrame([row])\n", " if not csv_path.exists():\n", " df_new.to_csv(csv_path, index=False)\n", " return\n", " df_existing = pd.read_csv(csv_path)\n", " df_out = pd.concat([df_existing, df_new], ignore_index=True)\n", " df_out.to_csv(csv_path, index=False)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "ensure_eval_csv(CSV_PATH, TEMPLATE_PATH)\n", "\n", "print(f\"Using device: {DEVICE}\")\n", "print(\"Running environment checks...\")\n", "\n", "try:\n", " hf_info = whoami()\n", " hf_name = hf_info.get(\"name\", \"\") if isinstance(hf_info, dict) else str(hf_info)\n", " print(f\"HF auth: ok (user={hf_name})\")\n", "except Exception as exc:\n", " token_hint = os.environ.get(\"HF_TOKEN\") or os.environ.get(\"HUGGING_FACE_HUB_TOKEN\")\n", " if token_hint:\n", " print(\"HF auth: token env var found, but whoami failed. You may need `hf auth login`.\")\n", " else:\n", " print(\"HF auth: not detected. Run `hf auth login` to access gated SAM3.1 checkpoints.\")\n", " print(f\"HF auth check detail: {exc}\")\n", "\n", "if SAM3_CHECKPOINT_PATH is not None:\n", " if not Path(SAM3_CHECKPOINT_PATH).exists():\n", " raise FileNotFoundError(f\"SAM3_CHECKPOINT_PATH does not exist: {SAM3_CHECKPOINT_PATH}\")\n", " print(f\"SAM3 checkpoint source: local ({SAM3_CHECKPOINT_PATH})\")\n", "else:\n", " print(f\"SAM3 checkpoint source: Hugging Face gated repo ({SAM3_REPO_ID})\")\n", "\n", "if DEVICE != \"cuda\":\n", " print(\"Warning: CUDA not available; SAM3.1 full checkpoint may be too slow on CPU.\")\n", "\n", "print(\"Loading SAM 3.1 model...\")\n", "if SAM3_CHECKPOINT_PATH is None:\n", " sam3_checkpoint_path = download_ckpt_from_hf(version=\"sam3.1\")\n", "else:\n", " sam3_checkpoint_path = SAM3_CHECKPOINT_PATH\n", "\n", "sam_model = build_sam3_image_model(\n", " device=DEVICE,\n", " checkpoint_path=sam3_checkpoint_path,\n", " load_from_HF=False,\n", " enable_inst_interactivity=False,\n", ")\n", "sam_model = sam_model.to(dtype=SAM3_TORCH_DTYPE)\n", "sam_model.eval()\n", "sam_processor = Sam3Processor(\n", " model=sam_model,\n", " device=DEVICE,\n", " confidence_threshold=SAM3_CONFIDENCE_THRESHOLD,\n", ")\n", "print(f\"SAM backend: SAM3.1 checkpoint at {sam3_checkpoint_path}\")\n", "print(f\"SAM dtype: {SAM3_TORCH_DTYPE}\")\n", "\n", "print(\"Loading SigCLIP model...\")\n", "sigclip_processor = AutoProcessor.from_pretrained(SIGCLIP_MODEL_ID)\n", "sigclip_model = AutoModel.from_pretrained(SIGCLIP_MODEL_ID).to(DEVICE)\n", "sigclip_model.eval()\n", "\n", "print(\"Loading DINOv3 model...\")\n", "dino_processor = AutoImageProcessor.from_pretrained(DINO_MODEL_ID)\n", "dino_model = AutoModel.from_pretrained(DINO_MODEL_ID).to(DEVICE)\n", "dino_model.eval()\n", "\n", "print(\"Building DINO prototype bank...\")\n", "dino_prototype_bank, missing_proto_paths = build_dino_prototype_bank(\n", " MATERIAL_PROTOTYPE_IMAGE_PATHS,\n", " dino_model=dino_model,\n", " dino_processor=dino_processor,\n", ")\n", "print(f\"DINO prototypes loaded: {len(dino_prototype_bank)} materials\")\n", "if missing_proto_paths:\n", " print(\"Missing prototype files:\")\n", " for p in missing_proto_paths:\n", " print(f\"- {p}\")\n", "\n", "print(\"Models ready.\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Configure one image and one-or-many objects in it.\n", "image_id = \"IMG002\"\n", "image_path = \"/content/images/example.jpg\"\n", "\n", "SELECTION_MODE = \"box\" # \"box\" or \"brush\"\n", "USE_ACCEPT_REDRAW = True\n", "MAX_SELECTION_ATTEMPTS = 5\n", "MULTI_OBJECT_MODE = True\n", "N_OBJECTS_TO_SELECT = 3 # used only in multi-object mode\n", "\n", "# Define object metadata in selection order. prompt_bbox_xyxy is optional in interactive mode.\n", "OBJECT_SPECS = [\n", " {\n", " \"sample_id\": \"S002\",\n", " \"ground_truth_object\": \"window\",\n", " \"ground_truth_material\": \"glass\",\n", " \"pred_object\": \"window\",\n", " },\n", " {\n", " \"sample_id\": \"S003\",\n", " \"ground_truth_object\": \"door\",\n", " \"ground_truth_material\": \"wood\",\n", " \"pred_object\": \"door\",\n", " },\n", " {\n", " \"sample_id\": \"S004\",\n", " \"ground_truth_object\": \"wall\",\n", " \"ground_truth_material\": \"brick\",\n", " \"pred_object\": \"wall\",\n", " },\n", "]\n", "\n", "img = load_image(image_path)\n", "MANUAL_SELECTION = False\n", "\n", "if SELECTION_MODE not in {\"box\", \"brush\"}:\n", " raise ValueError(\"SELECTION_MODE must be 'box' or 'brush'\")\n", "\n", "if MULTI_OBJECT_MODE and N_OBJECTS_TO_SELECT < 1:\n", " raise ValueError(\"N_OBJECTS_TO_SELECT must be >= 1 in multi-object mode\")\n", "\n", "if MANUAL_SELECTION:\n", " for obj in OBJECT_SPECS:\n", " if \"prompt_bbox_xyxy\" not in obj:\n", " raise ValueError(\"Manual mode requires prompt_bbox_xyxy for each object.\")\n", " obj.setdefault(\"prompt_points_xy\", [])\n", " OBJECT_RECORDS = OBJECT_SPECS if MULTI_OBJECT_MODE else [OBJECT_SPECS[0]]\n", "elif SELECTION_MODE == \"box\":\n", " if MULTI_OBJECT_MODE:\n", " if USE_ACCEPT_REDRAW:\n", " boxes = select_multiple_bboxes_with_redraw_colab(\n", " img,\n", " expected_boxes=N_OBJECTS_TO_SELECT,\n", " max_attempts_per_box=MAX_SELECTION_ATTEMPTS,\n", " )\n", " else:\n", " boxes = select_multiple_bboxes_colab(img, expected_boxes=N_OBJECTS_TO_SELECT)\n", " if len(boxes) > len(OBJECT_SPECS):\n", " raise ValueError(\"Add more OBJECT_SPECS entries to match drawn boxes.\")\n", " for i, box in enumerate(boxes):\n", " OBJECT_SPECS[i][\"prompt_bbox_xyxy\"] = [int(v) for v in box]\n", " OBJECT_SPECS[i][\"prompt_points_xy\"] = []\n", " OBJECT_RECORDS = OBJECT_SPECS[: len(boxes)]\n", " else:\n", " if USE_ACCEPT_REDRAW:\n", " box = select_bbox_with_redraw_colab(\n", " img,\n", " prompt=\"Draw box around target object\",\n", " max_attempts=MAX_SELECTION_ATTEMPTS,\n", " )\n", " else:\n", " box = select_bbox_colab(img, prompt=\"Draw box around target object\")\n", " OBJECT_SPECS[0][\"prompt_bbox_xyxy\"] = [int(v) for v in box]\n", " OBJECT_SPECS[0][\"prompt_points_xy\"] = []\n", " OBJECT_RECORDS = [OBJECT_SPECS[0]]\n", "else:\n", " if MULTI_OBJECT_MODE:\n", " if len(OBJECT_SPECS) < N_OBJECTS_TO_SELECT:\n", " raise ValueError(\"Add enough OBJECT_SPECS entries for brush selection.\")\n", " OBJECT_RECORDS = OBJECT_SPECS[:N_OBJECTS_TO_SELECT]\n", " for i, obj in enumerate(OBJECT_RECORDS):\n", " if USE_ACCEPT_REDRAW:\n", " hint = select_brush_hint_with_redraw_colab(\n", " img,\n", " prompt=f\"Brush object {i + 1}/{N_OBJECTS_TO_SELECT}\",\n", " max_attempts=MAX_SELECTION_ATTEMPTS,\n", " )\n", " else:\n", " hint = select_brush_hint_colab(\n", " img,\n", " prompt=f\"Brush object {i + 1}/{N_OBJECTS_TO_SELECT}\",\n", " )\n", " obj[\"prompt_bbox_xyxy\"] = [int(v) for v in hint[\"prompt_bbox_xyxy\"]]\n", " obj[\"prompt_points_xy\"] = [[int(x), int(y)] for x, y in hint[\"prompt_points_xy\"]]\n", " else:\n", " if USE_ACCEPT_REDRAW:\n", " hint = select_brush_hint_with_redraw_colab(\n", " img,\n", " prompt=\"Brush over target object\",\n", " max_attempts=MAX_SELECTION_ATTEMPTS,\n", " )\n", " else:\n", " hint = select_brush_hint_colab(img, prompt=\"Brush over target object\")\n", " OBJECT_SPECS[0][\"prompt_bbox_xyxy\"] = [int(v) for v in hint[\"prompt_bbox_xyxy\"]]\n", " OBJECT_SPECS[0][\"prompt_points_xy\"] = [[int(x), int(y)] for x, y in hint[\"prompt_points_xy\"]]\n", " OBJECT_RECORDS = [OBJECT_SPECS[0]]\n", "\n", "for obj in OBJECT_RECORDS:\n", " obj.setdefault(\"pred_object\", obj.get(\"ground_truth_object\", \"\"))\n", "\n", "print(f\"Prepared {len(OBJECT_RECORDS)} object(s) for image {image_id} (mode={SELECTION_MODE}).\")\n", "\n", "plt.figure(figsize=(9, 7))\n", "plt.imshow(img)\n", "ax = plt.gca()\n", "for idx, obj in enumerate(OBJECT_RECORDS, start=1):\n", " x0, y0, x1, y1 = [int(v) for v in obj[\"prompt_bbox_xyxy\"]]\n", " rect = patches.Rectangle((x0, y0), x1 - x0, y1 - y0, linewidth=2, edgecolor=\"red\", facecolor=\"none\")\n", " ax.add_patch(rect)\n", " for p in obj.get(\"prompt_points_xy\", []):\n", " ax.scatter([p[0]], [p[1]], c=\"cyan\", s=12)\n", " ax.text(x0 + 4, y0 + 14, f\"{idx}:{obj['sample_id']}\", color=\"yellow\", fontsize=9, bbox={\"facecolor\": \"black\", \"alpha\": 0.4, \"pad\": 2})\n", "plt.title(\"Prompt preview (bbox + optional brush points)\")\n", "plt.axis(\"off\")\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "RUN_ROWS = []\n", "\n", "for obj in OBJECT_RECORDS:\n", " sample_id = obj[\"sample_id\"]\n", " prompt_bbox_xyxy = [int(v) for v in obj[\"prompt_bbox_xyxy\"]]\n", " prompt_points_xy = [[int(p[0]), int(p[1])] for p in obj.get(\"prompt_points_xy\", [])]\n", " tap_x, tap_y = bbox_xyxy_to_center_xy(prompt_bbox_xyxy)\n", " pred_object = obj.get(\"pred_object\", obj.get(\"ground_truth_object\", \"\"))\n", "\n", " t0 = pd.Timestamp.utcnow()\n", "\n", " seg_start = pd.Timestamp.utcnow()\n", " seg = run_sam_guided_segmentation(\n", " image=img,\n", " prompt_bbox_xyxy=prompt_bbox_xyxy,\n", " prompt_points_xy=prompt_points_xy,\n", " sam_model=sam_model,\n", " sam_processor=sam_processor,\n", " )\n", " seg_end = pd.Timestamp.utcnow()\n", "\n", " crop_masked = crop_masked_rgb(img, seg[\"mask\"], seg[\"bbox_xywh\"], pad=8)\n", " crop_bbox = crop_bbox_rgb(img, seg[\"bbox_xywh\"], pad=16)\n", "\n", " clf_start = pd.Timestamp.utcnow()\n", " clf_sig = classify_material_sigclip(\n", " image_views=[crop_masked, crop_bbox, img],\n", " object_label=pred_object,\n", " sigclip_model=sigclip_model,\n", " sigclip_processor=sigclip_processor,\n", " top_k=3,\n", " view_weights=[0.5, 0.35, 0.15],\n", " )\n", " clf_dino = classify_material_dinov3(\n", " image_views=[crop_masked, crop_bbox],\n", " object_label=pred_object,\n", " dino_model=dino_model,\n", " dino_processor=dino_processor,\n", " prototype_bank=dino_prototype_bank,\n", " top_k=3,\n", " view_weights=[0.65, 0.35],\n", " )\n", " clf_fused = fuse_material_predictions(clf_sig, clf_dino)\n", " decision = suggest_material_decision(clf_fused)\n", "\n", " clf = {\n", " **clf_fused,\n", " \"sigclip\": clf_sig,\n", " \"dinov3\": clf_dino,\n", " \"suggested_mode\": decision[\"suggested_mode\"],\n", " \"suggested_single_material\": decision[\"suggested_single_material\"],\n", " \"suggested_mixed_materials\": decision[\"suggested_mixed_materials\"],\n", " \"low_confidence_flag\": decision[\"low_confidence_flag\"],\n", " \"requires_user_confirmation\": decision[\"requires_user_confirmation\"],\n", " \"margin12\": decision[\"margin12\"],\n", " }\n", " clf_end = pd.Timestamp.utcnow()\n", "\n", " t1 = pd.Timestamp.utcnow()\n", "\n", " segmentation_time_sec = (seg_end - seg_start).total_seconds()\n", " classification_time_sec = (clf_end - clf_start).total_seconds()\n", " end_to_end_time_sec = (t1 - t0).total_seconds()\n", "\n", " RUN_ROWS.append({\n", " \"sample_id\": sample_id,\n", " \"image_id\": image_id,\n", " \"image_path\": image_path,\n", " \"prompt_bbox_xyxy\": prompt_bbox_xyxy,\n", " \"prompt_points_xy\": prompt_points_xy,\n", " \"tap_x\": tap_x,\n", " \"tap_y\": tap_y,\n", " \"ground_truth_object\": obj.get(\"ground_truth_object\", \"\"),\n", " \"ground_truth_material\": obj.get(\"ground_truth_material\", \"\"),\n", " \"pred_object\": pred_object,\n", " \"seg\": seg,\n", " \"clf\": clf,\n", " \"crop_masked\": crop_masked,\n", " \"timing\": {\n", " \"segmentation_time_sec\": segmentation_time_sec,\n", " \"classification_time_sec\": classification_time_sec,\n", " \"end_to_end_time_sec\": end_to_end_time_sec,\n", " },\n", " })\n", "\n", " print(\n", " f\"[{sample_id}] bbox={seg['bbox_xywh']} area={seg['mask_area_px']} \"\n", " f\"sam={seg['sam_score']:.3f} sel={seg['selection_score']:.3f} \"\n", " f\"hint_iou={seg['selection_metrics']['hint_iou']:.3f} \"\n", " f\"pt_cov={seg['selection_metrics']['point_coverage']:.3f} \"\n", " f\"top3={clf['pred_material_top3']} mode={clf['suggested_mode']} \"\n", " f\"w=({clf['weights']['sigclip']:.2f},{clf['weights']['dinov3']:.2f})\"\n", " )\n", "\n", " fig, axes = plt.subplots(1, 2, figsize=(12, 5))\n", " axes[0].imshow(img)\n", " axes[0].imshow(seg[\"mask\"], alpha=0.35, cmap=\"spring\")\n", " px0, py0, px1, py1 = prompt_bbox_xyxy\n", " prompt_rect = patches.Rectangle((px0, py0), px1 - px0, py1 - py0, linewidth=2, edgecolor=\"cyan\", facecolor=\"none\")\n", " axes[0].add_patch(prompt_rect)\n", " if prompt_points_xy:\n", " xs = [p[0] for p in prompt_points_xy]\n", " ys = [p[1] for p in prompt_points_xy]\n", " axes[0].scatter(xs, ys, c=\"lime\", s=10)\n", " axes[0].set_title(f\"Segmentation Overlay ({sample_id})\")\n", " axes[0].axis(\"off\")\n", "\n", " axes[1].imshow(crop_masked)\n", " axes[1].set_title(f\"Masked Crop ({sample_id})\")\n", " axes[1].axis(\"off\")\n", " plt.tight_layout()\n", " plt.show()\n", "\n", "print(f\"Processed {len(RUN_ROWS)} object(s).\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Human review/edit step. Set defaults and optional per-sample overrides.\n", "DEFAULT_SEGMENTATION_ACCEPTED = 1\n", "DEFAULT_STATUS = \"ok\"\n", "DEFAULT_ERROR_TYPE = \"\"\n", "\n", "# Example override format:\n", "# USER_REVIEW_OVERRIDES = {\n", "# \"S003\": {\n", "# \"segmentation_accepted\": 1,\n", "# \"final_material_mode\": \"mixed\", # single|mixed|unknown\n", "# \"final_single_material\": None,\n", "# \"final_mixed_materials\": [{\"material\": \"wood\", \"percentage\": 60}, {\"material\": \"steel\", \"percentage\": 40}],\n", "# \"confirmed_by_user\": True,\n", "# \"edit_reason\": \"edge case\",\n", "# \"status\": \"ok\",\n", "# \"error_type\": \"\",\n", "# }\n", "# }\n", "USER_REVIEW_OVERRIDES = {}\n", "\n", "for run in RUN_ROWS:\n", " sid = run[\"sample_id\"]\n", " clf = run[\"clf\"]\n", " override = USER_REVIEW_OVERRIDES.get(sid, {})\n", "\n", " suggested_mode = clf[\"suggested_mode\"]\n", " suggested_single = clf.get(\"suggested_single_material\")\n", " suggested_mixed = clf.get(\"suggested_mixed_materials\", [])\n", "\n", " final_mode = override.get(\"final_material_mode\", suggested_mode)\n", " final_single = override.get(\"final_single_material\", suggested_single if final_mode == \"single\" else None)\n", " final_mixed = override.get(\"final_mixed_materials\", suggested_mixed if final_mode == \"mixed\" else [])\n", "\n", " requires_confirm = bool(clf.get(\"requires_user_confirmation\", False))\n", " confirmed_by_user = bool(override.get(\"confirmed_by_user\", not requires_confirm))\n", "\n", " if final_mode == \"mixed\":\n", " ok, err = validate_mixed_materials(final_mixed)\n", " if not ok:\n", " raise ValueError(f\"{sid}: invalid final_mixed_materials -> {err}\")\n", "\n", " user_edited = int(final_mode != suggested_mode)\n", "\n", " run[\"review\"] = {\n", " \"segmentation_accepted\": int(override.get(\"segmentation_accepted\", DEFAULT_SEGMENTATION_ACCEPTED)),\n", " \"final_material_mode\": final_mode,\n", " \"final_single_material\": final_single,\n", " \"final_mixed_materials\": final_mixed,\n", " \"requires_user_confirmation\": requires_confirm,\n", " \"confirmed_by_user\": confirmed_by_user,\n", " \"user_edited_material\": user_edited,\n", " \"edit_reason\": override.get(\"edit_reason\", \"\"),\n", " \"status\": override.get(\"status\", DEFAULT_STATUS),\n", " \"error_type\": override.get(\"error_type\", DEFAULT_ERROR_TYPE),\n", " }\n", "\n", "pd.DataFrame([\n", " {\n", " \"sample_id\": r[\"sample_id\"],\n", " \"suggested_mode\": r[\"clf\"][\"suggested_mode\"],\n", " \"final_mode\": r[\"review\"][\"final_material_mode\"],\n", " \"requires_confirm\": r[\"review\"][\"requires_user_confirmation\"],\n", " \"confirmed\": r[\"review\"][\"confirmed_by_user\"],\n", " \"user_edited_material\": r[\"review\"][\"user_edited_material\"],\n", " \"status\": r[\"review\"][\"status\"],\n", " }\n", " for r in RUN_ROWS\n", "])" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "for run in RUN_ROWS:\n", " review = run.get(\"review\", {})\n", " seg = run[\"seg\"]\n", " clf = run[\"clf\"]\n", " timing = run[\"timing\"]\n", "\n", " if review.get(\"requires_user_confirmation\", False) and not review.get(\"confirmed_by_user\", False):\n", " raise ValueError(f\"{run['sample_id']}: low-confidence result requires explicit user confirmation before save\")\n", "\n", " final_mode = review.get(\"final_material_mode\", \"unknown\")\n", " final_single = review.get(\"final_single_material\", None)\n", " final_mixed = review.get(\"final_mixed_materials\", [])\n", "\n", " if final_mode == \"mixed\":\n", " ok, err = validate_mixed_materials(final_mixed)\n", " if not ok:\n", " raise ValueError(f\"{run['sample_id']}: invalid mixed output -> {err}\")\n", " elif final_mode == \"single\":\n", " if not final_single:\n", " raise ValueError(f\"{run['sample_id']}: final_single_material is required for single mode\")\n", " elif final_mode == \"unknown\":\n", " final_single = None\n", " final_mixed = []\n", " else:\n", " raise ValueError(f\"{run['sample_id']}: unknown final_material_mode={final_mode}\")\n", "\n", " passport = {\n", " \"sample_id\": run[\"sample_id\"],\n", " \"image_id\": run[\"image_id\"],\n", " \"image_path\": run[\"image_path\"],\n", " \"interaction\": {\n", " \"selection_mode\": SELECTION_MODE,\n", " \"prompt_bbox_xyxy\": [int(v) for v in run[\"prompt_bbox_xyxy\"]],\n", " \"prompt_points_xy\": [[int(p[0]), int(p[1])] for p in run.get(\"prompt_points_xy\", [])],\n", " \"bbox_center_xy\": [int(run[\"tap_x\"]), int(run[\"tap_y\"])],\n", " },\n", " \"object\": {\n", " \"ground_truth\": run[\"ground_truth_object\"],\n", " \"predicted\": run[\"pred_object\"],\n", " },\n", " \"material\": {\n", " \"ground_truth\": run[\"ground_truth_material\"],\n", " \"predicted_topk\": [\n", " {\"material\": m, \"score\": float(s)}\n", " for m, s in zip(clf[\"topk_labels\"], clf[\"topk_scores\"])\n", " ],\n", " \"predicted_confidence_top1\": clf[\"pred_material_conf_top1\"],\n", " \"suggested_mode\": clf[\"suggested_mode\"],\n", " \"suggested_single_material\": clf.get(\"suggested_single_material\"),\n", " \"suggested_mixed_materials\": clf.get(\"suggested_mixed_materials\", []),\n", " \"final_material_mode\": final_mode,\n", " \"final_single_material\": final_single,\n", " \"final_mixed_materials\": final_mixed,\n", " \"low_confidence_flag\": bool(clf.get(\"low_confidence_flag\", False)),\n", " \"requires_user_confirmation\": bool(review.get(\"requires_user_confirmation\", False)),\n", " \"confirmed_by_user\": bool(review.get(\"confirmed_by_user\", False)),\n", " \"fusion_weights\": clf.get(\"weights\", {}),\n", " \"model_outputs\": {\n", " \"sigclip\": {\n", " \"topk_labels\": clf[\"sigclip\"].get(\"topk_labels\", []),\n", " \"topk_scores\": clf[\"sigclip\"].get(\"topk_scores\", []),\n", " },\n", " \"dinov3\": {\n", " \"topk_labels\": clf[\"dinov3\"].get(\"topk_labels\", []),\n", " \"topk_scores\": clf[\"dinov3\"].get(\"topk_scores\", []),\n", " \"prototype_coverage\": clf[\"dinov3\"].get(\"prototype_coverage\", 0.0),\n", " \"prototypes_available\": bool(clf[\"dinov3\"].get(\"prototypes_available\", False)),\n", " },\n", " },\n", " \"user_final\": encode_final_material_for_csv(final_mode, final_single, final_mixed),\n", " \"user_edited\": bool(review.get(\"user_edited_material\", 0)),\n", " \"edit_reason\": review.get(\"edit_reason\", \"\"),\n", " },\n", " \"segmentation\": {\n", " \"accepted\": bool(review.get(\"segmentation_accepted\", 1)),\n", " \"mask_area_px\": int(seg[\"mask_area_px\"]),\n", " \"bbox_xywh\": [int(v) for v in seg[\"bbox_xywh\"]],\n", " \"sam_score\": seg[\"sam_score\"],\n", " \"selection_score\": seg.get(\"selection_score\", seg[\"sam_score\"]),\n", " \"selection_metrics\": seg.get(\"selection_metrics\", {}),\n", " },\n", " \"timing\": {\n", " \"end_to_end_time_sec\": timing[\"end_to_end_time_sec\"],\n", " \"segmentation_time_sec\": timing[\"segmentation_time_sec\"],\n", " \"classification_time_sec\": timing[\"classification_time_sec\"],\n", " },\n", " \"status\": review.get(\"status\", \"ok\"),\n", " \"error_type\": review.get(\"error_type\", \"\"),\n", " \"model_name\": \"SigCLIP+DINOv3\",\n", " \"created_at_utc\": datetime.now(timezone.utc).isoformat(),\n", " }\n", "\n", " passport_path = save_passport(passport, PASSPORT_DIR)\n", "\n", " eval_row = {\n", " \"sample_id\": run[\"sample_id\"],\n", " \"image_id\": run[\"image_id\"],\n", " \"image_path\": run[\"image_path\"],\n", " \"tap_x\": run[\"tap_x\"],\n", " \"tap_y\": run[\"tap_y\"],\n", " \"ground_truth_object\": run[\"ground_truth_object\"],\n", " \"ground_truth_material\": run[\"ground_truth_material\"],\n", " \"pred_object\": run[\"pred_object\"],\n", " \"pred_material_top1\": clf[\"pred_material_top1\"],\n", " \"pred_material_top3\": clf[\"pred_material_top3\"],\n", " \"pred_material_conf_top1\": clf[\"pred_material_conf_top1\"],\n", " \"model_name\": \"SigCLIP+DINOv3\",\n", " \"segmentation_accepted\": review.get(\"segmentation_accepted\", 1),\n", " \"user_edited_material\": review.get(\"user_edited_material\", 0),\n", " \"user_final_material\": encode_final_material_for_csv(final_mode, final_single, final_mixed),\n", " \"edit_reason\": review.get(\"edit_reason\", \"\"),\n", " \"end_to_end_time_sec\": timing[\"end_to_end_time_sec\"],\n", " \"segmentation_time_sec\": timing[\"segmentation_time_sec\"],\n", " \"classification_time_sec\": timing[\"classification_time_sec\"],\n", " \"status\": review.get(\"status\", \"ok\"),\n", " \"error_type\": review.get(\"error_type\", \"\"),\n", " \"mask_area_px\": seg[\"mask_area_px\"],\n", " \"bbox_xywh\": \",\".join(str(v) for v in seg[\"bbox_xywh\"]),\n", " \"passport_json_path\": str(passport_path),\n", " }\n", " append_eval_row(CSV_PATH, eval_row)\n", " print(f\"[{run['sample_id']}] saved passport + eval row\")\n", "\n", "print(f\"Updated eval CSV: {CSV_PATH}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def compute_metrics(df_in: pd.DataFrame) -> dict:\n", " df = df_in.copy()\n", " df = df[df[\"status\"].notna()]\n", "\n", " if df.empty:\n", " raise ValueError(\"No rows with non-empty status.\")\n", "\n", " df[\"top1_correct\"] = (df[\"pred_material_top1\"] == df[\"ground_truth_material\"]).astype(int)\n", " df[\"top3_hit\"] = df.apply(\n", " lambda r: int(str(r[\"ground_truth_material\"]) in str(r[\"pred_material_top3\"]).split(\"|\")),\n", " axis=1,\n", " )\n", "\n", " metrics = {\n", " \"n_samples\": int(len(df)),\n", " \"segmentation_acceptance_rate\": float(df[\"segmentation_accepted\"].mean()),\n", " \"top1_accuracy\": float(df[\"top1_correct\"].mean()),\n", " \"top3_hit_rate\": float(df[\"top3_hit\"].mean()),\n", " \"manual_edit_rate\": float(df[\"user_edited_material\"].mean()),\n", " \"median_e2e_time_sec\": float(df[\"end_to_end_time_sec\"].median()),\n", " \"failure_rate\": float((df[\"status\"] != \"ok\").mean()),\n", " }\n", "\n", " go_no_go = (\n", " metrics[\"segmentation_acceptance_rate\"] >= THRESHOLDS[\"segmentation_acceptance_rate\"]\n", " and metrics[\"top1_accuracy\"] >= THRESHOLDS[\"top1_accuracy\"]\n", " and metrics[\"top3_hit_rate\"] >= THRESHOLDS[\"top3_hit_rate\"]\n", " and metrics[\"manual_edit_rate\"] <= THRESHOLDS[\"manual_edit_rate_max\"]\n", " and metrics[\"median_e2e_time_sec\"] <= THRESHOLDS[\"median_e2e_time_sec_max\"]\n", " and metrics[\"failure_rate\"] == 0.0\n", " )\n", "\n", " return {\"metrics\": metrics, \"go_no_go\": bool(go_no_go)}\n", "\n", "\n", "df_eval = pd.read_csv(CSV_PATH)\n", "result = compute_metrics(df_eval)\n", "print(\"Thresholds:\", THRESHOLDS)\n", "for k, v in result[\"metrics\"].items():\n", " print(f\"- {k}: {v}\")\n", "print(f\"Go/No-Go: {'GO' if result['go_no_go'] else 'NO-GO'}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "df_eval = pd.read_csv(CSV_PATH)\n", "if df_eval.empty:\n", " print(\"No evaluation rows yet.\")\n", "else:\n", " for name in sorted(df_eval[\"model_name\"].dropna().unique().tolist()):\n", " sub = df_eval[df_eval[\"model_name\"] == name]\n", " print(f\"Model={name}\")\n", " print(compute_metrics(sub)[\"metrics\"])" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3" } }, "nbformat": 4, "nbformat_minor": 5 }