feat: SAM2 for POM / F0 masking

app.py

@@ -19,6 +19,30 @@ from chord.util import get_positions, rgb_to_srgb
 from chord.io import load_torch_file
 from chord.minecraft_pbr import convert_to_labpbr, convert_to_bedrock, LABPBR_METAL_CHOICES
 
+# Try to import SAM 2 - it's optional
+SAM_AVAILABLE = False
+try:
+    from chord.sam_segmenter import SAM2Segmenter, create_mask_overlay, draw_points_on_image
+    SAM_AVAILABLE = True
+    print("SAM 2 segmenter available")
+except ImportError as e:
+    print(f"SAM 2 not available: {e}")
+    print("Install with: pip install sam2")
+    # Create dummy functions so the app doesn't crash
+    SAM2Segmenter = None
+    def create_mask_overlay(image, mask, color=(255, 100, 100), alpha=0.5):
+        return image
+    def draw_points_on_image(image, fg_points, bg_points=None, point_radius=6):
+        from PIL import ImageDraw
+        draw_img = image.copy()
+        draw = ImageDraw.Draw(draw_img)
+        r = point_radius
+        for x, y in (fg_points or []):
+            draw.ellipse([x - r, y - r, x + r, y + r], fill=(0, 255, 0), outline=(0, 180, 0), width=2)
+        for x, y in (bg_points or []):
+            draw.ellipse([x - r, y - r, x + r, y + r], fill=(255, 0, 0), outline=(180, 0, 0), width=2)
+        return draw_img
+
 
 def _load_examples(directory: str) -> list:
     """Load example images from a directory, returning empty list if not found."""
@@ -84,6 +108,359 @@ def relit(model, maps):
     rgb = model.model.compute_render(maps, camera, pos, light).squeeze(0).permute(0,3,1,2)  # GxBxHxWxC -> BxCxHxW
     return torch.clamp(rgb_to_srgb(rgb), 0, 1)
 
