import gradio as gr import os from dotenv import load_dotenv load_dotenv() from PIL import Image import torch import copy from omegaconf import OmegaConf from torchvision.transforms import v2 from torchvision.transforms.functional import to_pil_image from huggingface_hub import hf_hub_download, login import spaces from chord import ChordModel from chord.module import make 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.""" if os.path.isdir(directory): return [[f"{directory}/{f}"] for f in sorted(os.listdir(directory))] return [] EXAMPLES_USECASE_1 = _load_examples("examples/generated") EXAMPLES_USECASE_2 = _load_examples("examples/in_the_wild") EXAMPLES_USECASE_3 = _load_examples("examples/specular") MODEL_OBJ = None # Use local model if available, otherwise download from HuggingFace LOCAL_MODEL_PATH = "chord_v1.safetensors" if os.path.exists(LOCAL_MODEL_PATH): MODEL_CKPT_PATH = LOCAL_MODEL_PATH print(f"Using local model: {MODEL_CKPT_PATH}") else: hf_token = os.environ.get("HF_TOKEN") if not hf_token: raise EnvironmentError( "HF_TOKEN environment variable is required to download the model. " "Set it in a .env file or export it in your shell." ) login(token=hf_token) MODEL_CKPT_PATH = hf_hub_download(repo_id="Ubisoft/ubisoft-laforge-chord", filename="chord_v1.safetensors") print(f"Downloaded model to: {MODEL_CKPT_PATH}") def load_model(ckpt_path): print("Loading model from:", ckpt_path) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") config = OmegaConf.load("config/chord.yaml") model = ChordModel(config) state_dict = load_torch_file(ckpt_path) model.load_state_dict(state_dict) model.eval() model.to(device) return model def run_model(model, img: Image.Image): device = next(model.parameters()).device to_tensor = v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)]) image = to_tensor(img).to(device) x = v2.Resize(size=(1024, 1024), antialias=True)(image).unsqueeze(0) with torch.no_grad(), torch.autocast(device_type=device.type): output = model(x) return output def relit(model, maps): maps['metallic'] = maps.get('metalness', torch.zeros_like(maps['basecolor'])) device = next(model.parameters()).device h, w = maps["basecolor"].shape[-2:] light = make("point-light", {"position": [0, 0, 10]}).to(device) pos = get_positions(h, w, 10).to(device) camera = torch.tensor([0, 0, 10.0]).to(device) for key in maps: if maps[key].dim() == 3: maps[key] = maps[key].unsqueeze(0) maps[key] = maps[key].permute(0,2,3,1) # BxCxHxW -> BxHxWxC 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, output_format, seamless, ao_strength, ao_blur, include_height, height_low_freq, height_mid_freq, height_high_freq, height_intensity, compute_porosity, normalize_porosity, compute_sss, sss_curvature_weight, sss_ao_weight, sss_blur, compute_emission, emission_threshold, emission_knee, emission_bloom, hardcoded_metal, height_mask, metal_mask, ): """ Run Chord model and output shader-compatible textures. Args: output_format: "labpbr" for Java Edition shaders, "bedrock" for Bedrock RTX Returns: albedo: RGB albedo/basecolor texture packed: Specular (_s) for LabPBR or MER/MERS for Bedrock normal: Normal texture (_n for LabPBR, _normal for Bedrock) render: Relit preview image """ global MODEL_OBJ if MODEL_OBJ is None or getattr(MODEL_OBJ, "_ckpt", None) != MODEL_CKPT_PATH: MODEL_OBJ = load_model(MODEL_CKPT_PATH) MODEL_OBJ._ckpt = MODEL_CKPT_PATH # store path inside object if img is None: return None, None, None, None ori_h, ori_w = img.size[1], img.size[0] out = run_model(MODEL_OBJ, img) maps = copy.deepcopy(out) rendered = relit(MODEL_OBJ, maps) resize_back = v2.Resize(size=(ori_h, ori_w), antialias=True) # Resize all maps to original resolution basecolor = resize_back(out["basecolor"]) normal = resize_back(out["normal"]) 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( basecolor=basecolor, normal=normal, roughness=roughness, metalness=metalness, compute_sss=compute_sss, sss_curvature_weight=sss_curvature_weight, sss_ao_weight=sss_ao_weight, sss_blur=int(sss_blur), compute_emission=compute_emission, emission_threshold=emission_threshold, emission_knee=emission_knee, emission_bloom=int(emission_bloom), ) return ( result['albedo'], result['mer'], result['normal'], to_pil_image(resize_back(rendered).squeeze(0)), ) else: # Convert to LabPBR 1.3 format (default) result = convert_to_labpbr( basecolor=basecolor, normal=normal, roughness=roughness, metalness=metalness, derive_ao_height=True, include_height=include_height, height_low_freq=height_low_freq, 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), compute_porosity=compute_porosity, normalize_porosity=normalize_porosity, compute_sss=compute_sss, sss_curvature_weight=sss_curvature_weight, sss_ao_weight=sss_ao_weight, sss_blur=int(sss_blur), compute_emission=compute_emission, emission_threshold=emission_threshold, emission_knee=emission_knee, emission_bloom=int(emission_bloom), hardcoded_metal=hardcoded_metal, metal_mask=metal_mask_tensor, ) return ( result['albedo'], result['specular'], result['normal'], to_pil_image(resize_back(rendered).squeeze(0)), ) with gr.Blocks(title="Chord - PBR Material Estimation") as demo: gr.Markdown("# **Chord: PBR Material Estimation → Minecraft Shader Formats**") gr.Markdown(""" Upload an image to estimate PBR materials and export for Minecraft shaders. **Supported Formats:** - **LabPBR 1.3** (Java Edition): Specular (_s) + Normal (_n) with AO/height - **Bedrock RTX**: MER/MERS (_mer/_mers) + Normal (_normal) """) with gr.Row(): with gr.Column(): input_img = gr.Image(type="pil", label="Input Image", height=512) with gr.Accordion("Output Settings", open=True): gr.Markdown("#### Output Format") output_format = gr.Radio( choices=[("LabPBR 1.3 (Java Edition)", "labpbr"), ("Bedrock RTX", "bedrock")], value="labpbr", label="Format", info="Choose shader format for your Minecraft edition" ) gr.Markdown("#### General (LabPBR only)") seamless = gr.Checkbox(label="Seamless/Tileable", value=False, info="Enable for textures that should tile seamlessly") gr.Markdown("#### AO & Height Derivation (LabPBR only)") ao_strength = gr.Slider(minimum=0.5, maximum=5.0, value=2.0, step=0.1, label="AO Strength", info="AO contrast multiplier") ao_blur = gr.Slider(minimum=0, maximum=15, value=5, step=1, label="AO Blur", info="Gaussian blur radius for AO smoothing") include_height = gr.Checkbox(label="Include Height Map (for POM)", value=True, info="Derive height from normal map for parallax occlusion mapping. Disable for flat surfaces.") gr.Markdown("#### Height Relief Balance (LabPBR only)") height_low_freq = gr.Slider(minimum=0.0, maximum=2.0, value=1.0, step=0.01, label="Low Frequencies", info="Large-scale height variations") height_mid_freq = gr.Slider(minimum=0.0, maximum=2.0, value=1.0, step=0.01, label="Mid Frequencies", info="Medium surface features") height_high_freq = gr.Slider(minimum=0.0, maximum=2.0, value=1.0, step=0.01, label="High Frequencies", info="Fine surface detail") height_intensity = gr.Slider(minimum=0.0, maximum=1.0, value=1.0, step=0.01, label="Global Intensity", info="Overall height map strength (0 = flat)") gr.Markdown("#### Porosity (LabPBR Blue 0-64)") compute_porosity = gr.Checkbox(label="Compute Porosity", value=False, info="Calculate porosity from AO, smoothness, and F0 (LabPBR only)") normalize_porosity = gr.Checkbox(label="Normalize Porosity", value=True, info="Normalize to full range before LabPBR scaling") gr.Markdown("#### Subsurface Scattering") compute_sss = gr.Checkbox(label="Compute SSS", value=False, info="Calculate SSS thickness from normal curvature (LabPBR: blue 65-255, Bedrock: MERS alpha)") sss_curvature_weight = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, step=0.05, label="Curvature Weight", info="Weight for curvature contribution") sss_ao_weight = gr.Slider(minimum=0.0, maximum=1.0, value=0.3, step=0.05, label="AO Weight", info="Weight for inverted AO contribution") sss_blur = gr.Slider(minimum=0, maximum=10, value=2, step=1, label="SSS Blur", info="Gaussian blur for soft SSS look") gr.Markdown("#### Emission") compute_emission = gr.Checkbox(label="Compute Emission", value=False, info="Extract emission from bright areas of basecolor") emission_threshold = gr.Slider(minimum=0.5, maximum=1.0, value=0.85, step=0.01, label="Emission Threshold", info="Luminance threshold for detection") emission_knee = gr.Slider(minimum=0.0, maximum=0.3, value=0.1, step=0.01, label="Emission Knee", info="Soft knee width (0 = hard threshold)") emission_bloom = gr.Slider(minimum=0, maximum=21, value=0, step=1, label="Emission Bloom", info="Gaussian blur radius for bloom effect (0 = disabled)") gr.Markdown("#### Hardcoded Metal (LabPBR only)") hardcoded_metal = gr.Dropdown( choices=LABPBR_METAL_CHOICES, value="none", label="Metal Type", 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, inputs=[input_img], label="Examples (Generated Textures)" ) gr.Markdown("### Example Inputs — In The Wild Photographs") gr.Examples( examples=EXAMPLES_USECASE_2, inputs=[input_img], label="Examples (In The Wild Photographs)" ) gr.Markdown("### Example Inputs — Specular Textures") gr.Examples( examples=EXAMPLES_USECASE_3, inputs=[input_img], label="Examples (Specular Textures)" ) run_button = gr.Button("Run Estimation", variant="primary") with gr.Column(): gr.Markdown("### Output Textures") albedo_out = gr.Image(label="Albedo (basecolor)", height=340, format="png") packed_out = gr.Image(label="Specular/MER — LabPBR: _s (smoothness,F0,porosity,emission) | Bedrock: _mer (M,E,R,[S])", height=340, format="png") normal_out = gr.Image(label="Normal — LabPBR: _n (XY,AO,height) | Bedrock: _normal (XYZ DirectX)", height=340, format="png") 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=[ input_img, output_format, seamless, ao_strength, ao_blur, include_height, height_low_freq, height_mid_freq, height_high_freq, height_intensity, compute_porosity, normalize_porosity, compute_sss, sss_curvature_weight, sss_ao_weight, sss_blur, compute_emission, emission_threshold, emission_knee, emission_bloom, hardcoded_metal, current_mask_state, metal_combined_mask_state, ], outputs=[albedo_out, packed_out, normal_out, render_out] ) if __name__ == "__main__": demo.launch()