Spaces:
Running on Zero
Running on Zero
| """Gradio demo for UnReflectAnything: remove specular reflections from images.""" | |
| from __future__ import annotations | |
| import shutil | |
| import sys | |
| from pathlib import Path | |
| from typing import NamedTuple | |
| # Allow importing unreflectanything when run from gradio_space (e.g. HF Space with root dir) | |
| _REPO_ROOT = Path(__file__).resolve().parent.parent | |
| if _REPO_ROOT not in sys.path: | |
| sys.path.insert(0, str(_REPO_ROOT)) | |
| # Guard against missing '__main__' in worker threads (wandb/pydantic compat) | |
| if "__main__" not in sys.modules: | |
| import types | |
| sys.modules["__main__"] = types.ModuleType("__main__") | |
| _GRADIO_DIR = Path(__file__).resolve().parent | |
| try: | |
| import spaces | |
| except ModuleNotFoundError: | |
| spaces = None | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| from huggingface_hub import hf_hub_download, snapshot_download | |
| HF_REPO = "AlbeRota/UnReflectAnything" | |
| IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp") | |
| class HFAssets(NamedTuple): | |
| """Paths to assets downloaded from the Hugging Face repo.""" | |
| weights_path: str | |
| config_path: str | |
| logo_path: str | |
| sample_images_dir: Path | |
| def _download_from_hf() -> HFAssets: | |
| """Download weights, config, logo, and sample images from the HF repo. Returns paths to all assets.""" | |
| weights_path = hf_hub_download( | |
| repo_id=HF_REPO, | |
| revision="a0d3c7bff5ddb2c430e74b5f8ee67be7cc28fcbe", | |
| filename="weights/full_model_weights.pt", | |
| ) | |
| print("Weights path: ", weights_path) | |
| # config_path = hf_hub_download( | |
| # repo_id=HF_REPO, | |
| # filename="configs/pretrained_config.yaml", | |
| # ) | |
| logo_path = hf_hub_download( | |
| repo_id=HF_REPO, | |
| filename="assets/logo.png", | |
| ) | |
| sample_images_root = Path( | |
| snapshot_download( | |
| repo_id=HF_REPO, | |
| allow_patterns=["sample_images/*"], | |
| ) | |
| ) | |
| sample_images_dir = sample_images_root / "sample_images" | |
| return HFAssets( | |
| weights_path=weights_path, | |
| config_path=Path(__file__).parent / "pretrained_config.yaml", | |
| logo_path=logo_path, | |
| sample_images_dir=sample_images_dir, | |
| ) | |
| _cached_assets: HFAssets | None = None | |
| def _get_assets() -> HFAssets: | |
| """Return HF assets, downloading once and caching.""" | |
| global _cached_assets | |
| if _cached_assets is None: | |
| _cached_assets = _download_from_hf() | |
| return _cached_assets | |
| # Local copy of sample images under cwd so Gradio never needs allowed_paths for examples | |
| _SAMPLE_IMAGES_COPY_DIR: Path | None = None | |
| def _get_sample_image_paths() -> list[str]: | |
| """Return paths of sample images under cwd (copied from HF cache) so Gradio can use them without allowed_paths.""" | |
| global _SAMPLE_IMAGES_COPY_DIR | |
| assets = _get_assets() | |
| src = assets.sample_images_dir | |
| if not src.is_dir(): | |
| return [] | |
| dest = _GRADIO_DIR / "sample_images" | |
| dest.mkdir(parents=True, exist_ok=True) | |
| paths = [] | |
| for p in sorted(src.iterdir()): | |
| if not p.is_file() or p.suffix.lower() not in IMAGE_EXTENSIONS: | |
| continue | |
| dst_file = dest / p.name | |
| if not dst_file.exists() or dst_file.stat().st_mtime < p.stat().st_mtime: | |
| shutil.copy2(p, dst_file) | |
| paths.append(str(dst_file.resolve())) | |
| _SAMPLE_IMAGES_COPY_DIR = dest | |
| return paths | |
| def _get_sample_image_arrays() -> list[np.ndarray]: | |
| """Load sample images as numpy arrays (H, W, 3) uint8 for gr.Examples so the input Image shows a preview.""" | |
| from PIL import Image | |
| paths = _get_sample_image_paths() | |
| arrays = [] | |
| for p in paths: | |
| try: | |
| img = Image.open(p).convert("RGB") | |
| arrays.append(np.array(img)) | |
| except Exception: | |
| continue | |
| return arrays | |
| # Single model instance; loaded in background at app start or on first inference. | |
| _cached_ura_model = None | |
| _cached_device = None | |
| def _get_model(device: str): | |
| """Return the pretrained model, loading it once and moving to the requested device.""" | |
| global _cached_ura_model, _cached_device | |
| assets = _get_assets() | |
| from unreflectanything import model | |
| # If the model isn't loaded yet, initialize it | |
| if _cached_ura_model is None: | |
| print(f"Loading model initially on {device}...") | |
| _cached_ura_model = model( | |
| pretrained=True, | |
| weights_path=assets.weights_path, | |
| # config_path=assets.config_path, | |
| device=device, | |
| verbose=False, | |
| ) | |
| _cached_device = device | |
| # If the model is loaded but on the wrong device, move it | |
| if _cached_device != device: | |
| print(f"Moving model from {_cached_device} to {device}...") | |
| _cached_ura_model.to(device) | |
| _cached_device = device | |
| return _cached_ura_model | |
| def build_ui(): | |
| _get_assets() | |
| # PREVENT: _get_model("cuda") here. It will crash ZeroGPU during startup. | |
| print("UI building... Model will initialize on first inference.") | |
| # Note: Use the decorator directly on the function that does the heavy lifting | |
| def _extract_tokens_nc(tokens) -> torch.Tensor: | |
| """Extract [N, C] from list (last layer) or tensor [B, N, C] (first sample).""" | |
| t = tokens[-1] if isinstance(tokens, list) else tokens | |
| t = t[0].cpu().float() if t.dim() == 3 else t.cpu().float() | |
| return t.squeeze(0) if t.dim() == 3 else t # [N, C] | |
| def _tokens_pair_to_rgb( | |
| tokens_completed, | |
| tokens_input, | |
| h: int, | |
| w: int, | |
| ) -> tuple[np.ndarray, np.ndarray]: | |
| """PCA fit once on completed tokens; apply same mean and V to both; joint min/max norm.""" | |
| from PIL import Image as PILImage | |
| t_comp = _extract_tokens_nc(tokens_completed) # [N, C] | |
| t_inp = _extract_tokens_nc(tokens_input) # [N, C] | |
| mean = t_comp.mean(dim=0, keepdim=True) # [1, C] – fit on completed only | |
| centered_comp = t_comp - mean # [N, C] | |
| U, S, V = torch.svd_lowrank(centered_comp, q=3) # V: [C, 3] | |
| # Project both with same parameters (same mean, same V) | |
| proj_comp = (t_comp - mean) @ V # [N, 3] | |
| proj_inp = (t_inp - mean) @ V # [N, 3] | |
| # Joint min/max so both images share the same color scale | |
| lo = min(proj_comp.min().item(), proj_inp.min().item()) | |
| hi = max(proj_comp.max().item(), proj_inp.max().item()) | |
| eps = 1e-8 | |
| proj_comp = (proj_comp - lo) / (hi - lo + eps) | |
| proj_inp = (proj_inp - lo) / (hi - lo + eps) | |
| grid = int(t_comp.shape[0] ** 0.5) | |
| def to_img(proj: torch.Tensor) -> np.ndarray: | |
| arr = (proj.reshape(grid, grid, 3).numpy() * 255).clip(0, 255).astype(np.uint8) | |
| return np.array(PILImage.fromarray(arr).resize((w, h), PILImage.BILINEAR)) | |
| return to_img(proj_comp), to_img(proj_inp) | |
| def _gray_to_rgb(tensor_1c: torch.Tensor, h: int, w: int) -> np.ndarray: | |
| """Convert [B, 1, H_model, W_model] to resized [H, W, 3] uint8 grayscale-as-RGB.""" | |
| from torchvision.transforms import functional as TF | |
| resized = TF.resize(tensor_1c.cpu(), [h, w], antialias=True) # [B, 1, H, W] | |
| gray = (resized[0, 0].numpy().clip(0.0, 1.0) * 255).astype(np.uint8) # [H, W] | |
| return np.stack([gray] * 3, axis=-1) # [H, W, 3] | |
| def run_inference( | |
| image: np.ndarray | None, | |
| threshold: float = 0.3, | |
| dilation: int = 40, | |
| ) -> dict[str, np.ndarray] | None: | |
| """Run reflection removal; return all visualisable outputs as numpy arrays.""" | |
| if image is None: | |
| return None | |
| from torchvision.transforms import functional as TF | |
| import time | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| ura_model = _get_model(device) | |
| target_side = ura_model.image_size | |
| h, w = image.shape[:2] | |
| tensor = TF.to_tensor(image).unsqueeze(0) # [1, 3, H, W] | |
| tensor = TF.resize(tensor, [target_side, target_side], antialias=True) | |
| tensor = tensor.to(device, dtype=torch.float32) | |
| with torch.no_grad(): | |
| start_time = time.time() | |
| out = ura_model( | |
| images=tensor, | |
| threshold=threshold, | |
| dilation=int(dilation), | |
| return_dict=True, | |
| ) | |
| end_time = time.time() | |
| inference_time_ms = (end_time - start_time) * 1000 | |
| gr.Info(f"Inference complete in {inference_time_ms:.1f} ms") | |
| results: dict[str, np.ndarray] = {} | |
| # Diffuse: [1, 3, H, W] -> [H, W, 3] uint8 | |
| diffuse = TF.resize(out["diffuse"].cpu(), [h, w], antialias=True) | |
| results["diffuse"] = (diffuse[0].numpy().transpose(1, 2, 0).clip(0.0, 1.0) * 255).astype(np.uint8) | |
| # Detected highlight: RGBA overlay superimposed on darkened input | |
| hl_data = out.get("highlight") | |
| if hl_data is not None: | |
| m = TF.resize(hl_data.cpu(), [h, w], antialias=True)[0, 0].numpy().clip(0.0, 1.0) # [H, W] | |
| image_dark = (image.astype(np.float32) * 0.5).clip(0, 255) # [H, W, 3] base | |
| overlay_rgb = np.array([255, 200, 0], dtype=np.float32) # amber | |
| alpha = (0.5 * m)[:, :, np.newaxis] # [H, W, 1] | |
| comp_rgb = (1 - alpha) * image_dark + alpha * overlay_rgb # [H, W, 3] | |
| comp_uint8 = np.clip(comp_rgb, 0, 255).astype(np.uint8) | |
| results["highlight_overlay"] = np.concatenate( | |
| [comp_uint8, np.full((h, w, 1), 255, dtype=np.uint8)], axis=-1 | |
| ) # [H, W, 4] RGBA | |
| results["highlight_gray"] = _gray_to_rgb(hl_data, h, w) | |
| # Highlight mask (binary/dilated) | |
| mask_data = out.get("highlight_mask") | |
| if mask_data is not None: | |
| results["highlight_mask"] = _gray_to_rgb(mask_data, h, w) | |
| # DINOv3 tokens (PCA visualization) – same PCA fit for both, joint color scale | |
| tokens_completed_data = out.get("tokens_completed") | |
| tokens_input_data = out.get("tokens_input") | |
| if tokens_completed_data is not None and tokens_input_data is not None: | |
| img_comp, img_inp = _tokens_pair_to_rgb(tokens_completed_data, tokens_input_data, h, w) | |
| results["tokens_completed"] = img_comp | |
| results["tokens_input"] = img_inp | |
| elif tokens_completed_data is not None: | |
| t = _extract_tokens_nc(tokens_completed_data) | |
| mean = t.mean(dim=0, keepdim=True) | |
| V = torch.svd_lowrank(t - mean, q=3)[2] | |
| proj = (t - mean) @ V | |
| lo, hi = proj.min().item(), proj.max().item() | |
| proj = (proj - lo) / (hi - lo + 1e-8) | |
| grid = int(t.shape[0] ** 0.5) | |
| from PIL import Image as PILImage | |
| arr = (proj.reshape(grid, grid, 3).numpy() * 255).clip(0, 255).astype(np.uint8) | |
| results["tokens_completed"] = np.array(PILImage.fromarray(arr).resize((w, h), PILImage.BILINEAR)) | |
| results["tokens_input"] = results["tokens_completed"] | |
| elif tokens_input_data is not None: | |
| t = _extract_tokens_nc(tokens_input_data) | |
| mean = t.mean(dim=0, keepdim=True) | |
| V = torch.svd_lowrank(t - mean, q=3)[2] | |
| proj = (t - mean) @ V | |
| lo, hi = proj.min().item(), proj.max().item() | |
| proj = (proj - lo) / (hi - lo + 1e-8) | |
| grid = int(t.shape[0] ** 0.5) | |
| from PIL import Image as PILImage | |
| arr = (proj.reshape(grid, grid, 3).numpy() * 255).clip(0, 255).astype(np.uint8) | |
| results["tokens_input"] = np.array(PILImage.fromarray(arr).resize((w, h), PILImage.BILINEAR)) | |
| results["tokens_completed"] = results["tokens_input"] | |
| return results | |
| VIEW_MODES = ["Diffuse", "Highlight", "Inpaint mask", "DINOv3 space"] | |
| def run_inference_slider( | |
| image: np.ndarray | None, | |
| threshold: float, | |
| dilation: int, | |
| ) -> tuple: | |
| """Return 4 slider tuples: (left, right) for each view mode.""" | |
| results = run_inference(image, threshold, dilation) | |
| if results is None: | |
| return (None,) * 4 | |
| diffuse = results["diffuse"] | |
| # Darken input for slider so highlights are more visible; Gradio expects uint8 [0,255] | |
| image_dark = (image.astype(np.float32) * 0.5).clip(0, 255).astype(np.uint8) | |
| hl_overlay = results.get("highlight_overlay", diffuse) | |
| hl_gray = results.get("highlight_gray", diffuse) | |
| hl_mask = results.get("highlight_mask", diffuse) | |
| tok_comp = results.get("tokens_completed", diffuse) | |
| tok_inp = results.get("tokens_input", diffuse) | |
| return ( | |
| (image, diffuse), # Diffuse | |
| (image_dark, hl_overlay), # Detected highlight | |
| (hl_gray, hl_mask), # Highlight mask | |
| (tok_inp, tok_comp), # DINOv3 tokens | |
| ) | |
| assets = _get_assets() | |
| with gr.Blocks(title="UnReflectAnything") as demo: | |
| with gr.Row(elem_classes="mobile-stack"): | |
| with gr.Column(scale=0, min_width=100): | |
| if Path(assets.logo_path).is_file(): | |
| gr.Image( | |
| value=assets.logo_path, | |
| show_label=False, | |
| interactive=False, | |
| height=100, | |
| container=False, | |
| buttons=[], | |
| ) | |
| with gr.Column(scale=1): | |
| gr.Markdown( | |
| """ | |
| # UnReflectAnything | |
| UnReflectAnything inputs any RGB image and **removes specular highlights**, | |
| returning a clean diffuse-only outputs. We trained UnReflectAnything by synthetizing | |
| specularities and supervising in DINOv3 feature space. | |
| UnReflectAnything works on both natural indoor and **surgical/endoscopic** domain data. | |
| Visit the [Project Page](https://alberto-rota.github.io/UnReflectAnything/)! | |
| """ | |
| ) | |
| slider_labels = [ | |
| "Diffuse", | |
| "Highlight", | |
| "Inpaint", | |
| "DINOv3 Space", | |
| ] | |
| with gr.Row(elem_classes="mobile-stack"): | |
| inp = gr.Image( | |
| type="numpy", | |
| label="Input", | |
| height=600, | |
| ) | |
| sliders = [] | |
| for i, lbl in enumerate(slider_labels): | |
| sliders.append( | |
| gr.ImageSlider( | |
| label=lbl, | |
| type="numpy", | |
| height=600, | |
| show_label=True, | |
| visible=(i == 0), | |
| ) | |
| ) | |
| with gr.Row(elem_classes="mobile-stack"): | |
| threshold_slider = gr.Slider( | |
| minimum=0.0, maximum=1.0, value=0.4, step=0.01, | |
| label="Highlight Threshold", | |
| info="Brightness threshold for detecting specular highlights", | |
| ) | |
| dilation_slider = gr.Slider( | |
| minimum=0, maximum=100, value=11, step=1, | |
| label="Mask Dilation", | |
| info="Dilation (px) applied to the detected highlight mask", | |
| ) | |
| view_radio = gr.Radio( | |
| choices=VIEW_MODES, | |
| value=VIEW_MODES[0], | |
| label="Output view", | |
| ) | |
| run_btn = gr.Button("Run UnReflectAnything", variant="primary") | |
| run_btn.click( | |
| fn=run_inference_slider, | |
| inputs=[inp, threshold_slider, dilation_slider], | |
| outputs=sliders, | |
| ) | |
| view_radio.change( | |
| fn=lambda mode: [gr.update(visible=(m == mode)) for m in VIEW_MODES], | |
| inputs=view_radio, | |
| outputs=sliders, | |
| ) | |
| sample_arrays = _get_sample_image_arrays() | |
| if sample_arrays: | |
| gr.Examples( | |
| examples=[[arr] for arr in sample_arrays], | |
| inputs=inp, | |
| label="Pre-loaded examples", | |
| examples_per_page=20, | |
| ) | |
| gr.HTML("""<hr>""") | |
| gr.Markdown(""" | |
| [Project Page](https://alberto-rota.github.io/UnReflectAnything/) ⋅ | |
| [GitHub](https://github.com/alberto-rota/UnReflectAnything) ⋅ | |
| [Model Card](https://huggingface.co/AlbeRota/UnReflectAnything) ⋅ | |
| [Paper](https://arxiv.org/abs/2512.09583) ⋅ | |
| [Contact](mailto:alberto1.rota@polimi.it) | |
| """) | |
| return demo | |
| demo = build_ui() | |
| def _launch_allowed_paths(): | |
| """Paths Gradio is allowed to serve (e.g. for gr.Examples from HF cache).""" | |
| paths = [str(_GRADIO_DIR)] | |
| try: | |
| assets = _get_assets() | |
| sample_dir = assets.sample_images_dir | |
| if sample_dir.is_dir(): | |
| paths.append(str(sample_dir.resolve())) | |
| # Also allow parent (snapshot root) in case Gradio resolves paths from repo root | |
| parent = sample_dir.parent | |
| if parent.is_dir(): | |
| paths.append(str(parent.resolve())) | |
| except Exception as e: | |
| print(f"Warning: could not add HF sample_images to allowed_paths: {e}") | |
| return paths | |
| _MOBILE_CSS = """ | |
| @media (max-width: 768px) { | |
| .mobile-stack { flex-direction: column !important; } | |
| .mobile-stack > .gr-column, | |
| .mobile-stack > div { min-width: 100% !important; } | |
| } | |
| """ | |
| def _launch_kwargs(): | |
| """Default kwargs for launch() so allowed_paths are always set (e.g. when HF Spaces runs demo.launch()).""" | |
| return { | |
| "allowed_paths": _launch_allowed_paths(), | |
| "theme": gr.themes.Soft(primary_hue="orange", secondary_hue="blue"), | |
| "css": _MOBILE_CSS, | |
| } | |
| # Ensure launch() always receives allowed_paths (e.g. when HF Spaces runner calls demo.launch() without args) | |
| _original_launch = demo.launch | |
| def _launch_with_allowed_paths(*args, **kwargs): | |
| for key, value in _launch_kwargs().items(): | |
| if key not in kwargs: | |
| kwargs[key] = value | |
| return _original_launch(*args, **kwargs) | |
| demo.launch = _launch_with_allowed_paths | |
| # Replace your existing launch logic at the very bottom of the file with this: | |
| if __name__ == "__main__": | |
| demo.launch(ssr_mode=True, server_name="0.0.0.0", server_port=7860) | |
| else: | |
| demo.launch(ssr_mode=True, server_name="0.0.0.0", server_port=7860) |