+
+# =============================================================================
+# SAM 2 Segmentation Helpers
+# =============================================================================
+
+# Global cache for SAM2 model (loaded once, moved to GPU as needed)
+_SAM2_MODEL = None
+
+
+def _get_sam2_model():
+    """Get or create cached SAM2 model (on CPU for storage)."""
+    global _SAM2_MODEL
+    if _SAM2_MODEL is None:
+        print("Loading SAM 2 model (first time, will be cached)...")
+        from sam2.build_sam import build_sam2_hf
+        # Load to CPU first - will be moved to CUDA in @spaces.GPU function
+        _SAM2_MODEL = build_sam2_hf(
+            model_id="facebook/sam2.1-hiera-small",
+            device=torch.device("cpu")
+        )
+        print("SAM 2 model cached on CPU")
+    return _SAM2_MODEL
+
+
+def on_image_upload_for_mask(image):
+    """Reset SAM state when new image is uploaded."""
+    # Returns: fg_points, bg_points, mask, mask_preview_image
+    return [], [], None, image
+
+
+def feather_mask(mask, radius):
+    """Apply Gaussian blur to mask edges for soft feathering.
+
+    Args:
+        mask: Binary or soft mask as numpy array (H, W) with values 0-1
+        radius: Feather radius in pixels (0 = no feathering)
+
+    Returns:
+        Feathered mask with soft edges
+    """
+    if radius <= 0:
+        return mask
+
+    import numpy as np
+    from scipy.ndimage import gaussian_filter
+
+    # Gaussian blur creates soft edges
+    # Sigma is approximately radius/2 for natural-looking feather
+    sigma = radius / 2.0
+    feathered = gaussian_filter(mask.astype(np.float32), sigma=sigma)
+
+    return feathered
+
+
+@spaces.GPU
+def run_sam_prediction(image, fg_points, bg_points):
+    """Run SAM prediction on GPU. Must be in @spaces.GPU decorated function."""
+    import numpy as np
+    from sam2.sam2_image_predictor import SAM2ImagePredictor
+
+    # Get cached model and move to CUDA for this call
+    sam2_model = _get_sam2_model()
+    sam2_model = sam2_model.to("cuda")
+
+    predictor = SAM2ImagePredictor(sam2_model)
+
+    # Set image
+    image_np = np.array(image.convert("RGB"))
+    with torch.inference_mode():
+        predictor.set_image(image_np)
+
+    # Build point arrays
+    all_points = []
+    all_labels = []
+    for x, y in fg_points:
+        all_points.append([x, y])
+        all_labels.append(1)
+    if bg_points:
+        for x, y in bg_points:
+            all_points.append([x, y])
+            all_labels.append(0)
+
+    point_coords = np.array(all_points)
+    point_labels = np.array(all_labels)
+
+    # Predict
+    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+        masks, scores, _ = predictor.predict(
+            point_coords=point_coords,
+            point_labels=point_labels,
+            multimask_output=True,
+        )
+
+    # Return best mask (numpy array, already on CPU)
+    best_idx = np.argmax(scores)
+    return masks[best_idx].astype(np.float32), float(scores[best_idx])
+
+
+def on_mask_image_click(image, fg_points, bg_points, evt: gr.SelectData, point_mode, feather):
+    """Handle click on image for point annotation."""
+    import numpy as np
+
+    if image is None:
+        print("SAM: No image loaded")
+        return fg_points, bg_points, None, None
+
+    x, y = evt.index
+    print(f"SAM click at ({x}, {y}) - mode: {point_mode}, feather: {feather}")
+
+    if point_mode == "foreground":
+        fg_points = list(fg_points) + [(x, y)]
+    else:
+        bg_points = list(bg_points) + [(x, y)]
+
+    # Always draw points on preview so user sees their clicks
+    preview = draw_points_on_image(image, fg_points, bg_points)
+    mask = None
+
+    # Generate mask preview if we have foreground points
+    if fg_points:
+        try:
+            mask, score = run_sam_prediction(image, fg_points, bg_points)
+            print(f"SAM mask generated, score: {score:.3f}")
+
+            # Apply feathering to mask edges
+            mask = feather_mask(mask, int(feather))
+
+            # Create overlay visualization with mask
+            preview = create_mask_overlay(image, mask, color=(255, 100, 100), alpha=0.5)
+            preview = draw_points_on_image(preview, fg_points, bg_points)
+        except Exception as e:
+            print(f"SAM error: {e}")
+            import traceback
+            traceback.print_exc()
+            # Keep preview with just points drawn (already set above)
+
+    return fg_points, bg_points, mask, preview
+
+
+def clear_mask_state(original_image):
+    """Clear all mask-related state."""
+    # Returns: fg_points, bg_points, mask, mask_preview_image
+    return [], [], None, original_image
+
+
+def regenerate_mask_preview(image, fg_points, bg_points, feather):
+    """Regenerate mask preview with current points."""
+    if image is None:
+        return None, None
+
+    # Always draw points even without mask
+    preview = draw_points_on_image(image, fg_points, bg_points) if fg_points else image
+
+    if not fg_points:
+        return None, preview
+
+    try:
+        mask, score = run_sam_prediction(image, fg_points, bg_points)
+        print(f"SAM regenerate mask, score: {score:.3f}")
+
+        # Apply feathering to mask edges
+        mask = feather_mask(mask, int(feather))
+
+        preview = create_mask_overlay(image, mask, color=(255, 100, 100), alpha=0.5)
+        preview = draw_points_on_image(preview, fg_points, bg_points)
+
+        return mask, preview
+    except Exception as e:
+        print(f"SAM regenerate error: {e}")
+        import traceback
+        traceback.print_exc()
+        return None, preview
+
+
+# =============================================================================
+# Metal Mask SAM Helpers
+# =============================================================================
+
+# Color map for different metal types (for visualization)
+METAL_COLORS = {
+    "iron": (180, 180, 180),      # Gray
+    "gold": (255, 215, 0),        # Gold
+    "aluminum": (200, 200, 210),  # Light gray-blue
+    "chrome": (220, 220, 230),    # Silver
+    "copper": (184, 115, 51),     # Copper
+    "lead": (100, 100, 110),      # Dark gray
+    "platinum": (229, 228, 226),  # Platinum
+    "silver": (192, 192, 192),    # Silver
+    "custom": (255, 100, 255),    # Magenta for custom
+}
+
+
+def on_image_upload_for_metal_mask(image):
+    """Reset metal mask state when new image is uploaded."""
+    import numpy as np
+    # Returns: fg_points, bg_points, current_segment_mask, combined_metal_mask, preview
+    return [], [], None, None, image
+
+
+def on_metal_mask_image_click(image, fg_points, bg_points, evt: gr.SelectData, point_mode, metal_type, feather):
+    """Handle click on image for metal mask point annotation."""
+    import numpy as np
+
+    if image is None:
+        print("Metal SAM: No image loaded")
+        return fg_points, bg_points, None, None
+
+    x, y = evt.index
+    print(f"Metal SAM click at ({x}, {y}) - mode: {point_mode}, metal: {metal_type}, feather: {feather}")
+
+    if point_mode == "foreground":
+        fg_points = list(fg_points) + [(x, y)]
+    else:
+        bg_points = list(bg_points) + [(x, y)]
+
+    # Always draw points on preview so user sees their clicks
+    preview = draw_points_on_image(image, fg_points, bg_points)
+    mask = None
+
+    # Generate mask preview if we have foreground points
+    if fg_points:
+        try:
+            mask, score = run_sam_prediction(image, fg_points, bg_points)
+            print(f"Metal SAM mask generated, score: {score:.3f}")
+
+            # Apply feathering to mask edges
+            mask = feather_mask(mask, int(feather))
+
+            # Get color for current metal type
+            color = METAL_COLORS.get(metal_type, (255, 100, 100))
+
+            # Create overlay visualization with mask
+            preview = create_mask_overlay(image, mask, color=color, alpha=0.5)
+            preview = draw_points_on_image(preview, fg_points, bg_points)
+        except Exception as e:
+            print(f"Metal SAM error: {e}")
+            import traceback
+            traceback.print_exc()
+            # Keep preview with just points drawn (already set above)
+
+    return fg_points, bg_points, mask, preview
+
+
+def add_metal_region_to_mask(current_segment_mask, combined_metal_mask, metal_type, image_size):
+    """Add the current SAM segment to the combined metal mask with the selected metal type."""
+    import numpy as np
+    from chord.minecraft_pbr import LABPBR_METALS
+
+    if current_segment_mask is None:
+        return combined_metal_mask, "No segment to add. Click on the image to create a segment first."
+
+    # Get the metal ID value
+    metal_id = LABPBR_METALS.get(metal_type, 255)
+    if metal_id is None:
+        return combined_metal_mask, "Invalid metal type selected."
+
+    # Normalize to [0, 1] scale for the mask tensor
+    normalized_metal_value = metal_id / 255.0
+
+    # Initialize combined mask if needed
+    if combined_metal_mask is None:
+        combined_metal_mask = np.zeros(current_segment_mask.shape, dtype=np.float32)
+
+    # Add current segment with the metal value (overwrites existing values in the region)
+    combined_metal_mask = np.where(current_segment_mask > 0.5, normalized_metal_value, combined_metal_mask)
+
+    return combined_metal_mask, f"Added {metal_type} region (ID: {metal_id})"
+
+
+def regenerate_metal_mask_preview(image, fg_points, bg_points, metal_type, feather):
+    """Regenerate metal mask preview with current points and feather setting."""
+    if image is None:
+        return fg_points, bg_points, None, None
+
+    # Always draw points even without mask
+    preview = draw_points_on_image(image, fg_points, bg_points) if fg_points else image
+
+    if not fg_points:
+        return fg_points, bg_points, None, preview
+
+    try:
+        mask, score = run_sam_prediction(image, fg_points, bg_points)
+        print(f"Metal SAM regenerate mask, score: {score:.3f}")
+
+        # Apply feathering to mask edges
+        mask = feather_mask(mask, int(feather))
+
+        # Get color for current metal type
+        color = METAL_COLORS.get(metal_type, (255, 100, 100))
+
+        preview = create_mask_overlay(image, mask, color=color, alpha=0.5)
+        preview = draw_points_on_image(preview, fg_points, bg_points)
+
+        return fg_points, bg_points, mask, preview
+    except Exception as e:
+        print(f"Metal SAM regenerate error: {e}")
+        import traceback
+        traceback.print_exc()
+        return fg_points, bg_points, None, preview
+
+
+def clear_metal_mask_segment(original_image):
+    """Clear current segment points but keep combined mask."""
+    # Returns: fg_points, bg_points, current_segment_mask, preview (reset to original)
+    return [], [], None, original_image
+
+
+def clear_all_metal_masks(original_image):
+    """Clear all metal masks including combined mask."""
+    # Returns: fg_points, bg_points, current_segment_mask, combined_metal_mask, preview
+    return [], [], None, None, original_image
+
+
+def create_metal_mask_preview(image, combined_metal_mask):
+    """Create a preview visualization of all metal regions."""
+    import numpy as np
+    from PIL import Image as PILImage
+    from chord.minecraft_pbr import LABPBR_METALS
+
+    if image is None or combined_metal_mask is None:
+        return image
+
+    # Create inverse lookup: value -> metal name
+    value_to_metal = {v / 255.0: k for k, v in LABPBR_METALS.items() if v is not None}
+
+    # Convert image to numpy array
+    img_array = np.array(image).astype(np.float32)
+
+    # Create colored overlay for each metal type
+    overlay = np.zeros_like(img_array)
+    mask_active = combined_metal_mask > 0
+
+    # Find unique metal values in the mask
+    unique_values = np.unique(combined_metal_mask[mask_active])
+
+    for val in unique_values:
+        metal_name = value_to_metal.get(val, "custom")
+        color = METAL_COLORS.get(metal_name, (255, 100, 255))
+        region = np.abs(combined_metal_mask - val) < 0.01
+        for c in range(3):
+            overlay[:, :, c] = np.where(region, color[c], overlay[:, :, c])
+
+    # Blend overlay with original image
+    alpha = 0.5
+    result = np.where(
+        mask_active[:, :, np.newaxis],
+        img_array * (1 - alpha) + overlay * alpha,
+        img_array
+    )
+
+    return PILImage.fromarray(result.astype(np.uint8))
+
+
 @spaces.GPU
 def inference(
     img,
@@ -107,6 +484,8 @@ def inference(
     emission_knee,
     emission_bloom,
     hardcoded_metal,
+    height_mask,
+    metal_mask,
 ):
     """
     Run Chord model and output shader-compatible textures.
@@ -141,6 +520,29 @@ def inference(
     roughness = resize_back(out["roughness"].unsqueeze(0) if out["roughness"].dim() == 2 else out["roughness"])
     metalness = resize_back(out["metalness"].unsqueeze(0) if out["metalness"].dim() == 2 else out["metalness"])
 
+    # Get device from model output tensors
+    device = basecolor.device
+
+    # Convert height mask from numpy to tensor if provided
+    height_mask_tensor = None
+    if height_mask is not None:
+        import numpy as np
+        height_mask_tensor = torch.from_numpy(height_mask).float().to(device)
+        # Resize mask to match output resolution
+        if height_mask_tensor.dim() == 2:
+            height_mask_tensor = height_mask_tensor.unsqueeze(0).unsqueeze(0)
+        height_mask_tensor = v2.Resize(size=(ori_h, ori_w))(height_mask_tensor)
+
+    # Convert metal mask from numpy to tensor if provided
+    metal_mask_tensor = None
+    if metal_mask is not None:
+        import numpy as np
+        metal_mask_tensor = torch.from_numpy(metal_mask).float().to(device)
+        # Resize mask to match output resolution
+        if metal_mask_tensor.dim() == 2:
+            metal_mask_tensor = metal_mask_tensor.unsqueeze(0).unsqueeze(0)
+        metal_mask_tensor = v2.Resize(size=(ori_h, ori_w), interpolation=v2.InterpolationMode.NEAREST)(metal_mask_tensor)
+
     if output_format == "bedrock":
         # Convert to Bedrock RTX format (MER/MERS)
         result = convert_to_bedrock(
@@ -176,6 +578,7 @@ def inference(
             height_mid_freq=height_mid_freq,
             height_high_freq=height_high_freq,
             height_intensity=height_intensity,
+            height_mask=height_mask_tensor,
             seamless=seamless,
             ao_strength=ao_strength,
             ao_blur=int(ao_blur),
@@ -190,6 +593,7 @@ def inference(
             emission_knee=emission_knee,
             emission_bloom=int(emission_bloom),
             hardcoded_metal=hardcoded_metal,
+            metal_mask=metal_mask_tensor,
         )
         return (
             result['albedo'],
@@ -278,6 +682,127 @@ Upload an image to estimate PBR materials and export for Minecraft shaders.
                     info="Use predefined metal F0 values (230-237) for metallic areas"
                 )
 
+            with gr.Accordion("POM Height Mask (Optional - SAM 2)", open=False):
+                gr.Markdown("""
+Use SAM 2 to create a mask that flattens selected regions in the height map.
+Click on the image below to add points:
+- **Green points (Foreground)**: Include in mask - areas will be flattened
+- **Red points (Background)**: Exclude from mask - refine the selection
+""")
+
+                with gr.Row():
+                    point_mode = gr.Radio(
+                        choices=["foreground", "background"],
+                        value="foreground",
+                        label="Point Mode",
+                        info="Foreground = add to mask (flatten), Background = exclude from mask"
+                    )
+                    mask_feather = gr.Slider(
+                        minimum=0, maximum=50, value=0, step=1,
+                        label="Feather",
+                        info="Blur mask edges for soft transitions (0 = hard edges)"
+                    )
+
+                with gr.Row():
+                    mask_image = gr.Image(
+                        type="pil",
+                        label="Click to add points",
+                        interactive=True,
+                        height=256,
+                    )
+                    mask_preview = gr.Image(
+                        type="pil",
+                        label="Mask Preview (red = will be flattened)",
+                        interactive=False,
+                        height=256,
+                    )
+
+                with gr.Row():
+                    clear_mask_btn = gr.Button("Clear Mask", size="sm")
+                    regenerate_btn = gr.Button("Regenerate Preview", size="sm")
+
+                # Hidden state components
+                fg_points_state = gr.State([])
+                bg_points_state = gr.State([])
+                current_mask_state = gr.State(None)
+
+            with gr.Accordion("Metal Type Mask (Optional - SAM 2, LabPBR only)", open=False):
+                gr.Markdown("""
+Use SAM 2 to paint regions with specific LabPBR metal types (iron, gold, copper, etc.).
+This overrides the global "Hardcoded Metal" setting for selected regions.
+
+**Workflow:**
+1. Select a metal type from the dropdown
+2. Click on the image to add foreground points (areas to mark as this metal)
+3. Optionally add background points to refine the selection
+4. Click "Add Region" to save this metal region
+5. Repeat for other metal types if needed
+""")
+
+                metal_type_selector = gr.Dropdown(
+                    choices=[
+                        ("Iron (230)", "iron"),
+                        ("Gold (231)", "gold"),
+                        ("Aluminum (232)", "aluminum"),
+                        ("Chrome (233)", "chrome"),
+                        ("Copper (234)", "copper"),
+                        ("Lead (235)", "lead"),
+                        ("Platinum (236)", "platinum"),
+                        ("Silver (237)", "silver"),
+                        ("Custom Metal (255)", "custom"),
+                    ],
+                    value="iron",
+                    label="Metal Type to Paint",
+                    info="Select the metal type before clicking on the image"
+                )
+
+                with gr.Row():
+                    metal_point_mode = gr.Radio(
+                        choices=["foreground", "background"],
+                        value="foreground",
+                        label="Point Mode",
+                        info="Foreground = add to selection, Background = exclude from selection"
+                    )
+                    metal_mask_feather = gr.Slider(
+                        minimum=0, maximum=50, value=0, step=1,
+                        label="Feather",
+                        info="Blur mask edges for soft transitions (0 = hard edges)"
+                    )
+
+                with gr.Row():
+                    metal_mask_image = gr.Image(
+                        type="pil",
+                        label="Click to select metal regions",
+                        interactive=True,
+                        height=256,
+                    )
+                    metal_mask_preview = gr.Image(
+                        type="pil",
+                        label="Current Selection Preview",
+                        interactive=False,
+                        height=256,
+                    )
+
+                with gr.Row():
+                    add_metal_region_btn = gr.Button("Add Region", variant="primary", size="sm")
+                    clear_metal_segment_btn = gr.Button("Clear Selection", size="sm")
+                    clear_all_metals_btn = gr.Button("Clear All Metals", size="sm")
+
+                metal_status = gr.Textbox(label="Status", interactive=False, value="No metal regions defined")
+
+                combined_metal_preview = gr.Image(
+                    type="pil",
+                    label="Combined Metal Mask (all regions)",
+                    interactive=False,
+                    height=256,
+                )
+
+                # Hidden state components for metal mask
+                metal_fg_points_state = gr.State([])
+                metal_bg_points_state = gr.State([])
+                metal_current_segment_state = gr.State(None)
+                metal_combined_mask_state = gr.State(None)
+
             gr.Markdown("### Example Inputs — Generated Textures")
             gr.Examples(
                 examples=EXAMPLES_USECASE_1,
@@ -310,6 +835,122 @@ Upload an image to estimate PBR materials and export for Minecraft shaders.
             gr.Markdown("### Preview")
             render_out = gr.Image(label="Relit Preview (Point Light)", height=340, format="png")
 
+    # ==========================================================================
+    # SAM 2 Event Handlers
+    # ==========================================================================
+
+    # Sync input image to mask editor when uploaded
+    input_img.change(
+        fn=on_image_upload_for_mask,
+        inputs=[input_img],
+        outputs=[fg_points_state, bg_points_state, current_mask_state, mask_image]
+    )
+
+    # Handle clicks on mask image for point annotation
+    mask_image.select(
+        fn=on_mask_image_click,
+        inputs=[mask_image, fg_points_state, bg_points_state, point_mode, mask_feather],
+        outputs=[fg_points_state, bg_points_state, current_mask_state, mask_preview]
+    )
+
+    # Clear mask button
+    clear_mask_btn.click(
+        fn=clear_mask_state,
+        inputs=[input_img],
+        outputs=[fg_points_state, bg_points_state, current_mask_state, mask_preview]
+    )
+
+    # Regenerate mask preview button (also updates when feather changes)
+    regenerate_btn.click(
+        fn=regenerate_mask_preview,
+        inputs=[mask_image, fg_points_state, bg_points_state, mask_feather],
+        outputs=[current_mask_state, mask_preview]
+    )
+
+    # Auto-regenerate when feather slider changes
+    mask_feather.change(
+        fn=regenerate_mask_preview,
+        inputs=[mask_image, fg_points_state, bg_points_state, mask_feather],
+        outputs=[current_mask_state, mask_preview]
+    )
+
+    # ==========================================================================
+    # Metal Mask SAM 2 Event Handlers
+    # ==========================================================================
+
+    # Sync input image to metal mask editor when uploaded
+    input_img.change(
+        fn=on_image_upload_for_metal_mask,
+        inputs=[input_img],
+        outputs=[metal_fg_points_state, metal_bg_points_state, metal_current_segment_state,
+                 metal_combined_mask_state, metal_mask_image]
+    )
+
+    # Handle clicks on metal mask image for point annotation
+    metal_mask_image.select(
+        fn=on_metal_mask_image_click,
+        inputs=[metal_mask_image, metal_fg_points_state, metal_bg_points_state,
+                metal_point_mode, metal_type_selector, metal_mask_feather],
+        outputs=[metal_fg_points_state, metal_bg_points_state,
+                 metal_current_segment_state, metal_mask_preview]
+    )
+
+    # Auto-regenerate when metal feather slider changes
+    metal_mask_feather.change(
+        fn=regenerate_metal_mask_preview,
+        inputs=[metal_mask_image, metal_fg_points_state, metal_bg_points_state,
+                metal_type_selector, metal_mask_feather],
+        outputs=[metal_fg_points_state, metal_bg_points_state,
+                 metal_current_segment_state, metal_mask_preview]
+    )
+
+    # Add current region to combined metal mask
+    def add_region_handler(current_segment, combined_mask, metal_type, image):
+        new_mask, status = add_metal_region_to_mask(current_segment, combined_mask, metal_type, None)
+        preview = create_metal_mask_preview(image, new_mask)
+        # Clear current segment points after adding
+        return [], [], None, new_mask, status, image, preview
+
+    add_metal_region_btn.click(
+        fn=add_region_handler,
+        inputs=[metal_current_segment_state, metal_combined_mask_state,
+                metal_type_selector, input_img],
+        outputs=[metal_fg_points_state, metal_bg_points_state, metal_current_segment_state,
+                 metal_combined_mask_state, metal_status, metal_mask_preview, combined_metal_preview]
+    )
+
+    # Clear current segment (but keep combined mask)
+    clear_metal_segment_btn.click(
+        fn=clear_metal_mask_segment,
+        inputs=[input_img],
+        outputs=[metal_fg_points_state, metal_bg_points_state,
+                 metal_current_segment_state, metal_mask_preview]
+    )
+
+    # Clear all metal masks
+    clear_all_metals_btn.click(
+        fn=clear_all_metal_masks,
+        inputs=[input_img],
+        outputs=[metal_fg_points_state, metal_bg_points_state, metal_current_segment_state,
+                 metal_combined_mask_state, metal_mask_preview]
+    )
+
+    # Update combined preview when combined mask changes
+    def update_combined_preview(image, combined_mask):
+        if combined_mask is None:
+            return image
+        return create_metal_mask_preview(image, combined_mask)
+
+    metal_combined_mask_state.change(
+        fn=update_combined_preview,
+        inputs=[input_img, metal_combined_mask_state],
+        outputs=[combined_metal_preview]
+    )
+
+    # ==========================================================================
+    # Main Inference
+    # ==========================================================================
+
     run_button.click(
         inference,
         inputs=[
@@ -334,6 +975,8 @@ Upload an image to estimate PBR materials and export for Minecraft shaders.
             emission_knee,
             emission_bloom,
             hardcoded_metal,
+            current_mask_state,
+            metal_combined_mask_state,
         ],
         outputs=[albedo_out, packed_out, normal_out, render_out]
     )

chord/minecraft_pbr.py

@@ -74,6 +74,7 @@ def metalness_to_f0(
     metalness: torch.Tensor,
     threshold: float = 0.5,
     hardcoded_metal: str = "none",
+    metal_mask: torch.Tensor = None,
 ) -> torch.Tensor:
     """
     Convert metalness to LabPBR F0/metal channel.
@@ -91,6 +92,10 @@ def metalness_to_f0(
         threshold: Threshold above which material is considered metal
         hardcoded_metal: Name of predefined metal type ("none", "custom", "iron", "gold", etc.)
             When not "none", metallic areas use this specific metal value instead of 255.
+        metal_mask: Optional per-pixel metal type mask, shape matching metalness.
+            Values are LabPBR metal IDs (230-255) normalized to [0, 1].
+            Where mask > 0, uses the mask value directly as F0.
+            Where mask = 0, falls back to default metalness-based behavior.
 
     Returns:
         F0 channel values in range [0, 1] (stored LINEAR, scaled to 0-255 on save)
@@ -109,12 +114,26 @@ def metalness_to_f0(
 
     # Blend based on metalness (hard threshold for cleaner results)
     is_metal = (metalness > threshold).float()
-    return torch.lerp(
+    f0 = torch.lerp(
         torch.full_like(metalness, dielectric_f0),
         torch.full_like(metalness, metal_f0),
         is_metal
     )
 
+    # Apply per-pixel metal mask if provided
+    # Metal mask contains normalized metal IDs (e.g., 230/255 for iron)
+    # Where mask > 0, override with mask value
+    if metal_mask is not None:
+        # Ensure mask has same shape as f0
+        if metal_mask.shape != f0.shape:
+            # Try to broadcast or resize
+            if metal_mask.dim() < f0.dim():
+                metal_mask = metal_mask.unsqueeze(0)
+        mask_active = (metal_mask > 0).float()
+        f0 = torch.where(mask_active > 0.5, metal_mask, f0)
+
+    return f0
+
 
 def convert_normal_to_directx(
     normal: torch.Tensor,
@@ -464,6 +483,7 @@ def create_specular_texture(
     sss: torch.Tensor = None,
     emission: torch.Tensor = None,
     hardcoded_metal: str = "none",
+    metal_mask: torch.Tensor = None,
     sss_threshold: float = 0.01,
 ) -> torch.Tensor:
     """
@@ -486,6 +506,8 @@ def create_specular_texture(
         emission: Optional emission intensity map, range [0, 1] where 1 = max emission
             If provided, output is RGBA; if None, output is RGB
         hardcoded_metal: Predefined metal type for metallic areas ("none", "iron", "gold", etc.)
+        metal_mask: Optional per-pixel metal type mask with normalized LabPBR metal IDs.
+            Where mask > 0, uses mask value as F0 (overrides hardcoded_metal).
         sss_threshold: SSS intensity threshold for per-pixel priority (default 0.01)
             Pixels with SSS > threshold use SSS encoding, others use porosity
 
@@ -511,7 +533,7 @@ def create_specular_texture(
     smoothness = roughness_to_smoothness(roughness)
 
     # G: F0/Metal (stored LINEAR, uses hardcoded metal value for metallic areas if specified)
-    f0 = metalness_to_f0(metalness, hardcoded_metal=hardcoded_metal)
+    f0 = metalness_to_f0(metalness, hardcoded_metal=hardcoded_metal, metal_mask=metal_mask)
 
     # B: Porosity (0-64) and/or SSS (65-255) - per-pixel blending
     # SSS takes priority where it exceeds threshold
@@ -654,6 +676,7 @@ def convert_to_labpbr(
     height_high_freq: float = 1.0,
     height_intensity: float = 1.0,
     height_invert: bool = True,
+    height_mask: torch.Tensor = None,
     compute_porosity: bool = False,
     normalize_porosity: bool = True,
     compute_sss: bool = False,
@@ -665,6 +688,7 @@ def convert_to_labpbr(
     emission_knee: float = 0.1,
     emission_bloom: int = 0,
     hardcoded_metal: str = "none",
+    metal_mask: torch.Tensor = None,
     seamless: bool = False,
     ao_strength: float = 2.0,
     ao_blur: int = 5,
@@ -694,6 +718,8 @@ def convert_to_labpbr(
         height_intensity: Global height intensity/opacity (0.0 = flat, 1.0 = full height)
         height_invert: If True, invert height for correct POM direction (default True).
             This fixes the Frankot-Chellappa output so raised areas appear raised in POM.
+        height_mask: Optional mask tensor where 1=suppress height (flatten), 0=keep height.
+            Used with SAM segmentation for POM masking. Masked regions become flat (min_height).
         compute_porosity: If True, calculate porosity from AO, smoothness, and F0 when not provided
         normalize_porosity: If True, normalize porosity to full 0-1 range before LabPBR scaling
         compute_sss: If True, calculate SSS thickness from normal curvature and AO when not provided
@@ -706,6 +732,9 @@ def convert_to_labpbr(
         emission_bloom: Gaussian blur radius for emission bloom effect (0 = disabled)
         hardcoded_metal: Predefined metal type for specular G channel ("none", "iron", "gold", etc.)
             Uses metalness map as mask - metallic areas get this metal ID value.
+        metal_mask: Optional per-pixel metal type mask with normalized LabPBR metal IDs (0-1 scale).
+            Created via SAM segmentation. Where mask > 0, overrides hardcoded_metal with per-pixel values.
+            Values should be LabPBR metal IDs divided by 255 (e.g., 230/255 for iron).
         seamless: Whether the texture should tile seamlessly (for height derivation)
         ao_strength: AO intensity multiplier (higher = more contrast)
         ao_blur: Gaussian blur radius for AO smoothing
@@ -749,6 +778,7 @@ def convert_to_labpbr(
             height_intensity=height_intensity,
             height_min=0.25,  # Minecraft POM: black = 25% block depth
             height_invert=height_invert,
+            height_mask=height_mask,
         )
         if ao is None:
             ao = derived_ao
@@ -792,7 +822,8 @@ def convert_to_labpbr(
 
     # Create LabPBR textures
     specular_tex = create_specular_texture(
-        roughness, metalness, porosity, sss, emission, hardcoded_metal=hardcoded_metal
+        roughness, metalness, porosity, sss, emission,
+        hardcoded_metal=hardcoded_metal, metal_mask=metal_mask
     )
     normal_tex = create_normal_texture(
         normal, ao, height, flip_y=flip_normal_y, swap_xy=swap_normal_xy

chord/normal_utils.py

@@ -161,6 +161,7 @@ def normal_to_height(
     intensity: float = 1.0,
     min_height: float = 0.25,
     invert: bool = True,
+    height_mask: torch.Tensor = None,
 ) -> torch.Tensor:
     """
     Convert normal map to height map using Frankot-Chellappa algorithm.
@@ -177,6 +178,8 @@ def normal_to_height(
         invert: If True, invert height so raised areas in normal map appear raised
             in POM (default True). The Frankot-Chellappa integration can produce
             inverted heights depending on gradient sign conventions.
+        height_mask: Optional mask tensor where 1=suppress height (flatten),
+            0=keep height. Used with SAM segmentation for POM masking.
 
     Returns:
         Height map (B, 1, H, W) or (1, H, W), range [min_height, 1.0]
@@ -216,6 +219,30 @@ def normal_to_height(
         midpoint = (min_height + 1.0) / 2.0
         height = midpoint + (height - midpoint) * intensity
 
+    # Apply height mask (mask=1 suppresses height, mask=0 keeps height)
+    if height_mask is not None:
+        # Ensure mask has correct shape (B, 1, H, W)
+        if height_mask.dim() == 2:
+            height_mask = height_mask.unsqueeze(0).unsqueeze(0)
+        elif height_mask.dim() == 3:
+            height_mask = height_mask.unsqueeze(0)
+
+        # Resize mask to match height dimensions if needed
+        if height_mask.shape[-2:] != height.shape[-2:]:
+            height_mask = F.interpolate(
+                height_mask.float(),
+                size=height.shape[-2:],
+                mode='bilinear',
+                align_corners=False
+            )
+
+        # Move mask to same device/dtype as height
+        height_mask = height_mask.to(device=height.device, dtype=height.dtype)
+
+        # Blend: masked areas (mask=1) go to min_height (flat surface)
+        # Unmasked areas (mask=0) keep their derived height
+        height = height * (1.0 - height_mask) + min_height * height_mask
+
     if squeeze:
         height = height.squeeze(0)
 
@@ -336,6 +363,7 @@ def derive_ao_and_height(
     height_intensity: float = 1.0,
     height_min: float = 0.25,
     height_invert: bool = True,
+    height_mask: torch.Tensor = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Derive both AO and height from a normal map.
@@ -351,6 +379,8 @@ def derive_ao_and_height(
         height_intensity: Global height intensity/opacity (0.0 = flat, 1.0 = full)
         height_min: Minimum height value (default 0.25 for Minecraft POM)
         height_invert: If True, invert height for correct POM direction (default True)
+        height_mask: Optional mask tensor where 1=suppress height (flatten),
+            0=keep height. Used with SAM segmentation for POM masking.
 
     Returns:
         ao: Ambient occlusion (1.0 = no occlusion)
@@ -366,6 +396,7 @@ def derive_ao_and_height(
         intensity=height_intensity,
         min_height=height_min,
         invert=height_invert,
+        height_mask=height_mask,
     )
 
     return ao, height

chord/sam_segmenter.py

@@ -0,0 +1,207 @@
+"""
+SAM 2 Segmentation Wrapper for POM Height Masking.
+
+Provides interactive point-prompt segmentation using SAM 2.1 Hiera Small.
+Users click on images to segment regions for height mask creation.
+"""
+
+import torch
+import numpy as np
+from PIL import Image, ImageDraw
+from typing import Optional, Tuple, List
+
+# Global singleton for lazy loading
+_SAM2_PREDICTOR = None
+_SAM2_MODEL_CFG = None
+
+
+def _get_sam2_predictor():
+    """Lazy load SAM 2 predictor on first use."""
+    global _SAM2_PREDICTOR, _SAM2_MODEL_CFG
+
+    if _SAM2_PREDICTOR is None:
+        print("Loading SAM 2 model...")
+
+        from sam2.build_sam import build_sam2_hf
+        from sam2.sam2_image_predictor import SAM2ImagePredictor
+
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        # Use build_sam2_hf which handles HuggingFace model loading directly
+        sam2_model = build_sam2_hf(
+            model_id="facebook/sam2.1-hiera-small",
+            device=device
+        )
+        _SAM2_PREDICTOR = SAM2ImagePredictor(sam2_model)
+        _SAM2_MODEL_CFG = "sam2.1_hiera_s"
+
+        print(f"SAM 2 loaded on {device}")
+
+    return _SAM2_PREDICTOR
+
+
+class SAM2Segmenter:
+    """SAM 2 wrapper for interactive point-prompt segmentation."""
+
+    def __init__(self):
+        self.predictor = None
+        self.current_image = None
+        self._image_set = False
+
+    def set_image(self, image: Image.Image) -> None:
+        """Set the image for segmentation. Must be called before predict."""
+        if self.predictor is None:
+            self.predictor = _get_sam2_predictor()
+
+        # Convert PIL to numpy array (RGB)
+        image_np = np.array(image.convert("RGB"))
+
+        with torch.inference_mode():
+            self.predictor.set_image(image_np)
+
+        self.current_image = image
+        self._image_set = True
+
+    def predict_mask(
+        self,
+        fg_points: List[Tuple[int, int]],
+        bg_points: Optional[List[Tuple[int, int]]] = None,
+        multimask_output: bool = True,
+    ) -> Tuple[np.ndarray, float]:
+        """
+        Predict segmentation mask from point prompts.
+
+        Args:
+            fg_points: List of (x, y) foreground points (areas to include)
+            bg_points: List of (x, y) background points (areas to exclude)
+            multimask_output: If True, return best of 3 masks
+
+        Returns:
+            mask: Binary mask (H, W) as numpy array, 1=selected region
+            score: IoU prediction score
+        """
+        if not self._image_set:
+            raise RuntimeError("Must call set_image() before predict_mask()")
+
+        if not fg_points:
+            raise ValueError("At least one foreground point is required")
+
+        # Build point arrays
+        all_points = []
+        all_labels = []
+
+        for x, y in fg_points:
+            all_points.append([x, y])
+            all_labels.append(1)  # Foreground
+
+        if bg_points:
+            for x, y in bg_points:
+                all_points.append([x, y])
+                all_labels.append(0)  # Background
+
+        point_coords = np.array(all_points)
+        point_labels = np.array(all_labels)
+
+        with torch.inference_mode():
+            # Use autocast if on CUDA
+            device = self.predictor.device
+            if device.type == "cuda":
+                with torch.autocast("cuda", dtype=torch.bfloat16):
+                    masks, scores, _ = self.predictor.predict(
+                        point_coords=point_coords,
+                        point_labels=point_labels,
+                        multimask_output=multimask_output,
+                    )
+            else:
+                masks, scores, _ = self.predictor.predict(
+                    point_coords=point_coords,
+                    point_labels=point_labels,
+                    multimask_output=multimask_output,
+                )
+
+        # Return best mask (highest IoU score)
+        best_idx = np.argmax(scores)
+        return masks[best_idx].astype(np.float32), float(scores[best_idx])
+
+    def clear(self) -> None:
+        """Reset the segmenter state."""
+        self._image_set = False
+        self.current_image = None
+
+
+def create_mask_overlay(
+    image: Image.Image,
+    mask: np.ndarray,
+    color: Tuple[int, int, int] = (255, 100, 100),
+    alpha: float = 0.5,
+) -> Image.Image:
+    """
+    Create visualization overlay of mask on image.
+
+    Args:
+        image: Original PIL image
+        mask: Binary mask (H, W) with values 0-1
+        color: RGB color for mask overlay
+        alpha: Transparency of overlay (0-1)
+
+    Returns:
+        PIL image with mask overlay
+    """
+    image_np = np.array(image.convert("RGB")).astype(np.float32)
+
+    # Expand mask to 3 channels
+    mask_3ch = np.stack([mask, mask, mask], axis=-1)
+
+    # Create colored overlay
+    color_overlay = np.array(color, dtype=np.float32)
+
+    # Blend: original * (1 - mask*alpha) + color * (mask*alpha)
+    overlay = image_np * (1 - mask_3ch * alpha) + color_overlay * (mask_3ch * alpha)
+    overlay = np.clip(overlay, 0, 255).astype(np.uint8)
+
+    return Image.fromarray(overlay)
+
+
+def draw_points_on_image(
+    image: Image.Image,
+    fg_points: List[Tuple[int, int]],
+    bg_points: Optional[List[Tuple[int, int]]] = None,
+    point_radius: int = 6,
+) -> Image.Image:
+    """
+    Draw foreground (green) and background (red) points on image.
+
+    Args:
+        image: PIL image to draw on
+        fg_points: Foreground points (green)
+        bg_points: Background points (red)
+        point_radius: Radius of point circles
+
+    Returns:
+        PIL image with points drawn
+    """
+    draw_img = image.copy()
+    draw = ImageDraw.Draw(draw_img)
+
+    r = point_radius
+
+    # Draw foreground points (green)
+    for x, y in fg_points:
+        draw.ellipse(
+            [x - r, y - r, x + r, y + r],
+            fill=(0, 255, 0),
+            outline=(0, 180, 0),
+            width=2
+        )
+
+    # Draw background points (red)
+    if bg_points:
+        for x, y in bg_points:
+            draw.ellipse(
+                [x - r, y - r, x + r, y + r],
+                fill=(255, 0, 0),
+                outline=(180, 0, 0),
+                width=2
+            )
+
+    return draw_img

requirements.txt

@@ -10,4 +10,5 @@ omegaconf
 imageio
 gradio
 spaces
-python-dotenv
+python-dotenv
+sam2