diff options
| -rw-r--r-- | poc_eval_notebook.ipynb | 48 |
1 files changed, 30 insertions, 18 deletions
diff --git a/poc_eval_notebook.ipynb b/poc_eval_notebook.ipynb index 17f0de5..55f337f 100644 --- a/poc_eval_notebook.ipynb +++ b/poc_eval_notebook.ipynb @@ -35,6 +35,7 @@ "outputs": [], "source": [ "from pathlib import Path\n", + "from contextlib import nullcontext\n", "import os\n", "import base64\n", "import io\n", @@ -60,6 +61,7 @@ "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", "SIGCLIP_MODEL_ID = \"google/siglip-base-patch16-224\"\n", "\n", "CSV_PATH = Path(\"/content/poc_eval_sheet.csv\")\n", @@ -666,6 +668,12 @@ " 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", @@ -767,23 +775,24 @@ "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", - " 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", + " 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", @@ -832,7 +841,7 @@ " \"mask\": mask,\n", " \"mask_area_px\": int(mask.sum()),\n", " \"bbox_xywh\": bbox,\n", - " \"prompt_bbox_xyxy\": [int(x0), int(y0), int(x1), int(y1)],\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", @@ -976,12 +985,15 @@ " 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", |