Upload 13 files

README.md

@@ -1,14 +1,130 @@
----
-title: RADAR Demo
-emoji: 🐨
-colorFrom: pink
-colorTo: pink
-sdk: gradio
-sdk_version: 5.49.1
-app_file: app.py
-pinned: false
-license: apache-2.0
-short_description: Demo release for RADAR.
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+<h1 align="center"> Towards Reliable Identification of Diffusion-based Image Manipulations (NeurIPS 2025) </h1> 
+
+:rotating_light: This repository contains code snippets and checkpoints of our work "**Towards Reliable Identification of Diffusion-based Image Manipulations**" :rotating_light:
+ 
+by [Alex Costanzino](https://alex-costanzino.github.io/)*1, [Woody Bayliss](https://www.linkedin.com/in/woody-bayliss-750a27325/)*2, [Juil Sock](https://skf0321.wixsite.com/juil)*2, [Marc Gorriz Blanch](https://www.linkedin.com/in/marc-gorriz/)*2, [Danijela Horak](https://www.linkedin.com/in/danijela-horak-3b417a70/?originalSubdomain=uk)*2, [Ivan Laptev](https://mbzuai.ac.ae/study/faculty/ivan-laptev/)*3, [Philip Tor](https://eng.ox.ac.uk/people/philip-torr/)*4 and [Fabio Pizzati](https://fabvio.github.io/)*3.
+
+*1 University of Bologna, *2 BBC R&D, *3 MBZUAI, *4 University of Oxford
+
+
+<div class="alert alert-info">
+
+
+<h2 align="center"> 
+
+[Project Page](https://alex-costanzino.github.io/radar/) | [Paper (ArXiv)](https://www.arxiv.org/abs/2506.05466)
+</h2>
+
+
+## :bookmark_tabs: Table of Contents
+
+1. [Introduction](#clapper-introduction)
+2. [Dataset](#file_cabinet)
+3. [Checkpoints](#inbox_tray)
+4. [Code](#memo-code)
+6. [Contacts](#envelope-contacts)
+
+</div>
+
+## :clapper: Introduction
+
+Changing facial expressions, gestures, or background details may dramatically alter the meaning conveyed by an image. Notably, recent advances in diffusion models greatly improve the quality of image manipulation while also opening the door to misuse. Identifying changes made to authentic images, thus, becomes an important task, constantly challenged by new diffusion-based editing tools.
+
+To this end, we propose a novel approach for ReliAble iDentification of inpainted AReas (RADAR). RADAR builds on existing foundation models and combines features from different image modalities. It also incorporates an auxiliary contrastive loss that helps to isolate manipulated image patches.
+
+We demonstrate these techniques to significantly improve both the accuracy of our method and its generalisation to a large number of diffusion models. To support realistic evaluation, we further introduce BBC-PAIR, a new comprehensive benchmark, with images tampered by 28 diffusion models. Our experiments show that RADAR achieves excellent results, outperforming the state-of-the-art in detecting and localising image edits made by both seen and unseen diffusion models.
+
+<h4 align="center">
+
+</h4>
+
+<img src="./assets/architecture.jpg" alt="Alt text" style="width: 800px;" title="architecture">
+
+:fountain_pen: If you find this code useful in your research, please cite:
+
+```bibtex
+@article{costanzino2025radar,
+  author    = {Costanzino, Alex and Bayliss, Woody and Sock, Juil and Gorriz Blanch, Marc and Horak, Danijela and Laptev, Ivan and Torr, Philip and Pizzati, Fabio},
+  title     = {Towards Reliable Identification of Diffusion-based Image Manipulations},
+  journal   = {arXiv},
+  year      = {2025},
+}
+```
+
+<h2 id="file_cabinet"> :file_cabinet: Dataset </h2>
+
+In our experiments, we employed our benchmark: [BBC - PAIR (Paired Authentic and Inpainted References)](https://github.com/bbc/PAIR).
+
+
+<h2 id="inbox_tray"> :inbox_tray: Checkpoints </h2>
+
+Here, you can download the weights of the networks employed in the results our paper.
+
+To use these weights, please follow these steps:
+
+1. Create a folder named `checkpoints` in the project directory;
+2. Download the weights [[Download]](https://drive.google.com/drive/folders/1VYbLoK3HM238AU1_xszBM4DYd2nai0jf?usp=sharing);
+3. Copy the downloaded weights into the `checkpoints` folder.
+
+
+## :memo: Code
+
+<div class="alert alert-info">
+
+**Warning**:
+- The code utilizes `wandb` during training to log results. Please be sure to have a wandb account. Otherwise, if you prefer to not use `wandb`, disable it in `A_train.py` with the flag `mode = 'disabled'`.
+
+</div>
+
+### :hammer_and_wrench: Setup Instructions
+
+**Dependencies**: Ensure that you have installed all the necessary dependencies. The list of dependencies can be found in the `requirements.txt` file.
+
+### :rocket: Single-image Inference
+
+The `C_infer_single-image.py` run the framework on an image to generate a continous tampering map.
+
+You can specify the following arguments:
+   - `--checkpoint_savepath`: Path to the directory of the checkpoints, i.e., `checkpoints`;
+   - `--label`: A label to tell apart the experiments, i.e., `multimodal_15k_10inp`;
+   - `--image_path`: Path to the image to assess;
+   - `--epochs_no`: Number of epochs employed during the framework optimization.
+   - `--img_size`: Resolution employed during the framework optimization;
+   - `--batch_size`: Number of samples per batch employed during the framework optimization.
+
+If you haven't downloaded the checkpoints yet, you can find the download links in the **Checkpoints** section above.
+
+### :rocket: Full Inference [TBA]
+
+<!-- The `B_infer.py` script test the framework. It can be used to generate anomaly maps as well.
+
+You can specify the following options:
+   - `--dataset_path`: Path to the root directory of the dataset.
+   - `--checkpoint_folder`: Path to the directory of the checkpoints, i.e., `checkpoints/checkpoints_visa`.
+   - `--class_name`: Class on which the framework was trained, i.e., `candle`.
+   - `--epochs_no`: Number of epochs used in framework optimization.
+   - `--batch_size`: Number of samples per batch employed for framework optimization.
+   - `--qualitative_folder`: Folder on which the anomaly maps are saved.
+   - `--quantitative_folder`: Folder on which the metrics are saved.
+   - `--visualize_plot`: Flag to visualize qualitatived during inference.
+   - `--produce_qualitatives`: Flag to save qualitatived during inference.
+
+If you haven't downloaded the checkpoints yet, you can find the download links in the **Checkpoints** section above. -->
+
+### :rocket: Train [TBA]
+
+<!-- The `A_train.py` script train the framework.
+
+You can specify the following options:
+   - `--dataset_path`: Path to the root directory of the dataset.
+   - `--checkpoint_savepath`: Path to the directory on which checkpoints will be saved, i.e., `checkpoints/checkpoints_visa`.
+   - `--class_name`: Class on which the framework is trained, i.e., `candle`.
+   - `--epochs_no`: Number of epochs for framework optimization.
+   - `--img_size`: Resolution employed during training.
+   - `--batch_size`: Number of samples per batch for framework optimization.
+   - `--label`: A label to tell apart the experiments. -->
+
+## :envelope: Contacts
+
+For questions, please send an email to [email protected].

app.py

@@ -0,0 +1,250 @@
+import gradio as gr
+import torch
+import torchvision
+from PIL import Image
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib
+matplotlib.use('Agg')
+import io
+import os
+
+# Import your models
+from models.feature_extractor import FeatureExtractor, FeatureExtractorDepth
+from models.projector import SiameseProjector
+from models.fuser import DoubleCrossAttentionFusion
+from loaders.loader_utils import SquarePad
+
+# Configuration
+CHECKPOINT_PATH = './checkpoints'
+MODEL_LABEL = 'multimodal_15k_10inp'
+EPOCHS = 120
+BATCH_SIZE = 4
+IMAGE_SIZE = 896
+
+# Load models
+print("Loading models...")
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
+model_name = f'{MODEL_LABEL}_{EPOCHS}ep_{BATCH_SIZE}bs'
+
+rgb_transform = torchvision.transforms.Compose([
+    SquarePad(),
+    torchvision.transforms.Resize((IMAGE_SIZE, IMAGE_SIZE), 
+                                  interpolation=torchvision.transforms.InterpolationMode.BICUBIC),
+    torchvision.transforms.ToTensor(),
+    torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], 
+                                    std=[0.229, 0.224, 0.225]),
+    torchvision.transforms.Lambda(lambda img: img.unsqueeze(0)),
+])
+
+fe_rgb = FeatureExtractor().to(device).eval()
+fe_depth = FeatureExtractorDepth().to(device).eval()
+
+fusion_block = DoubleCrossAttentionFusion(hidden_dim=fe_rgb.embed_dim).to(device)
+fusion_block.load_state_dict(torch.load(
+    os.path.join(CHECKPOINT_PATH, f'fusion_block_{model_name}.pth'), 
+    weights_only=False,
+    map_location=device
+))
+fusion_block.eval()
+
+projector = SiameseProjector(inner_features=fe_rgb.embed_dim).to(device)
+projector.load_state_dict(torch.load(
+    os.path.join(CHECKPOINT_PATH, f'projector_{model_name}.pth'), 
+    weights_only=False,
+    map_location=device
+))
+projector.eval()
+
+print("Models loaded successfully!")
+
+def detect_manipulation(image):
+    """Process image and return heatmap"""
+    if image is None:
+        return None
+    
+    # Convert to PIL if needed
+    if isinstance(image, np.ndarray):
+        rgb_input = Image.fromarray(image).convert('RGB')
+    else:
+        rgb_input = image.convert('RGB')
+    
+    original_size = rgb_input.size
+    
+    # Transform and process
+    rgb = rgb_transform(rgb_input)
+    rgb = rgb.to(device)
+    
+    with torch.no_grad():
+        rgb_feat = fe_rgb(rgb)
+        depth_feat = fe_depth(rgb)
+        fused_feat = fusion_block(rgb_feat, depth_feat)
+        _, segmentation_map = projector(fused_feat)
+        segmentation_map = torch.sigmoid(segmentation_map)
+        
+        # Resize back to original
+        segmentation_map = torch.nn.functional.interpolate(
+            segmentation_map, 
+            size=[max(original_size), max(original_size)], 
+            mode='bilinear'
+        ).squeeze()
+        segmentation_map = torchvision.transforms.functional.center_crop(
+            segmentation_map, 
+            original_size[::-1]
+        )
+        
+        heatmap = segmentation_map.cpu().detach().numpy()
+    
+    # Create visualization
+    fig, ax = plt.subplots(figsize=(10, 10))
+    ax.imshow(heatmap, cmap='jet')
+    ax.axis('off')
+    plt.tight_layout(pad=0)
+    
+    # Convert to image
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0, dpi=150)
+    buf.seek(0)
+    result_image = Image.open(buf)
+    plt.close(fig)
+    
+    return result_image
+
+# Custom CSS to match your design
+custom_css = """
+#component-0 {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    padding: 20px;
+}
+
+.gradio-container {
+    max-width: 1200px !important;
+    margin: auto !important;
+}
+
+#title {
+    text-align: center;
+    color: white;
+    font-size: 3em;
+    margin-bottom: 10px;
+    text-shadow: 2px 2px 4px rgba(0,0,0,0.2);
+}
+
+#subtitle {
+    text-align: center;
+    color: white;
+    font-size: 1.2em;
+    margin-bottom: 30px;
+    opacity: 0.95;
+}
+
+.main-card {
+    background: white;
+    border-radius: 20px;
+    padding: 40px;
+    box-shadow: 0 20px 60px rgba(0,0,0,0.3);
+}
+
+#info-box {
+    background: #e8f4fd;
+    border-left: 4px solid #2196F3;
+    padding: 15px;
+    border-radius: 5px;
+    margin-bottom: 20px;
+    color: #1976D2;
+}
+
+.upload-container .transition {
+    border: 3px dashed #667eea !important;
+    border-radius: 15px !important;
+    background: #f8f9ff !important;
+}
+
+.upload-container .transition:hover {
+    border-color: #764ba2 !important;
+    background: #f0f2ff !important;
+}
+
+button.primary {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
+    border: none !important;
+    color: white !important;
+    border-radius: 25px !important;
+    padding: 12px 30px !important;
+    font-size: 1em !important;
+}
+
+button.primary:hover {
+    transform: translateY(-2px);
+    box-shadow: 0 5px 15px rgba(102, 126, 234, 0.4) !important;
+}
+
+.output-image {
+    border-radius: 15px;
+    box-shadow: 0 5px 15px rgba(0,0,0,0.1);
+}
+"""
+
+# Create Gradio interface
+with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
+    gr.HTML("""
+        <div id="title">🎯 RADAR</div>
+        <div id="subtitle">ReliAble iDentification of inpainted AReas</div>
+    """)
+    
+    with gr.Column(elem_classes="main-card"):
+        gr.HTML("""
+            <div id="info-box">
+                <strong>ℹ️ About RADAR:</strong> Upload an image to detect and localize regions 
+                that have been manipulated using diffusion-based inpainting models. 
+                The output shows a heatmap where red areas indicate detected manipulations.
+            </div>
+        """)
+        
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(
+                    label="Upload Image",
+                    type="pil",
+                    elem_classes="upload-container"
+                )
+                submit_btn = gr.Button("🔍 Detect Manipulations", variant="primary", size="lg")
+            
+            with gr.Column():
+                output_image = gr.Image(
+                    label="Manipulation Heatmap",
+                    type="pil",
+                    elem_classes="output-image"
+                )
+        
+        gr.Examples(
+            examples=[
+                # Add paths to your example images here
+                # ["examples/example1.png"],
+                # ["examples/example2.png"],
+            ],
+            inputs=input_image,
+            outputs=output_image,
+            fn=detect_manipulation,
+            cache_examples=False,
+        )
+    
+    # Connect the button
+    submit_btn.click(
+        fn=detect_manipulation,
+        inputs=input_image,
+        outputs=output_image
+    )
+    
+    # Also trigger on image upload
+    input_image.change(
+        fn=detect_manipulation,
+        inputs=input_image,
+        outputs=output_image
+    )
+
+# Launch
+if __name__ == "__main__":
+    demo.launch()

checkpoints/fusion_block_multimodal_15k_10inp_120ep_4bs.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:70601e0d63edddbfe42da48674dc9053dc626176c5fcda273a2018b10b80bd10
+size 23646170

checkpoints/projector_multimodal_15k_10inp_120ep_4bs.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e5ec7007ebeae7f08054de2ab223899386bb1b977e028545175241efae8ce388
+size 14208328

loaders/dataloader.py

@@ -0,0 +1,167 @@
+import os
+import torch
+from PIL import Image
+from torchvision import transforms
+import torchvision
+
+import json
+
+from torch.utils.data import Dataset
+from torch.utils.data import DataLoader
+
+class SquarePad:
+    def __call__(self, image):
+        max_wh = max(image.size)
+        p_left, p_top = [(max_wh - s) // 2 for s in image.size]
+        p_right, p_bottom = [max_wh - (s+pad) for s, pad in zip(image.size, [p_left, p_top])]
+        padding = (p_left, p_top, p_right, p_bottom)
+        return transforms.functional.pad(image, padding, padding_mode = 'edge')
+
+class BaseDataset(Dataset):
+    def __init__(self, img_size, dataset_path, inpainter):
+        self.IMAGENET_MEAN = [0.485, 0.456, 0.406]
+        self.IMAGENET_STD = [0.229, 0.224, 0.225]
+
+        self.img_size = img_size
+        self.dataset_path = dataset_path
+        self.inpainter = inpainter
+
+        self.json_path = os.path.join(dataset_path, 'DFDS_V2/DFDS_V2.0_2Percent.json')
+        # self.json_path = os.path.join(dataset_path, 'DFDS_V2.0_2Percent.json')
+
+        self.data = self.load_json()
+        self.data_train = self.data[0:500]
+
+        self.rgb_transform = transforms.Compose([
+            SquarePad(),
+            transforms.Resize((img_size, img_size), interpolation = transforms.InterpolationMode.BICUBIC),
+            transforms.ToTensor(),
+            transforms.Normalize(mean = self.IMAGENET_MEAN, std = self.IMAGENET_STD)
+            ])
+
+    def load_json(self):
+        with open(self.json_path, 'r') as file:
+            data = json.load(file)
+        return data
+
+class TrainDataset(BaseDataset):
+    def __init__(self, img_size, dataset_path):
+        super().__init__(img_size = img_size, dataset_path = dataset_path, inpainter = None)
+
+        self.gt_transform = transforms.Compose([
+            SquarePad(),
+            transforms.Resize((img_size, img_size), interpolation = transforms.InterpolationMode.BICUBIC),
+            transforms.ToTensor()]
+            )
+
+        self.img_paths_pos, self.img_paths_neg, self.mask_paths_neg = self.load_dataset()
+
+    def load_dataset(self):
+        positive_imgs = [os.path.join(self.dataset_path, data['base_image_location'].lstrip('/')) for data in self.data_train] + \
+                        [os.path.join(self.dataset_path, data['base_image_location'].lstrip('/')) for data in self.data_train] + \
+                        [os.path.join(self.dataset_path, data['base_image_location'].lstrip('/')) for data in self.data_train] + \
+                        [os.path.join(self.dataset_path, data['base_image_location'].lstrip('/')) for data in self.data_train] + \
+                        [os.path.join(self.dataset_path, data['base_image_location'].lstrip('/')) for data in self.data_train] + \
+                        [os.path.join(self.dataset_path, data['base_image_location'].lstrip('/')) for data in self.data_train] + \
+                        [os.path.join(self.dataset_path, data['base_image_location'].lstrip('/')) for data in self.data_train] + \
+                        [os.path.join(self.dataset_path, data['base_image_location'].lstrip('/')) for data in self.data_train] + \
+                        [os.path.join(self.dataset_path, data['base_image_location'].lstrip('/')) for data in self.data_train] + \
+                        [os.path.join(self.dataset_path, data['base_image_location'].lstrip('/')) for data in self.data_train]
+
+        negative_imgs = [os.path.join(self.dataset_path, data['masks'][0]['inpainters']['SD1_Inpaint'].lstrip('/')) for data in self.data_train] + \
+                        [os.path.join(self.dataset_path, data['masks'][0]['inpainters']['SD1.5_Inpaint'].lstrip('/')) for data in self.data_train] + \
+                        [os.path.join(self.dataset_path, data['masks'][0]['inpainters']['SD2_Inpaint'].lstrip('/')) for data in self.data_train] + \
+                        [os.path.join(self.dataset_path, data['masks'][0]['inpainters']['SDXL_Inpaint'].lstrip('/')) for data in self.data_train] + \
+                        [os.path.join(self.dataset_path, data['masks'][0]['inpainters']['SD3_Inpaint'].lstrip('/')) for data in self.data_train] + \
+                        [os.path.join(self.dataset_path, data['masks'][0]['inpainters']['SD3.5_Inpaint'].lstrip('/')) for data in self.data_train] + \
+                        [os.path.join(self.dataset_path, data['masks'][0]['inpainters']['kadinsky2.2_Inpaint'].lstrip('/')) for data in self.data_train] + \
+                        [os.path.join(self.dataset_path, data['masks'][0]['inpainters']['kadinsky3.1_Inpaint'].lstrip('/')) for data in self.data_train] + \
+                        [os.path.join(self.dataset_path, data['masks'][0]['inpainters']['FLUX_SHNELL_Inpaint'].lstrip('/')) for data in self.data_train] + \
+                        [os.path.join(self.dataset_path, data['masks'][0]['inpainters']['FLUX_DEV_FILL'].lstrip('/')) for data in self.data_train]
+
+        negative_masks = [os.path.join(self.dataset_path, data['masks'][0]['edited_mask_location'].lstrip('/')) for data in self.data_train] + \
+                         [os.path.join(self.dataset_path, data['masks'][0]['edited_mask_location'].lstrip('/')) for data in self.data_train] + \
+                         [os.path.join(self.dataset_path, data['masks'][0]['edited_mask_location'].lstrip('/')) for data in self.data_train] + \
+                         [os.path.join(self.dataset_path, data['masks'][0]['edited_mask_location'].lstrip('/')) for data in self.data_train] + \
+                         [os.path.join(self.dataset_path, data['masks'][0]['edited_mask_location'].lstrip('/')) for data in self.data_train] + \
+                         [os.path.join(self.dataset_path, data['masks'][0]['edited_mask_location'].lstrip('/')) for data in self.data_train] + \
+                         [os.path.join(self.dataset_path, data['masks'][0]['edited_mask_location'].lstrip('/')) for data in self.data_train] + \
+                         [os.path.join(self.dataset_path, data['masks'][0]['edited_mask_location'].lstrip('/')) for data in self.data_train] + \
+                         [os.path.join(self.dataset_path, data['masks'][0]['edited_mask_location'].lstrip('/')) for data in self.data_train] + \
+                         [os.path.join(self.dataset_path, data['masks'][0]['edited_mask_location'].lstrip('/')) for data in self.data_train]
+
+        return positive_imgs, negative_imgs, negative_masks
+
+    def __len__(self):
+        return len(self.data_train) * 3
+
+    def __getitem__(self, idx):
+        img_path_pos, img_path_neg, gt_neg = self.img_paths_pos[idx], self.img_paths_neg[idx], self.mask_paths_neg[idx]
+
+        img_pos = Image.open(img_path_pos.replace('/Open_V7/','')).convert('RGB')
+        img_neg = Image.open(img_path_neg.replace('/Open_V7/','')).convert('RGB')
+
+        rgb_pos = self.rgb_transform(img_pos)
+        rgb_neg = self.rgb_transform(img_neg)
+
+        gt_pos = torch.zeros([1, img_pos.size[1], img_pos.size[0]])
+        gt_pos = torchvision.transforms.functional.to_pil_image(gt_pos)
+        gt_pos = self.gt_transform(gt_pos)
+
+        gt_neg = Image.open(gt_neg.replace('/Open_V7/','')).convert('L')
+        gt_neg = self.gt_transform(gt_neg)
+        gt_neg = torch.where(gt_neg > 0.5, 1., .0)
+
+        return rgb_pos, gt_pos, rgb_neg, gt_neg
+
+class TestDataset(BaseDataset):
+    def __init__(self, img_size, dataset_path, inpainter):
+        super().__init__(img_size = img_size, dataset_path = dataset_path, inpainter = inpainter)
+
+        self.gt_transform = transforms.Compose([
+            transforms.ToTensor()])
+
+        self.img_paths, self.mask_paths, self.labels = self.load_dataset()
+
+    def load_dataset(self):
+        positive_imgs = [os.path.join(self.dataset_path, data['base_image_location'].lstrip('/')) for data in self.data[500:600]]
+        positive_masks = [None for data in self.data[500:600]]
+
+        negative_imgs = [os.path.join(self.dataset_path, data['masks'][0]['inpainters'][self.inpainter].lstrip('/')) for data in self.data[600:700]]
+        negative_masks = [os.path.join(self.dataset_path, data['masks'][0]['edited_mask_location'].lstrip('/')) for data in self.data[600:700]]
+
+        labels = [0.0 for data in self.data[500:600]] + [1.0 for data in self.data[600:700]]
+
+        return positive_imgs + negative_imgs, positive_masks + negative_masks, labels
+
+    def __len__(self):
+        return len(self.data[500:700])
+
+    def __getitem__(self, idx):
+        img_path, gt, label = self.img_paths[idx], self.mask_paths[idx], self.labels[idx]
+
+        img = Image.open(img_path.replace('/Open_V7/','').replace('data/', 'data/DFDS_V2/')).convert('RGB')
+
+        rgb = self.rgb_transform(img)
+
+        if gt == None:
+            gt = torch.zeros(
+                [1, img.size[1], img.size[0]])
+        else:
+            gt = Image.open(gt.replace('/Open_V7/','').replace('data/', 'data/DFDS_V2/')).convert('L')
+            gt = self.gt_transform(gt)
+            gt = torch.where(gt > 0.5, 1., .0)
+
+        return rgb, label, gt, img_path
+
+
+def get_data_loader(split, img_size, batch_size, dataset_path, inpainter = None):
+    if split == 'train':
+        dataset = TrainDataset(img_size, dataset_path)
+        data_loader = DataLoader(dataset = dataset, batch_size = batch_size, shuffle = True, num_workers = 8, drop_last = True, pin_memory = False)
+
+    elif split == 'test':
+        dataset = TestDataset(img_size, dataset_path, inpainter)
+        data_loader = DataLoader(dataset = dataset, batch_size = batch_size, shuffle = False, num_workers = 8, drop_last = False, pin_memory = False)
+
+    return data_loader

loaders/loader_utils.py

@@ -0,0 +1,14 @@
+from torchvision import transforms
+
+imagenet_denormalize = transforms.Compose([
+    transforms.Normalize(mean = [0., 0., 0.], std = [1/0.229, 1/0.224, 1/0.225]),
+    transforms.Normalize(mean = [-0.485, -0.456, -0.406], std = [1., 1., 1.])
+    ])
+
+class SquarePad:
+    def __call__(self, image):
+        max_wh = max(image.size)
+        p_left, p_top = [(max_wh - s) // 2 for s in image.size]
+        p_right, p_bottom = [max_wh - (s+pad) for s, pad in zip(image.size, [p_left, p_top])]
+        padding = (p_left, p_top, p_right, p_bottom)
+        return transforms.functional.pad(image, padding, padding_mode = 'edge')

models/feature_extractor.py

@@ -0,0 +1,27 @@
+import torch
+from transformers import AutoModelForDepthEstimation
+
+class FeatureExtractor(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.fe = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14_reg')
+
+        self.patch_size = self.fe.patch_size
+        self.embed_dim = self.fe.embed_dim
+
+    def forward(self, x):
+        return self.fe.forward_features(x)['x_norm_patchtokens']
+    
+class FeatureExtractorDepth(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.fe = AutoModelForDepthEstimation.from_pretrained("depth-anything/Depth-Anything-V2-Base-hf")
+
+        self.patch_size = 14
+        self.embed_dim = 768
+
+    def forward(self, x):
+        x = self.fe(x, output_hidden_states=True).hidden_states
+        return x[-1][:,1:,:]

models/fuser.py

@@ -0,0 +1,65 @@
+import torch
+
+class DoubleCrossAttentionFusion(torch.nn.Module):
+    def __init__(self, hidden_dim=768, num_heads=8, dropout=0.1):
+        super().__init__()
+        
+        # 1. Per-modality normalization.
+        self.norm_rgb = torch.nn.LayerNorm(hidden_dim)
+        self.norm_depth = torch.nn.LayerNorm(hidden_dim)
+        
+        # 2. Cross-attention.
+        self.cross_attn_depth = torch.torch.nn.MultiheadAttention(
+            embed_dim=hidden_dim,
+            num_heads=num_heads,
+            dropout=dropout,
+            batch_first=True,
+        )
+
+        self.cross_attn_rgb = torch.torch.nn.MultiheadAttention(
+            embed_dim=hidden_dim,
+            num_heads=num_heads,
+            dropout=dropout,
+            batch_first=True,
+        )
+
+        # 3. Mixing.
+        self.mixer = torch.nn.Sequential(
+            torch.nn.Linear(hidden_dim * 2, hidden_dim),
+            torch.nn.GELU(),
+            torch.nn.Dropout(dropout)
+        )
+
+        # 4. Output normalisation.
+        self.out_norm = torch.nn.LayerNorm(hidden_dim)
+
+    def forward(self, rgb_features, depth_features):
+        # 1. Normalize inputs.
+        rgb = self.norm_rgb(rgb_features)
+        depth = self.norm_depth(depth_features)
+        
+        # 2a. Cross-attention (depth -> rgb).
+        attn_out_depth, _ = self.cross_attn_depth(
+            query=depth,
+            key=rgb,
+            value=rgb,
+            need_weights=False
+        )
+
+        # 2b. Cross-attention (rgb -> depth).
+        attn_out_rgb, _ = self.cross_attn_rgb(
+            query=rgb,
+            key=depth,
+            value=depth,
+            need_weights=False
+        )
+
+        # 3a. Residuals.
+        depth_attn = depth + attn_out_depth
+        rgb_attn = rgb + attn_out_rgb
+
+        # 3b. Mixing.
+        fused = self.mixer(torch.cat([depth_attn, rgb_attn], dim=-1))
+
+        # 4. Output normalisation.
+        return self.out_norm(fused)

models/losses.py

@@ -0,0 +1,105 @@
+"""
+Author: Yonglong Tian ([email protected])
+Date: May 07, 2020
+"""
+from __future__ import print_function
+
+import torch
+import torch.nn as nn
+
+
+class SupervisedContrastiveLoss(nn.Module):
+    """Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf.
+    It also supports the unsupervised contrastive loss in SimCLR"""
+    def __init__(self, temperature=0.1, contrast_mode='all', base_temperature=0.1):
+        super(SupervisedContrastiveLoss, self).__init__()
+        self.temperature = temperature
+        self.contrast_mode = contrast_mode
+        self.base_temperature = base_temperature
+
+    def forward(self, features, labels=None, mask=None):
+        """Compute loss for model. If both `labels` and `mask` are None,
+        it degenerates to SimCLR unsupervised loss:
+        https://arxiv.org/pdf/2002.05709.pdf
+
+        Args:
+            features: hidden vector of shape [bsz, n_views, ...].
+            labels: ground truth of shape [bsz].
+            mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j
+                has the same class as sample i. Can be asymmetric.
+        Returns:
+            A loss scalar.
+        """
+        device = (torch.device('cuda')
+                  if features.is_cuda
+                  else torch.device('cpu'))
+
+        if len(features.shape) < 3:
+            raise ValueError('`features` needs to be [bsz, n_views, ...],'
+                             'at least 3 dimensions are required')
+        if len(features.shape) > 3:
+            features = features.view(features.shape[0], features.shape[1], -1)
+
+        batch_size = features.shape[0]
+        if labels is not None and mask is not None:
+            raise ValueError('Cannot define both `labels` and `mask`')
+        elif labels is None and mask is None:
+            mask = torch.eye(batch_size, dtype=torch.float32).to(device)
+        elif labels is not None:
+            labels = labels.contiguous().view(-1, 1)
+            if labels.shape[0] != batch_size:
+                raise ValueError('Num of labels does not match num of features')
+            mask = torch.eq(labels, labels.T).float().to(device)
+        else:
+            mask = mask.float().to(device)
+
+        contrast_count = features.shape[1]
+        contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0)
+        if self.contrast_mode == 'one':
+            anchor_feature = features[:, 0]
+            anchor_count = 1
+        elif self.contrast_mode == 'all':
+            anchor_feature = contrast_feature
+            anchor_count = contrast_count
+        else:
+            raise ValueError('Unknown mode: {}'.format(self.contrast_mode))
+
+        # compute logits
+        anchor_dot_contrast = torch.div(
+            torch.matmul(anchor_feature, contrast_feature.T),
+            self.temperature)
+        # for numerical stability
+        logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)
+        logits = anchor_dot_contrast - logits_max.detach()
+
+        # tile mask
+        mask = mask.repeat(anchor_count, contrast_count)
+        # mask-out self-contrast cases
+        logits_mask = torch.scatter(
+            torch.ones_like(mask),
+            1,
+            torch.arange(batch_size * anchor_count).view(-1, 1).to(device),
+            0
+        )
+        mask = mask * logits_mask
+
+        # compute log_prob
+        exp_logits = torch.exp(logits) * logits_mask
+        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))
+
+        # compute mean of log-likelihood over positive
+        # modified to handle edge cases when there is no positive pair
+        # for an anchor point. 
+        # Edge case e.g.:- 
+        # features of shape: [4,1,...]
+        # labels:            [0,1,1,2]
+        # loss before mean:  [nan, ..., ..., nan] 
+        mask_pos_pairs = mask.sum(1)
+        mask_pos_pairs = torch.where(mask_pos_pairs < 1e-6, 1, mask_pos_pairs)
+        mean_log_prob_pos = (mask * log_prob).sum(1) / mask_pos_pairs
+
+        # loss
+        loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos
+        loss = loss.view(anchor_count, batch_size).mean()
+
+        return loss

models/projector.py

@@ -0,0 +1,58 @@
+import torch
+
+class SiameseProjector(torch.nn.Module):
+    def __init__(self, inner_features = None, act_layer = torch.nn.GELU):
+        super().__init__()
+        
+        self.inner_features = inner_features
+        self.act_fcn = act_layer()
+
+        # Localisation branch.
+        self.input = torch.nn.Linear(in_features=inner_features, out_features=inner_features)
+        self.projection = torch.nn.Linear(in_features=inner_features, out_features=inner_features)
+        self.output = torch.nn.Linear(in_features=inner_features, out_features=inner_features)
+
+        # Contrastive branch.
+        self.input_con = torch.nn.Linear(in_features=inner_features, out_features=inner_features)
+        self.projection_con = torch.nn.Linear(in_features=inner_features, out_features=inner_features)
+        self.output_con = torch.nn.Linear(in_features=inner_features, out_features=inner_features)
+
+        # Localisation head.
+        self.probe = torch.nn.Conv2d(in_channels=inner_features, out_channels=1, kernel_size=3)
+
+    def forward(self, x):
+
+        # Localisation branch.
+        x = self.input(x)
+        x = self.act_fcn(x)
+        x = self.projection(x)
+        x = self.act_fcn(x)
+        x = self.output(x)
+
+        # Localisation head.
+        seg = self.probe(x.permute(0,2,1).reshape(x.shape[0], self.inner_features, int(x.shape[1]**0.5), int(x.shape[1]**0.5)))
+
+        # Contrastive branch.
+        y = self.input_con(x)
+        y = self.act_fcn(y)
+        y = self.projection_con(y)
+        y = self.act_fcn(y)
+
+        # Contrastive head.
+        feat = self.output_con(y)
+        
+        return feat, seg
+
+    def forward_segmentation(self, x):
+
+        # Localisation branch.
+        x = self.input(x)
+        x = self.act_fcn(x)
+        x = self.projection(x)
+        x = self.act_fcn(x)
+        x = self.output(x)
+        
+        # Localisation head.
+        seg = self.probe(x.permute(0,2,1).reshape(x.shape[0], self.inner_features, int(x.shape[1]**0.5), int(x.shape[1]**0.5)))
+
+        return seg

requirements.txt

@@ -0,0 +1,12 @@
+gradio
+datasets==4.0.0
+matplotlib==3.10.7
+numpy==2.3.3
+Pillow==11.3.0
+scikit_learn==1.7.2
+segmentation_models_pytorch==0.5.0
+torch==2.8.0
+torchvision==0.23.0
+tqdm==4.67.1
+transformers==4.57.1
+wandb==0.22.2