updated with SAM3.1 and env checks

1 files changed, 109 insertions, 42 deletions

diff --git a/poc_eval_notebook.ipynb b/poc_eval_notebook.ipynb
index 8ea57e4..17f0de5 100644
--- a/poc_eval_notebook.ipynb
+++ b/poc_eval_notebook.ipynb
@@ -7,14 +7,14 @@
         "# Material Hunters POC Notebook\n",
         "\n",
         "This notebook runs a minimal end-to-end baseline:\n",
-        "- box-prompt segmentation with SAM (`facebook/sam-vit-base`)\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",
-        "When SAM2/SAM3 integration is ready, replace the SAM helper function in this notebook."
+        "Requires access to the gated SAM 3.1 checkpoint on Hugging Face."
       ]
     },
     {
@@ -24,7 +24,8 @@
       "outputs": [],
       "source": [
         "# Run once in Colab if needed.\n",
-        "%pip -q install transformers accelerate pillow matplotlib pandas"
+        "%pip -q install transformers accelerate pillow matplotlib pandas huggingface_hub\n",
+        "%pip -q install git+https://github.com/facebookresearch/sam3.git"
       ]
     },
     {
@@ -34,6 +35,7 @@
       "outputs": [],
       "source": [
         "from pathlib import Path\n",
+        "import os\n",
         "import base64\n",
         "import io\n",
         "from datetime import datetime, timezone\n",
@@ -46,15 +48,18 @@
         "import matplotlib.patches as patches\n",
         "from PIL import Image\n",
         "\n",
-        "from transformers import (\n",
-        "    SamModel,\n",
-        "    SamProcessor,\n",
-        "    AutoModel,\n",
-        "    AutoProcessor,\n",
-        ")\n",
+        "from transformers import 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",
-        "SAM_MODEL_ID = \"facebook/sam-vit-base\"\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",
         "SIGCLIP_MODEL_ID = \"google/siglip-base-patch16-224\"\n",
         "\n",
         "CSV_PATH = Path(\"/content/poc_eval_sheet.csv\")\n",
@@ -661,6 +666,21 @@
         "    return ((x0 + x1) // 2, (y0 + y1) // 2)\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",
@@ -745,37 +765,46 @@
         "\n",
         "\n",
         "def run_sam_guided_segmentation(image: Image.Image, prompt_bbox_xyxy, prompt_points_xy, sam_model, sam_processor):\n",
-        "    x0, y0, x1, y1 = [float(v) for v in prompt_bbox_xyxy]\n",
-        "    if prompt_points_xy:\n",
-        "        pts = [[[float(x), float(y)] for x, y in prompt_points_xy]]\n",
-        "        labels = [[1 for _ in prompt_points_xy]]\n",
-        "        inputs = sam_processor(\n",
-        "            image,\n",
-        "            input_boxes=[[[x0, y0, x1, y1]]],\n",
-        "            input_points=pts,\n",
-        "            input_labels=labels,\n",
-        "            return_tensors=\"pt\",\n",
-        "        )\n",
-        "    else:\n",
-        "        inputs = sam_processor(\n",
-        "            image,\n",
-        "            input_boxes=[[[x0, y0, x1, y1]]],\n",
-        "            return_tensors=\"pt\",\n",
-        "        )\n",
+        "    img_w, img_h = image.size\n",
         "\n",
-        "    inputs = {k: v.to(DEVICE) if hasattr(v, \"to\") else v for k, v in inputs.items()}\n",
+        "    state = sam_processor.set_image(image, state={})\n",
         "\n",
-        "    with torch.no_grad():\n",
-        "        outputs = sam_model(**inputs)\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",
-        "    masks = sam_processor.image_processor.post_process_masks(\n",
-        "        outputs.pred_masks.cpu(),\n",
-        "        inputs[\"original_sizes\"].cpu(),\n",
-        "        inputs[\"reshaped_input_sizes\"].cpu(),\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",
-        "    sam_scores = outputs.iou_scores[0, 0].detach().cpu().numpy()\n",
-        "    candidate_masks = masks[0][0].numpy().astype(bool)\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",
@@ -786,7 +815,7 @@
         "    for idx in range(candidate_masks.shape[0]):\n",
         "        metrics = score_sam_candidate(\n",
         "            mask=candidate_masks[idx],\n",
-        "            sam_score=float(sam_scores[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",
@@ -911,10 +940,48 @@
         "ensure_eval_csv(CSV_PATH, TEMPLATE_PATH)\n",
         "\n",
         "print(f\"Using device: {DEVICE}\")\n",
-        "print(\"Loading SAM model...\")\n",
-        "sam_processor = SamProcessor.from_pretrained(SAM_MODEL_ID)\n",
-        "sam_model = SamModel.from_pretrained(SAM_MODEL_ID).to(DEVICE)\n",
-        "sam_model.eval()\n",
+        "print(\"Running environment checks...\")\n",
+        "\n",
+        "try:\n",
+        "    hf_info = whoami()\n",
+        "    hf_name = hf_info.get(\"name\", \"<unknown>\") 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_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",
         "\n",
         "print(\"Loading SigCLIP model...\")\n",
         "sigclip_processor = AutoProcessor.from_pretrained(SIGCLIP_MODEL_ID)\n",