Remove mask output; keep heatmap overlay only

app.py

@@ -93,27 +93,10 @@ def infer(ref_files, test_file, use_sam3, sam_prompt):
     model.fit(ref_paths, n_ref=len(ref_paths))
     image, amap, score = model.predict(test_path)
 
-    # Default empty mask for UI (in case SAM3 is off or fails)
-    fg_mask_img = Image.fromarray(np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8))
-    if segmenter is not None:
-        # Compute full-resolution mask for visualization and for suppressing background
-        fg_mask = segmenter.get_object_mask(image)
-        # Resize mask to anomaly map resolution
-        fg_mask_resized = fg_mask
-        if fg_mask.shape != amap.shape:
-            fg_mask_resized = cv2.resize(
-                fg_mask.astype(np.float32),
-                (amap.shape[1], amap.shape[0]),
-                interpolation=cv2.INTER_NEAREST,
-            ) > 0.5
-        # Apply mask to anomaly map to suppress background leakage
-        amap = amap * fg_mask_resized
-        fg_mask_img = Image.fromarray((fg_mask.astype(np.uint8) * 255))
-
     overlay = _make_overlay(image, amap)
 
     status_text = "\n".join(status_lines) if status_lines else f"Ran on {device}."
-    return overlay, fg_mask_img, float(score), status_text
+    return overlay, float(score), status_text
 
 
 def build_demo():
@@ -138,14 +121,13 @@ def build_demo():
 
         run_btn = gr.Button("Run", variant="primary")
         overlay_out = gr.Image(label="Heatmap overlay", type="pil")
-        mask_out = gr.Image(label="Foreground mask (if SAM3)", type="pil", image_mode="L")
         score_out = gr.Number(label="Image anomaly score (mean top 1% distance)")
         status_out = gr.Markdown(label="Status / tips")
 
         run_btn.click(
             infer,
             inputs=[ref_in, test_in, use_sam, sam_prompt],
-            outputs=[overlay_out, mask_out, score_out, status_out],
+            outputs=[overlay_out, score_out, status_out],
         )
 
     return demo