Upload weights, notebooks, sample images

README.md

@@ -1,114 +1,108 @@
+---
+license: mit
+tags:
+  - image-to-image
+  - reflection-removal
+  - highlight-removal
+  - computer-vision
+  - dinov3
+  - surgical-imaging
+---
+
 # UnReflectAnything
 
 [![Project](https://img.shields.io/badge/Project-Webpage-ff611b?logo=googlehome&logoColor=ff611b)](https://alberto-rota.github.io/UnReflectAnything/)
-[![PyPI](https://img.shields.io/pypi/v/unreflectanything?color=76b1f3&label=pip%20install&logo=python&logoColor=76b1f3)]
+[![PyPI](https://img.shields.io/pypi/v/unreflectanything?color=76b1f3&label=pip%20install&logo=python&logoColor=76b1f3)](https://pypi.org/project/unreflectanything/)
 [![Paper](https://img.shields.io/badge/Paper-arXiv-B31B1B?logo=arxiv&logoColor=B31B1B)](https://arxiv.org/abs/2512.09583)
 [![Demo](https://img.shields.io/badge/Demo-HF%20-FFD21E?logo=huggingface&logoColor=FFD21E)](https://huggingface.co/spaces/AlbeRota/UnReflectAnything)
 [![Modelcard](https://img.shields.io/badge/Model%20Card-HF%20-FFD21E?logo=huggingface&logoColor=FFD21E)](https://huggingface.co/AlbeRota/UnReflectAnything)
 [![Wiki](https://img.shields.io/badge/API-Wiki-9187FF?logo=wikipedia&logoColor=9187FF)](https://github.com/alberto-rota/UnReflectAnything/wiki)
 [![Licence](https://img.shields.io/badge/MIT-License-1E811F)](https://mit-license.org/)
-### RGB-Only Highlight Removal by Rendering Synthetic Specular Supervision
-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 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. 
 
 ---
-![examples](https://raw.githubusercontent.com/alberto-rota/UnReflectAnything/refs/heads/main/assets/header.png)
 
-## Installation
-```bash
-pip install unreflectanything
-```
-Install UnReflectAnything as a Python Package. 
+## Architecture
+![Architecture](https://raw.githubusercontent.com/alberto-rota/UnReflectAnything/refs/heads/main/assets/architecture.png)
 
-The minimum required Python version is 3.11, but development and all experiments have been bases on **Python 3.12**.
 
-For GPU support, make sure PyTorch comes with CUDA version for your system (see [PyTorch Get Started](https://pytorch.org/get-started/locally/)).
 
+* **<font color="#a001e0">Encoder</font> ($\mathit{\textcolor{a001e0}{E}}$ )**: Processes the input image $\mathbf{I}$ to extract a rich latent representation, $\mathbf{F}_\ell$. This is the off-the-shelf pretrained [DINOv3-large](https://huggingface.co/facebook/dinov3-vitl16-pretrain-lvd1689m) 
 
-## Setting up
-After pip-installing, you can use the `unreflectanything` CLI command, which is also aliased to `unreflect` and `ura`. The three commands are equivalent.
+* **<font color="#0167ff">Reflection Predictor</font> ($\mathit{\textcolor{0167ff}{H}}$ )**: Predicts a soft highlight mask (**H**), identifying areas of specular highlights.
 
-With the CLI you can already download the model weights with
-```bash
-unreflectanything download --weights
-```
-and some sample images with 
-```bash
-unreflectanything download --images
-```
+* **Masking Operation</font> ($\mathit{P}$ )**: A binary mask **P** is derived from the prediction and applied to the feature map: $(1-\mathbf{P}) \odot \mathbf{F}_\ell$. This removes features contaminated by reflections, leaving "holes" in the data.
 
-Weights are stored by default in `~/.cache/unreflectanything/weights` (or `$XDG_CACHE_HOME/unreflectanything/weights` if set ; `%LOCALAPPDATA%\unreflectanything` for Windows). Use `--output-dir` to choose another location. 
+* **<font color="#23ac2c">Token Inpainter</font> ($\mathit{\textcolor{23ac2c}{T}}$ )**: Acts as a neural in-painter. It processes the masked features and uses the surrounding clean context prior and a learned mask token to synthesize the missing information in embedding space, producing the completed feature map $\mathbf{F}_{\text{comp}}$.
 
-Both the weights and images are stored on the [HuggingFace Model Repo](https://huggingface.co/spaces/AlbeRota/UnReflectAnything).
+* **<font color="#ff7700">Decoder</font> ($\mathit{\textcolor{ff7700}{D}}$ )**: Project the completed features back into the pixel space to generate the final, reflection-free image $\mathbf{I}_{\text{diff}}$.
 
-## Enable shell completion
-Shell completion is available for the `bash` and `zsh` shells. Run
-```bash
-unreflectanything completion bash
-```
-and execute the `echo ...` command that gets printed.
+---
 
+## Training Strategy
+We train UnReflectAnything with **Synthetic Specular Supervision** by inferring 3D geometry from [MoGe-2](https://wangrc.site/MoGe2Page/) and rendering highlights with a Blinn-Phong reflection model. We randomly sample the light source position in 3D space at every training iteration enhance etherogeneity.
 
-## Command Line Interface
-Get an overview of the available CLI endpoints with 
+![SupervisionExamples](https://raw.githubusercontent.com/alberto-rota/UnReflectAnything/refs/heads/main/assets/highlights.png)
+
+We train the model in two stages
+1.  **DPT Decoder Pre-Training**: The **<font color="#ff7700">Decoder</font>** is first pre-trained in an autoencoder configuration ($\min_{\theta} \mathcal{L}(M_{\theta}(\mathbf{I}), \mathbf{I})$) to ensure it can reconstruct realistic RGB textures from the DINOV3 latent space.
+2.  **End-to-End Refinement**: The full pipeline is then trained to predict reflection masks from $\mathit{\textcolor{0167ff}{H}}$, and fill them using the **<font color="#38761D">Token Inpainter</font>**, ensuring the final output is both visually consistent and physically accurate. The decoder is also fine-tuned at this stage
+
+
+
+## Weights
+Install the API and CLI on a **Python>=3.11** environment with 
+```bash
+pip install unreflectanything
 ```
-unreflectanything --help   # alias 'unreflect --help' alias 'ura --help'
+then run 
+```bash
+unreflectanything download --weights
 ```
-Refer to the [Wiki](https://github.com/alberto-rota/UnReflectAnything/wiki) to get detailed documentation about each endpoint. We report a summary of the available subcommands. Remember that `ura` is aliased to the `unreflectanything` command
+to download the `.pth` weights in the package cache dir. The cache dir is usually at `.cache/unreflectanything`
 
-| Subcommand | Description | Command |
-|------------|-------------|-------------|
-| `inference` | Run inference on an image directory  |`ura inference --input /path/to/images --output /path/to/unref_images` |
-| `train` | Run training | `ura train --config config_train.yaml`|
-| `test` | Run evaluation on a trained model |`ura test --config config_test.yaml`|
-| `download` | Download checkpoint weights, sample images, notebooks |`ura download --weights`|
-| `verify` | Verify weights installation and compatibility, as well as dataset directory structure | `ura verify --dataset /path/to/dataset`|
-| `evaluate` | Compute metrics on output data | `ura evaluate --output /path/to/unref_images --gt /path/to/groundtruth_images/`|
-| `completion` | Print shell completion (bash/zsh): |`ura completion bash` |
-| `cite` | Print shell completion (bash/zsh)| `ura cite --bibtex` |
-
-## Python API
+---
 
-The same endpoints above are exposed as a Python API. Refer to the [Wiki](https://github.com/alberto-rota/UnReflectAnything/wiki) to get detailed documentation about each endpoint. A few examples are reported below
+### Basic Python Usage
 
 ```python
-import unreflectanything as ura
+import unreflectanything
 import torch
 
-# Get the model class (e.g. for custom setup or training)
-ModelClass = ura.model()
+# Load the pretrained model (uses cached weights)
+unreflect_model = unreflectanything.model() 
 
-# Get a pretrained model (torch.nn.Module) and run on batched RGB
-uramodel = ura.model(pretrained=True)  # uses cached weights; run 'ura download --weights' first
-images = torch.rand(2, 3, 448, 448, device="cuda")  # [B, 3, H, W], values in [0, 1]
-model_out = uramodel(images)  # [B, 3, H, W] diffuse tensor
+# Run inference on a tensor [B, 3, H, W] in range [0, 1]
+images = torch.rand(2, 3, 448, 448).cuda()
+diffuse_output = unreflect_model(images) 
 
-# File-based or tensor-based inference (one-shot, no model handle)
-ura.inference("input.png", output="output.png")
-result = ura.inference(images)  # tensor input returns tensor
+# Simple file-based inference
+unreflectanything.inference("input_with_highlights.png", output="diffuse_result.png")
+```
+Refer to the [Wiki](https://github.com/alberto-rota/UnReflectAnything/wiki) for all details on the API endpoints
 
-# Run training or testing
-ura.run_pipeline(mode="train")   # or mode="test"
+---
 
-# Run inference from options
-options = ura.InferenceOptions(
-    weights_path="path/to/full_model_weights.pt",
-    input_dir="path/to/input/images",
-    output_dir="path/to/output/diffuse",
-)
-ura.run_inference(options)
-```
+### CLI Overview
+
+The package provides a comprehensive command-line interface via `ura`, `unreflect`, or `unreflectanything`.
+
+* **Inference**: `ura inference --input /path/to/images --output /path/to/output`
+* **Evaluation**: `ura evaluate --output /path/to/results --gt /path/to/groundtruth`
+* **Verification**: `ura verify --dataset /path/to/dataset`
+
+Refer to the [Wiki](https://github.com/alberto-rota/UnReflectAnything/wiki) for all details on the CLI endpoints
+
+---
 
 ## Citation
-If you include UnReflectAnything in your pipline or research work, we encourage you cite our work. 
-Get the citation entry with 
-```bash
-unreflectanything cite --bibtex
-```
-or copy it directly from below
-```
+
+If you use UnReflectAnything in your research or pipeline, please cite our paper:
+
+```bibtex
 @misc{rota2025unreflectanythingrgbonlyhighlightremoval,
       title={UnReflectAnything: RGB-Only Highlight Removal by Rendering Synthetic Specular Supervision}, 
       author={Alberto Rota and Mert Kiray and Mert Asim Karaoglu and Patrick Ruhkamp and Elena De Momi and Nassir Navab and Benjamin Busam},
@@ -116,6 +110,9 @@ or copy it directly from below
       eprint={2512.09583},
       archivePrefix={arXiv},
       primaryClass={cs.CV},
-      url={https://arxiv.org/abs/2512.09583}, 
+      url={[https://arxiv.org/abs/2512.09583](https://arxiv.org/abs/2512.09583)}, 
 }
+
 ```
+
+---