Add files using upload-large-folder tool

Pixal3D/torch_hub/hub/valeoai_NAF_main/config/loss/base.yaml

@@ -0,0 +1,3 @@
+mse:
+  _target_: src.loss.Loss
+  loss_type: mse

Pixal3D/torch_hub/hub/valeoai_NAF_main/config/model/anyup.yaml

@@ -0,0 +1,2 @@
+_target_: src.model.AnyUpsampler
+name: anyup

Pixal3D/torch_hub/hub/valeoai_NAF_main/config/model/bilinear.yaml

@@ -0,0 +1,2 @@
+_target_: src.model.Bilinear
+name: bilinear

Pixal3D/torch_hub/hub/valeoai_NAF_main/config/model/featup.yaml

@@ -0,0 +1,5 @@
+_target_: src.model.FeatUp
+feature_dim: ${get_feature:${backbone.name}}
+backbone_name: ${backbone.name}
+name: featup
+ratio: 16

Pixal3D/torch_hub/hub/valeoai_NAF_main/config/model/ircnn.yaml

@@ -0,0 +1,5 @@
+_target_: src.model.IRCNN
+in_nc: 3
+out_nc: 3
+nc: 64
+name: ircnn

Pixal3D/torch_hub/hub/valeoai_NAF_main/config/model/jafar.yaml

@@ -0,0 +1,10 @@
+defaults:
+  - _self_
+
+_target_: src.model.JAFAR
+dim: 128
+v_dim: ${model.feature_dim}
+feature_dim: ${get_feature:${backbone.name}}
+kernel_size: 1
+num_heads: 4
+name: jafar

Pixal3D/torch_hub/hub/valeoai_NAF_main/config/model/jbf.yaml

@@ -0,0 +1,5 @@
+_target_: src.model.JBF
+kernel_size: 5
+sigma_color: 0.1
+sigma_spatial: 1.5
+name: jbf

Pixal3D/torch_hub/hub/valeoai_NAF_main/config/model/jbu.yaml

@@ -0,0 +1,5 @@
+_target_: src.model.JBU
+dim: 256
+radius: 5
+combine: False
+name: jbu

Pixal3D/torch_hub/hub/valeoai_NAF_main/config/model/naf.yaml

@@ -0,0 +1,16 @@
+defaults:
+  - _self_
+
+_target_: src.model.naf.NAF
+dim: 256
+heads_attn: 4
+heads_rope: 4
+kernel_size: 9
+
+# ImageEncoder options
+use_semencoder: true
+use_encoder: true
+img_layers: 2
+rope_rescale: 2
+
+name: naf

Pixal3D/torch_hub/hub/valeoai_NAF_main/config/model/nearest.yaml

@@ -0,0 +1,3 @@
+_target_: src.model.Nearest
+feature_dim: ${get_feature:${backbone.name}}
+name: nearest

Pixal3D/torch_hub/hub/valeoai_NAF_main/config/model/rednet.yaml

@@ -0,0 +1,5 @@
+_target_: src.model.REDNet
+num_layers: 15
+num_features: 64
+input_dim: 3
+name: dncnn

Pixal3D/torch_hub/hub/valeoai_NAF_main/config/model/restormer.yaml

@@ -0,0 +1,4 @@
+_target_: src.model.Restormer
+inp_channels: 3
+out_channels: 3
+name: restormer

Pixal3D/torch_hub/hub/valeoai_NAF_main/config/optimizer/adamw.yaml

@@ -0,0 +1,4 @@
+_target_: torch.optim.AdamW
+lr: 2e-4
+betas: [0.9, 0.999]
+weight_decay: 1e-5

Pixal3D/torch_hub/hub/valeoai_NAF_main/docs/DATASETS.md

@@ -0,0 +1,25 @@
+# Datasets
+
+Create a data folder corresponding to the root data directory represented in the configuration file by `data_root: data/`.
+
+We train NAF using the ImageNet dataset, but it can be trained on any dataset of natural images. Then we evaluate NAF on several downstream tasks using the following datasets:
+- Semantic Segmentation: Cityscapes, PASCAL VOC 2012, ADE20K, KITTI-360
+- Depth Estimation: NYU Depth V2
+- Video Object Segmentation: DAVIS 2017
+
+To download the datasets used for training and evaluation, please refer to the [docs/data.sh](data.sh) file or see the official dataset documentation.
+
+The data folder should have the following structure:
+
+```
+data/
+(training)
+└── ImageNet/ 
+(evaluation)
+├── cityscapes/ 
+├── VOCdevkit/
+├── ADEChallengeData2016/ 
+├── KITTI-360/ 
+├── DAVIS/
+└── ...
+```

Pixal3D/torch_hub/hub/valeoai_NAF_main/docs/EVALUATIONS.md

@@ -0,0 +1,8 @@
+# 🔄 Evaluation 
+
+You should put the loaded weight file `naf_release.pth` in the `output/` folder.
+
+To evaluate the upsampler on semantic segmentation, simply run:
+``` python
+python evaluation/eval_seg_probing.py dataset=voc eval.model_ckpt=output/naf_release.pth
+```

Pixal3D/torch_hub/hub/valeoai_NAF_main/docs/INSTALL.md

@@ -0,0 +1,39 @@
+# Installation Instructions
+## Install Natten
+NAF is based on natten. Natten is build for specific torch versions. You simply need to install the compatible natten version for your torch version. For instance:
+
+```bash
+# Version used for developping.
+pip install natten==0.17.4+torch240cu118 -f https://shi-labs.com/natten/wheels/
+```
+For newer versions of torch or natten check this [url](https://github.com/SHI-Labs/NATTEN/blob/main/docs/install.md).
+For older versions of torch or natten check this [url](https://shi-labs.com/natten/wheels/).
+
+We have verified the following compatibilities:
+```base
+# For torch 2.7.0 + cu118
+pip install torch==2.7.0 torchvision==0.22.0 --index-url https://download.pytorch.org/whl/cu128
+pip install natten==0.20.1+torch270cu128 -f https://whl.natten.org
+
+# For torch 2.4.0 + cu118
+pip install torch==2.0.0 torchvision==0.15.1 --index-url https://download.pytorch.org/whl/cu118
+pip install natten==0.17.3+torch200cu117 -f https://shi-labs.com/natten/wheels/
+```
+
+Do not hesitate to open an issue if you face any problem installing natten.
+
+## Complete environment
+We used the following environment for developping NAF. We recommend using [micromamba](https://mamba.readthedocs.io/en/latest/installation/micromamba-installation.html) to create the environment and using [uv](https://github.com/astral-sh/uv) to speed up pip installs.
+
+``` bash
+micromamba create -n naf python==3.10.14  -y -c defaults
+micromamba activate naf
+pip install uv
+
+# Install torch and natten
+uv pip install torch==2.4.0 torchvision==0.19.0 --index-url https://download.pytorch.org/whl/cu118 
+pip install natten==0.17.4+torch240cu118 -f https://shi-labs.com/natten/wheels/
+
+# Install other dependencies
+uv pip install einops==0.8.0 numpy==1.24.4 timm==1.0.22 plotly==6.0.0 tensorboard==2.20.0 hydra-core==1.3.2 plotly==6.0.0 matplotlib==3.7.0 rich==14.2.0 torchmetrics==1.6.2 scipy==1.15.2 kornia==0.8.2 ipykernel ipympl pytest
+```

Pixal3D/torch_hub/hub/valeoai_NAF_main/docs/README.md

@@ -0,0 +1,18 @@
+# 📁 Repository Structure
+
+The repository contains the important folders:
+```
+|-- configs/                # Configuration files
+|-- docs/                   # Documentation
+|-- evaluation/             # Scripts to reproduce results and datasets initialization
+|-- notebooks/              # Jupyter notebooks for inference (of any VFM at any scale) and attention map visualizations
+|-- src/                    # Source code of the project
+|-- test/                   # Unit tests to compare model efficiency
+```
+
+
+# 🔍 Tests
+
+We provide unit tests to evaluate the efficiency of different model configurations, including forward/backward runtime, GPU memory usage, GFLOPs, and number of parameters.  
+All tests are available in the `test/` directory and we provide the test_results.json file with pre-computed results for reference computed on a 1 A100 40GB GPU.
+

Pixal3D/torch_hub/hub/valeoai_NAF_main/docs/TRAINING.md

@@ -0,0 +1,23 @@
+# Reproducibility 
+
+⚠️ Before launching commands, please make sure to set the `dataroot` variable in [config/base.yaml](config/base.yaml) to the path where the datasets are stored.
+
+General instructions:
+- Change the VFM: add `model.backbone=<backbone_name>` to the command line (e.g. `model.backbone=franca_vitb16`).
+- Change the dataset: add `dataset=<dataset_name>` to the command line (e.g. `dataset=cityscapes`).
+
+For debugging, add `HYDRA_FULL_ERROR=1` before the command to get full stack traces.
+
+# 🔄 Training
+Training NAF take less than 2 hours consuming less than 8GB of GPU memory on a single NVIDIA A100.
+
+To train the VFM-agnostic upsampler using NAF, simply run:
+``` python
+python train.py hydra.run.dir=output/naf
+```
+
+To train the denoising model using NAF, simply run:
+``` python
+python denoising.py denoising.noise_type=gaussian
+```
+You can select salt-and-pepper noise by setting `denoising.noise_type=salt_pepper`.

Pixal3D/torch_hub/hub/valeoai_NAF_main/docs/data.sh

@@ -0,0 +1,51 @@
+# ---------- COCOStuff ----------
+# Follow
+# https://github.com/nightrome/cocostuff
+# Download everything
+mkdir COCOStuff
+cd COCOStuff
+wget --directory-prefix=downloads http://images.cocodataset.org/zips/train2017.zip
+wget --directory-prefix=downloads http://images.cocodataset.org/zips/val2017.zip
+wget --directory-prefix=downloads http://calvin.inf.ed.ac.uk/wp-content/uploads/data/cocostuffdataset/stuffthingmaps_trainval2017.zip
+
+# Unpack everything
+mkdir -p dataset/images
+mkdir -p dataset/annotations
+unzip downloads/train2017.zip -d dataset/images/
+unzip downloads/val2017.zip -d dataset/images/
+unzip downloads/stuffthingmaps_trainval2017.zip -d dataset/annotations/
+
+# https://github.com/xu-ji/IIC/blob/master/datasets/README.txt
+wget https://www.robots.ox.ac.uk/~xuji/datasets/COCOStuff164kCurated.tar.gz
+tar -xzf COCOStuff164kCurated.tar.gz
+mv COCO/COCOStuff164k ./currated
+rmdir COCO
+rm COCOStuff164kCurated.tar.gz
+
+# ---------- Cityscapes ----------
+# https://github.com/cemsaz/city-scapes-script
+wget --keep-session-cookies --save-cookies=cookies.txt --post-data 'username=XXXX&password=XXX&submit=Login' https://www.cityscapes-dataset.com/login/
+wget --load-cookies cookies.txt --content-disposition https://www.cityscapes-dataset.com/file-handling/?packageID=1
+wget --load-cookies cookies.txt --content-disposition https://www.cityscapes-dataset.com/file-handling/?packageID=3
+unzip gtFine_trainvaltest.zip 
+unzip leftImg8bit_trainvaltest.zip
+
+
+# ---------- VOC ----------
+wget http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar
+# generate folder VOCdevkit/VOC2012
+tar -xvf VOCtrainval_11-May-2012.tar
+
+# ---------- ADE20K ----------
+wget http://data.csail.mit.edu/places/ADEchallenge/ADEChallengeData2016.zip
+unzip ADEChallengeData2016.zip
+
+
+# ---------- DAVIS ----------
+mkdir DAVIS
+cd DAVIS
+wget https://data.vision.ee.ethz.ch/csergi/share/davis/DAVIS-2017-trainval-480p.zip
+unzip DAVIS-2017-trainval-480p.zip
+rm DAVIS-2017-trainval-480p.zip
+mv DAVIS/* ./
+rmdir DAVIS

Pixal3D/torch_hub/hub/valeoai_NAF_main/evaluation/dataset/ade20k.py

@@ -0,0 +1,231 @@
+import os
+from typing import Optional
+
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+
+
+class ADE20KDataset(Dataset):
+    split_to_dir = {"train": "training", "val": "validation"}
+
+    def __init__(
+        self,
+        root,
+        transform,
+        target_transform,
+        split="train",
+        skip_other_class=False,
+        file_set=None,
+        num_classes=None,
+        tag=None,
+    ):
+        super().__init__()
+        self.transform = transform
+        self.target_transform = target_transform
+        self.split = split
+        self.root = root
+        self.skip_other_class = skip_other_class
+        self.file_set = file_set
+
+        # Collect the data
+        self.data = self.collect_data()
+
+        if split == "train":
+            assert self.__len__() == 20210
+        elif split == "val":
+            assert self.__len__() == 2000
+
+    def get_class_names(self):
+        return [
+            "background",
+            "wall",
+            "building",
+            "sky",
+            "floor",
+            "tree",
+            "ceiling",
+            "road",
+            "bed ",
+            "windowpane",
+            "grass",
+            "cabinet",
+            "sidewalk",
+            "person",
+            "earth",
+            "door",
+            "table",
+            "mountain",
+            "plant",
+            "curtain",
+            "chair",
+            "car",
+            "water",
+            "painting",
+            "sofa",
+            "shelf",
+            "house",
+            "sea",
+            "mirror",
+            "rug",
+            "field",
+            "armchair",
+            "seat",
+            "fence",
+            "desk",
+            "rock",
+            "wardrobe",
+            "lamp",
+            "bathtub",
+            "railing",
+            "cushion",
+            "base",
+            "box",
+            "column",
+            "signboard",
+            "chest of drawers",
+            "counter",
+            "sand",
+            "sink",
+            "skyscraper",
+            "fireplace",
+            "refrigerator",
+            "grandstand",
+            "path",
+            "stairs",
+            "runway",
+            "case",
+            "pool table",
+            "pillow",
+            "screen door",
+            "stairway",
+            "river",
+            "bridge",
+            "bookcase",
+            "blind",
+            "coffee table",
+            "toilet",
+            "flower",
+            "book",
+            "hill",
+            "bench",
+            "countertop",
+            "stove",
+            "palm",
+            "kitchen island",
+            "computer",
+            "swivel chair",
+            "boat",
+            "bar",
+            "arcade machine",
+            "hovel",
+            "bus",
+            "towel",
+            "light",
+            "truck",
+            "tower",
+            "chandelier",
+            "awning",
+            "streetlight",
+            "booth",
+            "television receiver",
+            "airplane",
+            "dirt track",
+            "apparel",
+            "pole",
+            "land",
+            "bannister",
+            "escalator",
+            "ottoman",
+            "bottle",
+            "buffet",
+            "poster",
+            "stage",
+            "van",
+            "ship",
+            "fountain",
+            "conveyer belt",
+            "canopy",
+            "washer",
+            "plaything",
+            "swimming pool",
+            "stool",
+            "barrel",
+            "basket",
+            "waterfall",
+            "tent",
+            "bag",
+            "minibike",
+            "cradle",
+            "oven",
+            "ball",
+            "food",
+            "step",
+            "tank",
+            "trade name",
+            "microwave",
+            "pot",
+            "animal",
+            "bicycle",
+            "lake",
+            "dishwasher",
+            "screen",
+            "blanket",
+            "sculpture",
+            "hood",
+            "sconce",
+            "vase",
+            "traffic light",
+            "tray",
+            "ashcan",
+            "fan",
+            "pier",
+            "crt screen",
+            "plate",
+            "monitor",
+            "bulletin board",
+            "shower",
+            "radiator",
+            "glass",
+            "clock",
+            "flag",
+        ]
+
+    def collect_data(self):
+        # Get the image and annotation dirs
+        image_dir = os.path.join(self.root, f"images/{self.split_to_dir[self.split]}")
+        annotation_dir = os.path.join(self.root, f"annotations/{self.split_to_dir[self.split]}")
+
+        # Collect the filepaths
+        if self.file_set is None:
+            image_paths = [os.path.join(image_dir, f) for f in sorted(os.listdir(image_dir))]
+            annotation_paths = [os.path.join(annotation_dir, f) for f in sorted(os.listdir(annotation_dir))]
+        else:
+            image_paths = [os.path.join(image_dir, f"{f}.jpg") for f in sorted(self.file_set)]
+            annotation_paths = [os.path.join(annotation_dir, f"{f}.png") for f in sorted(self.file_set)]
+
+        data = list(zip(image_paths, annotation_paths))
+        return data
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        # Get the  paths
+        image_path, annotation_path = self.data[index]
+
+        # Load
+        image = Image.open(image_path).convert("RGB")
+        target = Image.open(annotation_path)
+
+        # Augment
+        image = self.transform(image).squeeze(0)
+        target = self.target_transform(target)
+        if self.skip_other_class == True:
+            target = target * 255.0
+            target[target.type(torch.int64) == 0] = 255.0
+            target /= 255.0
+        target = target.squeeze(0)
+
+        batch = {"image": image, "label": target}
+        return batch

Pixal3D/torch_hub/hub/valeoai_NAF_main/evaluation/dataset/cityscapes.py

@@ -0,0 +1,115 @@
+import os
+import random
+
+import numpy as np
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+from torchvision.datasets import Cityscapes
+
+
+class CityscapesDataset(Dataset):
+    BASE_DIR = "cityscapes"
+    NUM_CLASS = 19
+
+    def __init__(
+        self,
+        root,
+        split="train",
+        transform=None,
+        target_transform=None,
+        include_labels=True,
+        num_classes=19,
+        tag=None,
+        **kwargs
+    ):
+        super().__init__()
+        self.root = root
+        self.split = split
+        self.transform = transform
+        self.target_transform = target_transform
+        self.include_labels = include_labels
+        self.num_classes = num_classes
+        self.tag = tag
+
+        # fmt: off
+        self._key = np.array([-1, -1, -1, -1, -1, -1,
+                              -1, -1, 0, 1, -1, -1,
+                              2, 3, 4, -1, -1, -1,
+                              5, -1, 6, 7, 8, 9,
+                              10, 11, 12, 13, 14, 15,
+                              -1, -1, 16, 17, 18])
+        self._mapping = np.array(range(-1, len(self._key) - 1)).astype("int32")
+        # fmt: on
+
+        # Get image/mask path pairs
+        self.cityscapes_dataset = Cityscapes(
+            root=root,
+            split=split,
+            mode="fine",
+            target_type=["instance", "semantic"],
+            transform=None,
+            target_transform=None,
+        )
+        # fmt: on
+
+        if split == "train":
+            assert self.__len__() == 2975
+        elif split == "val":
+            assert self.__len__() == 500
+
+    def get_class_names(
+        self,
+    ):
+        return [
+            "road",
+            "sidewalk",
+            "building",
+            "wall",
+            "fence",
+            "pole",
+            "traffic_light",
+            "traffic_sign",
+            "vegetation",
+            "terrain",
+            "sky",
+            "person",
+            "rider",
+            "car",
+            "truck",
+            "bus",
+            "train",
+            "motorcycle",
+            "bicycle",
+        ]
+
+    def __len__(self):
+        return len(self.cityscapes_dataset)
+
+    def __getitem__(self, index):
+        batch = {}
+
+        # Load image
+        data = self.cityscapes_dataset[index]
+        img = data[0]
+        mask = data[1][1]
+
+        if self.transform:
+            img = self.transform(img)
+        batch["image"] = img
+
+        if self.include_labels:
+            if self.target_transform:
+                mask = self.target_transform(mask)
+            # Convert to tensor and match COCO/ADE20K format
+            mask = torch.from_numpy(self._class_to_index(np.array(mask))).to(torch.uint8)
+            batch["label"] = mask.squeeze()
+
+        return batch
+
+    def _class_to_index(self, mask):
+        values = np.unique(mask)
+        for value in values:
+            assert value in self._mapping
+        index = np.digitize(mask.ravel(), self._mapping, right=True)
+        return self._key[index].reshape(mask.shape).astype(np.int64)

Pixal3D/torch_hub/hub/valeoai_NAF_main/evaluation/dataset/coco.py

@@ -0,0 +1,317 @@
+import random
+from os.path import join
+
+import numpy as np
+import torch
+import torch.multiprocessing
+from PIL import Image
+from torch.utils.data import Dataset
+
+
+def bit_get(val, idx):
+    """Gets the bit value.
+    Args:
+      val: Input value, int or numpy int array.
+      idx: Which bit of the input val.
+    Returns:
+      The "idx"-th bit of input val.
+    """
+    return (val >> idx) & 1
+
+
+def create_pascal_label_colormap():
+    """Creates a label colormap used in PASCAL VOC segmentation benchmark.
+    Returns:
+      A colormap for visualizing segmentation results.
+    """
+    colormap = np.zeros((512, 3), dtype=int)
+    ind = np.arange(512, dtype=int)
+
+    for shift in reversed(list(range(8))):
+        for channel in range(3):
+            colormap[:, channel] |= bit_get(ind, channel) << shift
+        ind >>= 3
+
+    return colormap
+
+
+class COCODataset(Dataset):
+
+    def __init__(
+        self,
+        root,
+        split,
+        transform,
+        target_transform,
+        include_labels=True,
+        coarse_labels=False,
+        exclude_things=False,
+        subset=None,
+        num_classes=None,
+        tag=None,
+    ):
+        super().__init__()
+        self.split = split
+        self.root = root
+        self.coarse_labels = coarse_labels
+        self.transform = transform
+        self.label_transform = target_transform
+        self.subset = subset
+        self.exclude_things = exclude_things
+        self.include_labels = include_labels
+
+        if self.subset is None:
+            self.image_list = "Coco164kFull_Stuff_Coarse.txt"
+        elif self.subset == 6:  # IIC Coarse
+            self.image_list = "Coco164kFew_Stuff_6.txt"
+        elif self.subset == 7:  # IIC Fine
+            self.image_list = "Coco164kFull_Stuff_Coarse_7.txt"
+
+        assert self.split in ["train", "val", "train+val"]
+        split_dirs = {
+            "train": ["train2017"],
+            "val": ["val2017"],
+            "train+val": ["train2017", "val2017"],
+        }
+
+        self.image_files = []
+        self.label_files = []
+        for split_dir in split_dirs[self.split]:
+            with open(join(self.root, "curated", split_dir, self.image_list), "r") as f:
+                img_ids = [fn.rstrip() for fn in f.readlines()]
+                for img_id in img_ids:
+                    self.image_files.append(join(self.root, "images", split_dir, img_id + ".jpg"))
+                    self.label_files.append(join(self.root, "annotations", split_dir, img_id + ".png"))
+
+        self.fine_to_coarse = {
+            0: 9,
+            1: 11,
+            2: 11,
+            3: 11,
+            4: 11,
+            5: 11,
+            6: 11,
+            7: 11,
+            8: 11,
+            9: 8,
+            10: 8,
+            11: 8,
+            12: 8,
+            13: 8,
+            14: 8,
+            15: 7,
+            16: 7,
+            17: 7,
+            18: 7,
+            19: 7,
+            20: 7,
+            21: 7,
+            22: 7,
+            23: 7,
+            24: 7,
+            25: 6,
+            26: 6,
+            27: 6,
+            28: 6,
+            29: 6,
+            30: 6,
+            31: 6,
+            32: 6,
+            33: 10,
+            34: 10,
+            35: 10,
+            36: 10,
+            37: 10,
+            38: 10,
+            39: 10,
+            40: 10,
+            41: 10,
+            42: 10,
+            43: 5,
+            44: 5,
+            45: 5,
+            46: 5,
+            47: 5,
+            48: 5,
+            49: 5,
+            50: 5,
+            51: 2,
+            52: 2,
+            53: 2,
+            54: 2,
+            55: 2,
+            56: 2,
+            57: 2,
+            58: 2,
+            59: 2,
+            60: 2,
+            61: 3,
+            62: 3,
+            63: 3,
+            64: 3,
+            65: 3,
+            66: 3,
+            67: 3,
+            68: 3,
+            69: 3,
+            70: 3,
+            71: 0,
+            72: 0,
+            73: 0,
+            74: 0,
+            75: 0,
+            76: 0,
+            77: 1,
+            78: 1,
+            79: 1,
+            80: 1,
+            81: 1,
+            82: 1,
+            83: 4,
+            84: 4,
+            85: 4,
+            86: 4,
+            87: 4,
+            88: 4,
+            89: 4,
+            90: 4,
+            91: 17,
+            92: 17,
+            93: 22,
+            94: 20,
+            95: 20,
+            96: 22,
+            97: 15,
+            98: 25,
+            99: 16,
+            100: 13,
+            101: 12,
+            102: 12,
+            103: 17,
+            104: 17,
+            105: 23,
+            106: 15,
+            107: 15,
+            108: 17,
+            109: 15,
+            110: 21,
+            111: 15,
+            112: 25,
+            113: 13,
+            114: 13,
+            115: 13,
+            116: 13,
+            117: 13,
+            118: 22,
+            119: 26,
+            120: 14,
+            121: 14,
+            122: 15,
+            123: 22,
+            124: 21,
+            125: 21,
+            126: 24,
+            127: 20,
+            128: 22,
+            129: 15,
+            130: 17,
+            131: 16,
+            132: 15,
+            133: 22,
+            134: 24,
+            135: 21,
+            136: 17,
+            137: 25,
+            138: 16,
+            139: 21,
+            140: 17,
+            141: 22,
+            142: 16,
+            143: 21,
+            144: 21,
+            145: 25,
+            146: 21,
+            147: 26,
+            148: 21,
+            149: 24,
+            150: 20,
+            151: 17,
+            152: 14,
+            153: 21,
+            154: 26,
+            155: 15,
+            156: 23,
+            157: 20,
+            158: 21,
+            159: 24,
+            160: 15,
+            161: 24,
+            162: 22,
+            163: 25,
+            164: 15,
+            165: 20,
+            166: 17,
+            167: 17,
+            168: 22,
+            169: 14,
+            170: 18,
+            171: 18,
+            172: 18,
+            173: 18,
+            174: 18,
+            175: 18,
+            176: 18,
+            177: 26,
+            178: 26,
+            179: 19,
+            180: 19,
+            181: 24,
+        }
+
+        self._label_names = [
+            "ground-stuff",
+            "plant-stuff",
+            "sky-stuff",
+        ]
+        self.cocostuff3_coarse_classes = [23, 22, 21]
+        self.first_stuff_index = 12
+        if split == "train":
+            assert self.__len__() == 97702
+        elif split == "val":
+            assert self.__len__() == 4172
+
+    def __len__(self):
+        return len(self.image_files)
+
+    def __getitem__(self, index):
+        image_path = self.image_files[index]
+        label_path = self.label_files[index]
+        batch = {}
+
+        img = self.transform(Image.open(image_path).convert("RGB"))
+        batch["image"] = img
+        batch["img_path"] = image_path
+
+        if self.include_labels:
+            label = self.label_transform(Image.open(label_path)).squeeze(0)
+            label[label == 255] = -1  # to be consistent with 10k
+            batch["label"] = label
+
+            coarse_label = torch.zeros_like(label)
+            for fine, coarse in self.fine_to_coarse.items():
+                coarse_label[label == fine] = coarse
+            coarse_label[label == -1] = -1
+
+            if self.coarse_labels:
+                coarser_labels = -torch.ones_like(label)
+                for i, c in enumerate(self.cocostuff3_coarse_classes):
+                    coarser_labels[coarse_label == c] = i
+                batch["label"] = coarser_labels
+            else:
+                if self.exclude_things:
+                    batch["label"] = coarse_label - self.first_stuff_index
+                else:
+                    batch["label"] = coarse_label
+
+        batch["label"] = batch["label"].to(torch.uint8)
+        return batch

Pixal3D/torch_hub/hub/valeoai_NAF_main/evaluation/dataset/davis.py

@@ -0,0 +1,61 @@
+import glob
+import os
+
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+from torchvision import transforms
+
+
+class DAVIS(Dataset):
+    def __init__(self, root, split="train", transform=None, target_transform=None, tag=None):
+        """
+        Args:
+            root (string): Directory with all the videos.
+            transform (callable, optional): Optional transform to be applied on a sample.
+        """
+        videos = self.load_split(root, split)
+
+        frames = []
+        for video in videos:
+            video_path = os.path.join(root, f"JPEGImages/480p/{video}/*/*.jpg")
+            video_frames = sorted(glob.glob(video_path))
+            frames += video_frames
+        self.frames = frames
+        self.transform = transform
+        self.target_transform = target_transform
+        self.tag = tag
+
+    def load_split(self, root, split):
+        """
+        Load the split files for DAVIS dataset.
+        Args:
+            root (string): Directory with all the videos.
+            split (string): Split name (train, val, test).
+        Returns:
+            list: List of video paths.
+        """
+        split_path = os.path.join(root, "ImageSets/2017/", f"{split}.txt")
+        with open(split_path, "r") as f:
+            videos = [line.strip() for line in f.readlines()]
+        return videos
+
+    def __len__(self):
+        return len(self.frames)
+
+    def __getitem__(self, index):
+        image_path = self.frames[index]
+        annotation_path = image_path.replace("JPEGImages", "Annotations")
+        batch = {}
+
+        # Load
+        image = Image.open(image_path).convert("RGB")
+        target = Image.open(annotation_path)
+
+        # Augment
+        image = self.transform(image).squeeze(0)
+        target = self.target_transform(target)
+        target = target.squeeze(0)
+
+        batch = {"image": image, "label": target}
+        return batch

Pixal3D/torch_hub/hub/valeoai_NAF_main/evaluation/dataset/image_dataset.py

@@ -0,0 +1,118 @@
+import os
+import sys
+
+import torch
+import torchvision.transforms as T
+from torchvision.datasets import folder
+from torchvision.datasets.vision import VisionDataset
+
+
+def remove_prefix(s, prefix):
+    if s.startswith(prefix):
+        s = s[len(prefix) :]
+    return s
+
+
+class ImageDataset(VisionDataset):
+    """
+    modified from: https://pytorch.org/docs/stable/_modules/torchvision/datasets/folder.html#ImageFolder
+    uses cached directory listing if available rather than walking directory
+     Attributes:
+        classes (list): List of the class names.
+        class_to_idx (dict): Dict with items (class_name, class_index).
+        samples (list): List of (sample path, class_index) tuples
+        targets (list): The class_index value for each image in the dataset
+    """
+
+    def __init__(
+        self,
+        root,
+        root_cache,
+        loader=folder.default_loader,
+        extensions=folder.IMG_EXTENSIONS,
+        transform=None,
+        is_valid_file=None,
+        include_labels=False,
+        **kwargs,
+    ):
+        super(ImageDataset, self).__init__(
+            root,
+            transform=transform,
+        )
+        self.root = root
+        classes, class_to_idx = self._find_classes(self.root)
+        if root_cache is None:
+            root_cache = root
+        cache = root_cache.rstrip("/") + ".txt"
+        print(cache)
+        # cache = './val_paths.txt'
+        if os.path.isfile(cache):
+            print("Using directory list at: %s" % cache)
+            with open(cache) as f:
+                samples = []
+                for line in f:
+                    (path, idx) = line.strip().split(";")
+                    samples.append((os.path.join(path), int(idx)))
+        else:
+            print("Walking directory: %s" % self.root)
+            samples = folder.make_dataset(self.root, class_to_idx, extensions, is_valid_file)
+            with open(cache, "w") as f:
+                for line in samples:
+                    path, label = line
+                    f.write("%s;%d\n" % (remove_prefix(path, self.root).lstrip("/"), label))
+
+        if len(samples) == 0:
+            raise (
+                RuntimeError(
+                    "Found 0 files in subfolders of: " + self.root + "\n" "Supported extensions are: " + ",".join(extensions)
+                )
+            )
+
+        self.loader = loader
+        self.extensions = extensions
+        self.classes = classes
+        self.class_to_idx = class_to_idx
+        self.samples = samples
+        self.targets = [s[1] for s in samples]
+        self.transform = transform
+        self.include_labels = include_labels
+
+    def _find_classes(self, dir):
+        """
+        Finds the class folders in a dataset.
+        Ensures:
+            No class is a subdirectory of another.
+        """
+        if sys.version_info >= (3, 5):
+            # Faster and available in Python 3.5 and above
+            classes = [d.name for d in os.scandir(dir) if d.is_dir()]
+        else:
+            classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))]
+        classes.sort()
+        class_to_idx = {classes[i]: i for i in range(len(classes))}
+        return classes, class_to_idx
+
+    def __getitem__(self, index):
+        path, target = self.samples[index]
+        path = os.path.join(self.root, path)
+        sample = self.loader(path)
+        if self.transform is not None:
+            sample = self.transform(sample)
+
+        batch = {
+            "index": index,
+            "image": sample,
+            "target": target,
+            "path": path,
+        }
+
+        if self.include_labels:
+            target = self.targets[index]
+            if self.target_transform is not None:
+                target = self.target_transform(target)
+            batch["label"] = target
+
+        return batch
+
+    def __len__(self):
+        return len(self.samples)

Pixal3D/torch_hub/hub/valeoai_NAF_main/evaluation/dataset/kitti360.py

@@ -0,0 +1,196 @@
+import json
+import os
+import random
+from collections import namedtuple
+from typing import Optional
+
+import numpy as np
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+
+# https://github.com/autonomousvision/kitti360Scripts/blob/32f6d64eef27c32b52c542b4e95be8af9c9c6444/kitti360scripts/helpers/labels.py
+# a label and all meta information
+Label = namedtuple(
+    "Label",
+    ["name", "id", "kittiId", "trainId", "category", "categoryId", "hasInstances", "ignoreInEval", "ignoreInInst", "color"],
+)
+
+
+labels = [
+    #       name                     id    kittiId,    trainId   category            catId     hasInstances   ignoreInEval   ignoreInInst   color
+    Label("unlabeled", 0, -1, 255, "void", 0, False, True, True, (0, 0, 0)),
+    Label("ego vehicle", 1, -1, 255, "void", 0, False, True, True, (0, 0, 0)),
+    Label("rectification border", 2, -1, 255, "void", 0, False, True, True, (0, 0, 0)),
+    Label("out of roi", 3, -1, 255, "void", 0, False, True, True, (0, 0, 0)),
+    Label("static", 4, -1, 255, "void", 0, False, True, True, (0, 0, 0)),
+    Label("dynamic", 5, -1, 255, "void", 0, False, True, True, (111, 74, 0)),
+    Label("ground", 6, -1, 255, "void", 0, False, True, True, (81, 0, 81)),
+    Label("road", 7, 1, 0, "flat", 1, False, False, False, (128, 64, 128)),
+    Label("sidewalk", 8, 3, 1, "flat", 1, False, False, False, (244, 35, 232)),
+    Label("parking", 9, 2, 255, "flat", 1, False, True, True, (250, 170, 160)),
+    Label("rail track", 10, 10, 255, "flat", 1, False, True, True, (230, 150, 140)),
+    Label("building", 11, 11, 2, "construction", 2, True, False, False, (70, 70, 70)),
+    Label("wall", 12, 7, 3, "construction", 2, False, False, False, (102, 102, 156)),
+    Label("fence", 13, 8, 4, "construction", 2, False, False, False, (190, 153, 153)),
+    Label("guard rail", 14, 30, 255, "construction", 2, False, True, True, (180, 165, 180)),
+    Label("bridge", 15, 31, 255, "construction", 2, False, True, True, (150, 100, 100)),
+    Label("tunnel", 16, 32, 255, "construction", 2, False, True, True, (150, 120, 90)),
+    Label("pole", 17, 21, 5, "object", 3, True, False, True, (153, 153, 153)),
+    Label("polegroup", 18, -1, 255, "object", 3, False, True, True, (153, 153, 153)),
+    Label("traffic light", 19, 23, 6, "object", 3, True, False, True, (250, 170, 30)),
+    Label("traffic sign", 20, 24, 7, "object", 3, True, False, True, (220, 220, 0)),
+    Label("vegetation", 21, 5, 8, "nature", 4, False, False, False, (107, 142, 35)),
+    Label("terrain", 22, 4, 9, "nature", 4, False, False, False, (152, 251, 152)),
+    Label("sky", 23, 9, 10, "sky", 5, False, False, False, (70, 130, 180)),
+    Label("person", 24, 19, 11, "human", 6, True, False, False, (220, 20, 60)),
+    Label("rider", 25, 20, 12, "human", 6, True, False, False, (255, 0, 0)),
+    Label("car", 26, 13, 13, "vehicle", 7, True, False, False, (0, 0, 142)),
+    Label("truck", 27, 14, 14, "vehicle", 7, True, False, False, (0, 0, 70)),
+    Label("bus", 28, 34, 15, "vehicle", 7, True, False, False, (0, 60, 100)),
+    Label("caravan", 29, 16, 255, "vehicle", 7, True, True, True, (0, 0, 90)),
+    Label("trailer", 30, 15, 255, "vehicle", 7, True, True, True, (0, 0, 110)),
+    Label("train", 31, 33, 16, "vehicle", 7, True, False, False, (0, 80, 100)),
+    Label("motorcycle", 32, 17, 17, "vehicle", 7, True, False, False, (0, 0, 230)),
+    Label("bicycle", 33, 18, 18, "vehicle", 7, True, False, False, (119, 11, 32)),
+    Label("garage", 34, 12, 2, "construction", 2, True, True, True, (64, 128, 128)),
+    Label("gate", 35, 6, 4, "construction", 2, False, True, True, (190, 153, 153)),
+    Label("stop", 36, 29, 255, "construction", 2, True, True, True, (150, 120, 90)),
+    Label("smallpole", 37, 22, 5, "object", 3, True, True, True, (153, 153, 153)),
+    Label("lamp", 38, 25, 255, "object", 3, True, True, True, (0, 64, 64)),
+    Label("trash bin", 39, 26, 255, "object", 3, True, True, True, (0, 128, 192)),
+    Label("vending machine", 40, 27, 255, "object", 3, True, True, True, (128, 64, 0)),
+    Label("box", 41, 28, 255, "object", 3, True, True, True, (64, 64, 128)),
+    Label("unknown construction", 42, 35, 255, "void", 0, False, True, True, (102, 0, 0)),
+    Label("unknown vehicle", 43, 36, 255, "void", 0, False, True, True, (51, 0, 51)),
+    Label("unknown object", 44, 37, 255, "void", 0, False, True, True, (32, 32, 32)),
+    Label("license plate", -1, -1, -1, "vehicle", 7, False, True, True, (0, 0, 142)),
+]
+
+
+class KITTIDataset(Dataset):
+    def __init__(
+        self,
+        root: str,
+        split: str = "train",
+        transform: Optional[callable] = None,
+        target_transform: Optional[callable] = None,
+        include_labels: bool = True,
+        num_classes: int = 19,
+        tag: Optional[str] = None,
+        **kwargs,
+    ):
+        super().__init__()
+        self.root = root
+        self.split = split
+        self.transform = transform
+        self.target_transform = target_transform
+        self.include_labels = include_labels
+        self.num_classes = num_classes
+        self.tag = tag
+
+        # Create mapping from id to trainId
+        self.id_to_trainid = {label.id: label.trainId for label in labels}
+
+        # Initialize paths
+        self.image_paths = []
+        self.label_paths = []
+
+        # Set up split management
+        self.split_file = os.path.join(os.path.dirname(__file__), "kitti360", f"{split}_split.json")
+
+        # Create split if needed
+        if not os.path.exists(self.split_file):
+            self._create_split()
+
+        # Load paths from split file
+        self._load_split()
+
+    def _create_split(self):
+        """Create and save train/val split if it doesn't exist"""
+        # Collect all valid paths first
+        all_image_paths = []
+        all_label_paths = []
+
+        raw_base = os.path.join(self.root, "data_2d_raw")
+        sem_base = os.path.join(self.root, "data_2d_semantics", "train")
+
+        for drive in os.listdir(raw_base):
+            drive_path = os.path.join(raw_base, drive)
+
+            for camera in ["image_00"]:
+                img_dir = os.path.join(drive_path, camera, "data_rect")
+                lbl_dir = os.path.join(sem_base, drive, camera, "semantic")
+
+                if os.path.exists(img_dir) and os.path.exists(lbl_dir):
+                    for fname in os.listdir(img_dir):
+                        if fname.endswith(".png"):
+                            base_name = os.path.splitext(fname)[0]
+                            img_path = os.path.join(img_dir, fname)
+                            lbl_path = os.path.join(lbl_dir, f"{base_name}.png")
+
+                            if os.path.exists(lbl_path):
+                                all_image_paths.append(img_path)
+                                all_label_paths.append(lbl_path)
+
+        # Create 80-20 split
+        random.seed(42)  # For reproducibility
+        combined = list(zip(all_image_paths, all_label_paths))
+        random.shuffle(combined)
+        split_idx = int(0.8 * len(combined))
+
+        train_split = combined[:split_idx]
+        val_split = combined[split_idx:]
+
+        # Create directory if needed
+        split_dir = os.path.dirname(self.split_file)
+        os.makedirs(split_dir, exist_ok=True)
+
+        # Save splits
+        with open(os.path.join(split_dir, "train_split.json"), "w") as f:
+            json.dump(train_split, f)
+
+        with open(os.path.join(split_dir, "val_split.json"), "w") as f:
+            json.dump(val_split, f)
+
+    def _load_split(self):
+        """Load paths from existing split file"""
+        with open(self.split_file, "r") as f:
+            split_data = json.load(f)
+
+        self.image_paths, self.label_paths = zip(*split_data)
+
+    def __len__(self):
+        return len(self.image_paths)
+
+    def __getitem__(self, index):
+        batch = {}
+        seed = random.randint(0, 2**32 - 1)
+
+        # Load image
+        random.seed(seed)
+        torch.manual_seed(seed)
+        img = Image.open(self.image_paths[index]).convert("RGB")
+        if self.transform:
+            img = self.transform(img)
+        batch["image"] = img
+        batch["img_path"] = self.image_paths[index]
+
+        if self.include_labels:
+            # Load and process label
+            random.seed(seed)
+            torch.manual_seed(seed)
+            label = Image.open(self.label_paths[index])
+
+            if self.target_transform:
+                label = self.target_transform(label)
+
+            # Convert from id to trainId
+            label_array = np.array(label)
+            for id_val, trainid_val in self.id_to_trainid.items():
+                label_array[label_array == id_val] = trainid_val
+
+            label = torch.from_numpy(label_array).to(torch.uint8)
+            batch["label"] = label.squeeze(0)
+
+        return batch

Pixal3D/torch_hub/hub/valeoai_NAF_main/evaluation/dataset/voc.py

@@ -0,0 +1,113 @@
+import random
+from typing import Optional
+
+import numpy as np
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+from torchvision.datasets import VOCSegmentation
+
+
+class VOCDataset(Dataset):
+    def __init__(
+        self,
+        root: str,
+        split: str = "train",
+        transform: Optional[callable] = None,
+        target_transform: Optional[callable] = None,
+        include_labels: bool = True,
+        num_classes: int = 21,
+        tag: Optional[str] = None,
+        year: str = "2012",
+        download: bool = False,
+        **kwargs
+    ):
+        super().__init__()
+        self.root = root
+        self.split = split
+        self.transform = transform
+        self.target_transform = target_transform
+        self.include_labels = include_labels
+        self.num_classes = num_classes
+        self.tag = tag
+
+        # Initialize torchvision's VOC dataset
+        self.voc_dataset = VOCSegmentation(
+            root=root,
+            year=year,
+            image_set="train" if split == "train" else "val",
+            download=download,
+            transform=None,
+            target_transform=None,
+        )
+
+        if split == "train":
+            assert self.__len__() == 1464
+        elif split == "val":
+            assert self.__len__() == 1449
+
+    def __len__(self):
+        return len(self.voc_dataset)
+
+    def __getitem__(self, index):
+        batch = {}
+        seed = random.randint(0, 2**32 - 1)
+
+        # Load image
+        img_path, mask_path = (
+            self.voc_dataset.images[index],
+            self.voc_dataset.masks[index],
+        )
+
+        random.seed(seed)
+        torch.manual_seed(seed)
+        img = Image.open(img_path).convert("RGB")
+        if self.transform:
+            img = self.transform(img)
+        batch["image"] = img
+        batch["img_path"] = img_path
+
+        if self.include_labels:
+            # Load and process mask
+            random.seed(seed)
+            torch.manual_seed(seed)
+            mask = Image.open(mask_path)
+
+            if self.target_transform:
+                mask = self.target_transform(mask)
+
+            # Convert to tensor and match COCO/ADE20K format
+            mask = torch.from_numpy(np.array(mask)).to(torch.uint8)
+            batch["label"] = mask.squeeze()
+
+        return batch
+
+    @property
+    def pred_offset(self):
+        return 0
+
+    def get_class_names(self):
+        """Return the class names for the VOC dataset."""
+        return [
+            "background",
+            "aeroplane",
+            "bicycle",
+            "bird",
+            "boat",
+            "bottle",
+            "bus",
+            "car",
+            "cat",
+            "chair",
+            "cow",
+            "diningtable",
+            "dog",
+            "horse",
+            "motorbike",
+            "person",
+            "pottedplant",
+            "sheep",
+            "sofa",
+            "train",
+            "tvmonitor",
+        ]

Pixal3D/torch_hub/hub/valeoai_NAF_main/evaluation/eval_depth_probing.py

@@ -0,0 +1,448 @@
+import datetime
+import os
+import sys
+from pathlib import Path
+
+import hydra
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as T
+from einops import rearrange
+from hydra.core.global_hydra import GlobalHydra
+from hydra.core.hydra_config import HydraConfig
+from hydra.utils import instantiate
+from omegaconf import OmegaConf
+from rich.console import Console
+from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn
+from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import random
+
+from src.loss import DepthLoss
+from utils.training import get_batch, get_dataloaders, load_multiple_backbones
+
+LOG_INTERVAL = 100
+
+
+def eval_metrics(gt, pred, min_depth=1e-3, max_depth=10):
+    mask_1 = gt > min_depth
+    mask_2 = gt < max_depth
+    mask = np.logical_and(mask_1, mask_2)
+
+    gt = gt[mask]
+    pred = pred[mask]
+
+    thresh = np.maximum((gt / pred), (pred / gt))
+    d1 = (thresh < 1.25).mean()
+    d2 = (thresh < 1.25**2).mean()
+    d3 = (thresh < 1.25**3).mean()
+
+    abs_rel = np.mean(np.abs(gt - pred) / gt)
+    sq_rel = np.mean(((gt - pred) ** 2) / gt)
+
+    rmse = (gt - pred) ** 2
+    rmse = np.sqrt(rmse.mean())
+
+    rmse_log = (np.log(gt) - np.log(pred)) ** 2
+    rmse_log = np.sqrt(rmse_log.mean())
+
+    err = np.log(pred) - np.log(gt)
+    silog = np.sqrt(np.mean(err**2) - np.mean(err) ** 2) * 100
+
+    log_10 = (np.abs(np.log10(gt) - np.log10(pred))).mean()
+    return dict(
+        d1=d1,
+        d2=d2,
+        d3=d3,
+        abs_rel=abs_rel,
+        rmse=rmse,
+        log_10=log_10,
+        rmse_log=rmse_log,
+        silog=silog,
+        sq_rel=sq_rel,
+    )
+
+
+class UpsamplerEvaluator(nn.Module):
+
+    def __init__(self, model, backbone, device, cfg, writer, console):
+        super().__init__()
+        self.model, self.backbone, self.device, self.cfg, self.writer, self.console = (
+            model,
+            backbone,
+            device,
+            cfg,
+            writer,
+            console,
+        )
+
+        # Use backbone-specific normalization (like train_jafar.py)
+        self.mean_bck = backbone.config["mean"]
+        self.std_bck = backbone.config["std"]
+
+        # Upsampler normalization (ImageNet standard)
+        self.mean_ups = (0.485, 0.456, 0.406)
+        self.std_ups = (0.229, 0.224, 0.225)
+
+        # Initialize depth estimation components
+        self.classifier = nn.Conv2d(backbone.embed_dim, 256, kernel_size=1, padding=0, stride=1).to(device)
+        self.depthloss = DepthLoss(max_depth=cfg.metrics.depth.max_depth)
+        self.register_buffer("bins", torch.linspace(cfg.metrics.depth.min_depth, cfg.metrics.depth.max_depth, 256))
+
+    def set_up_classifier(self, checkpoint_path):
+        """Load classifier weights from a checkpoint."""
+        if Path(checkpoint_path).exists():
+            checkpoint = torch.load(checkpoint_path)
+            self.classifier.load_state_dict(checkpoint)
+            self.console.print(f"Loaded classifier from checkpoint: {checkpoint_path}")
+        else:
+            raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
+
+    def set_optimizer(self, cfg, loader):
+        params = []
+        params_classifier = self.classifier.parameters()
+        params_model = self.model.parameters()
+
+        params = list(params_classifier)
+        if self.cfg.eval.supervise_model == True:
+            params += params_model
+        optimizer = instantiate(cfg.optimizer, params=params)
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=cfg.num_epochs * len(loader))
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+
+        num_params = sum(p.numel() for p in params if p.requires_grad)
+        self.log_print(f"[bold cyan]Number of optimized parameters: {num_params:,}[/bold cyan]")
+
+    def log_print(self, *args, **kwargs):
+        Console(force_terminal=True).print(*args, **kwargs)
+        self.console.print(*args, **kwargs)
+        if hasattr(self.console, "file") and self.console.file:
+            self.console.file.flush()
+
+    def log_tensorboard(self, step, loss=None, metrics=None):
+        if loss is not None:
+            self.writer.add_scalar("Loss/Step", loss, step)
+        if metrics is not None:
+            for metric_name, metric_value in metrics.items():
+                self.writer.add_scalar(f"Metrics/{metric_name}", metric_value, step)
+
+    def process_batch(self, image_batch, target, is_training=True, batch_idx=0):
+        H, W = target.shape[-2:]
+
+        # Apply separate normalization like in train_jafar.py
+        image_batch_bck = T.functional.normalize(image_batch, mean=self.mean_bck, std=self.std_bck)
+        image_batch_ups = T.functional.normalize(image_batch, mean=self.mean_ups, std=self.std_ups)
+
+        with torch.no_grad():
+            pred = self.backbone(image_batch_bck)
+
+        if self.cfg.eval.supervise_model:
+            pred = self.model(image_batch_ups, pred, (H, W))
+        else:
+            with torch.no_grad():
+                pred = self.model(image_batch_ups, pred, (H, W)).detach()
+
+        if batch_idx == 0:
+            self.log_print(f"[red]Model output size {pred.shape[-2:]} for expected size {(H, W)}[/red]")
+        pred = self.classifier(pred)
+
+        if pred.shape[-2:] != (H, W):
+            pred = F.interpolate(pred, size=(H, W), mode="bilinear")
+
+        # https://github.com/mbanani/probe3d/blob/c52d00b069d949b2f00c544d4991716df68d5233/evals/models/probes.py#L108
+        pred = F.relu(pred)
+        eps = 0.1
+        pred = pred + eps
+        pred = pred / pred.sum(dim=1, keepdim=True)
+        pred = torch.einsum("ikmn,k->imn", [pred, self.bins]).unsqueeze(dim=1)
+        return pred, target
+
+    def train(
+        self,
+        train_dataloader,
+        progress,
+        epoch,
+        start_time,
+    ):
+        self.log_print(f"[yellow]Training model epoch {epoch+1}...[/yellow]")
+        self.backbone.eval()
+        if self.cfg.eval.supervise_model:
+            self.model.train()
+        else:
+            self.model.eval()
+        self.classifier.train()
+
+        epoch_task = progress.add_task(
+            f"Epoch {epoch+1}/{self.cfg.num_epochs}",
+            total=len(train_dataloader),
+            loss=0.0,
+            step=0,
+        )
+        total_loss = 0
+
+        for batch_idx, batch in enumerate(train_dataloader):
+            # Process batch using get_batch
+            batch = get_batch(batch, self.device)
+            image_batch = batch["image"]
+            target = batch["depth"].to(self.device)
+
+            if random.random() < 0.5:
+                image_batch = torch.flip(image_batch, dims=[3])
+                target = torch.flip(target, dims=[3])
+
+            self.optimizer.zero_grad()
+
+            pred, target = self.process_batch(image_batch, target, is_training=True, batch_idx=batch_idx)
+
+            loss = self.depthloss(pred, target)
+
+            loss.backward()
+            self.optimizer.step()
+            total_loss += loss.item()
+
+            avg_loss = total_loss / (batch_idx + 1)
+
+            if (batch_idx + 1) % LOG_INTERVAL == 0 or batch_idx == len(train_dataloader) - 1:
+                elapsed_time = datetime.datetime.now() - start_time
+                elapsed_str = str(elapsed_time).split(".")[0]
+                current_lr = self.optimizer.param_groups[0]["lr"]
+
+                progress.update(epoch_task, advance=LOG_INTERVAL, loss=avg_loss, step=batch_idx + 1)
+
+                self.log_print(
+                    f"[cyan]Iteration {batch_idx + 1}[/cyan] - "
+                    f"Loss: {avg_loss:.6f} - "
+                    f"LR: {current_lr:.5e} - "
+                    f"Elapsed Time: {elapsed_str}"
+                )
+                if self.console and hasattr(self.console, "file"):
+                    self.console.file.flush()
+
+                progress.refresh()
+
+                self.log_tensorboard(epoch * len(train_dataloader) + batch_idx, loss=avg_loss)
+
+            if self.cfg.sanity and batch_idx == 0:
+                break
+
+            self.scheduler.step()
+
+            if self.cfg.sanity:
+                break
+
+        current_lr = self.optimizer.param_groups[0]["lr"]
+        self.log_print(
+            f"[bold cyan]Epoch {epoch+1} Summary:[/bold cyan] " f"Loss = {avg_loss:.6f} - " f"LR = {current_lr:.2e}"
+        )
+
+        return
+
+    def save_checkpoint(self, checkpoint_path):
+        console = self.console
+        # Save final model state after training
+        torch.save(self.classifier.state_dict(), checkpoint_path)
+        self.log_print(f"[bold green]Training completed. Model saved at: {checkpoint_path}[/bold green]")
+        console.file.flush()
+        return
+
+    @torch.inference_mode()
+    def evaluate(self, dataloader, epoch):
+        self.log_print("[yellow]Evaluating model...[/yellow]")
+        torch.cuda.empty_cache()
+
+        self.backbone.eval()
+        self.model.eval()
+        self.classifier.eval()
+
+        # Initialize depth metrics
+        results = {"d1": 0, "d2": 0, "d3": 0, "abs_rel": 0, "sq_rel": 0, "rmse": 0, "rmse_log": 0, "log_10": 0, "silog": 0}
+
+        nsamples = 0
+
+        for batch_idx, batch in enumerate(tqdm(dataloader, desc="Evaluating")):
+            batch = get_batch(batch, self.device)
+            image_batch = batch["image"]
+            target = batch["depth"].to(self.device)
+
+            # Process batch and get masked predictions and targets
+            pred, target = self.process_batch(image_batch, target, is_training=False, batch_idx=batch_idx)
+
+            cur_results = eval_metrics(
+                target.cpu().detach().numpy(),
+                pred.cpu().detach().numpy(),
+                self.cfg.metrics.depth.min_depth,
+                self.cfg.metrics.depth.max_depth,
+            )
+
+            for k in results.keys():
+                results[k] += cur_results[k]
+            nsamples += 1
+
+            if self.cfg.sanity and batch_idx == 0:
+                break
+
+        metrics = {}
+        for k in results.keys():
+            metrics[k] = results[k] / nsamples
+
+        # Log metrics to TensorBoard
+        self.log_tensorboard(step=epoch, metrics=metrics)
+
+        self.log_print(f"[bold green]Results: {metrics}[/bold green]")
+        return
+
+    @torch.inference_mode()
+    def simple_inference(self, image_batch):
+        self.backbone.eval()
+        self.model.eval()
+        self.classifier.eval()
+
+        H, W = image_batch.shape[-2:]
+        with torch.no_grad():
+            lr_feats = self.backbone(image_batch)
+            hr_feats = self.model(image_batch, lr_feats, (H, W))
+
+        labels = self.classifier(hr_feats)  # Pass through the classifier
+
+        labels = F.relu(labels)
+        eps = 0.1
+        labels = labels + eps
+        labels = labels / labels.sum(dim=1, keepdim=True)
+        labels = torch.einsum("ikmn,k->imn", [labels, self.bins]).unsqueeze(dim=1)
+        return hr_feats, labels, lr_feats
+
+
+@hydra.main(config_path="../config", config_name="eval_probing")
+def main(cfg):
+    # Use Hydra's current working directory for all outputs
+    task = "depth"
+    model_name = cfg.model.name if hasattr(cfg.model, "name") else "base"
+
+    # Get Hydra's current working directory
+    hydra_cfg = HydraConfig.get()
+    current_run_dir = hydra_cfg.runtime.output_dir
+
+    # Setup Classifier checkpoint in current Hydra directory
+    checkpoint_filename = "linear_probe.pth"
+    checkpoint_path = os.path.join(current_run_dir, checkpoint_filename)
+
+    # Create persistent consoles instead of creating new ones each time
+    terminal_console = Console()  # Terminal output
+    tag = ""
+    if cfg.eval.model_ckpt:
+        try:
+            tag = cfg.eval.model_ckpt.split("/")[-3]
+        except:
+            tag = ""
+
+    # Create log file names
+    if Path(checkpoint_path).exists():
+        file_name = f"eval_{model_name}_{tag}_{task}_prob_from_ups.log"
+    else:
+        file_name = f"train_{model_name}_{tag}_{task}_prob_from_ups.log"
+
+    # All outputs in Hydra's current directory
+    log_file_path = os.path.join(current_run_dir, file_name)
+    tb_dir = os.path.join(current_run_dir, "tb")
+    os.makedirs(tb_dir, exist_ok=True)
+
+    file_console = Console(file=open(log_file_path, "w"))
+
+    # Initialize TensorBoard writer in current Hydra directory
+    writer = SummaryWriter(log_dir=tb_dir)
+
+    def log_print(*args, **kwargs):
+        """Log to both terminal and file with immediate flushing"""
+        terminal_console.print(*args, **kwargs)
+        file_console.print(*args, **kwargs)
+        file_console.file.flush()
+
+    # Start logging
+    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    log_print(f"\n[bold blue]{'='*50}[/bold blue]")
+    log_print(f"[bold blue]Starting at {timestamp}[/bold blue]")
+    log_print(f"[bold green]Configuration:[/bold green]")
+    log_print(OmegaConf.to_yaml(cfg))
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    log_print(f"[bold yellow]Using device: {device}[/bold yellow]")
+
+    log_print(f"\n[bold cyan]Processing {task} task:[/bold cyan]")
+    log_print(f"\n[bold cyan]Image size: {cfg.img_size}[/bold cyan]")
+
+    # Setup Backbone
+    backbones, backbone_names, *_ = load_multiple_backbones(cfg, cfg.backbone, device)
+    backbone = backbones[0]
+    backbone.requires_grad_(False)
+    backbone.eval()
+
+    # Setup Model
+    model = instantiate(cfg.model).to(device)
+    if cfg.eval.model_ckpt:
+        checkpoint = torch.load(cfg.eval.model_ckpt, map_location=device, weights_only=False)
+        if cfg.model.name == "featup":
+            new_ckpts = {
+                k.replace("model.1.", "norm."): v
+                for k, v in checkpoint["state_dict"].items()
+                if "upsampler" in k or "model.1.norm" in k
+            }
+            model.load_state_dict(new_ckpts, strict=True)
+        else:
+            model.load_state_dict(checkpoint, strict=True)
+        log_print(f"[green]Loaded model from checkpoint: {cfg.eval.model_ckpt}[/green]")
+    else:
+        log_print(f"[yellow]No model checkpoint provided, using untrained model[/yellow]")
+
+    if cfg.eval.supervise_model == False:
+        model.requires_grad_(False)
+        model.eval()
+    else:
+        model.requires_grad_(True)
+        model.train()
+
+    # Setup Dataloaders - Modified to match train_jafar.py approach
+    train_loader, val_loader = get_dataloaders(cfg, shuffle=False)
+    log_print(f"[bold cyan]Train Dataset size: {len(train_loader.dataset)}[/bold cyan]")
+    log_print(f"[bold cyan]Val Dataset size: {len(val_loader.dataset)}[/bold cyan]")
+
+    # Setup Evaluator
+    evaluator = UpsamplerEvaluator(model, backbone, device, cfg, writer, file_console)
+    evaluator.to(device)
+
+    log_print("[yellow]Training classifier...[/yellow]\n")
+    evaluator.set_optimizer(cfg, train_loader)
+
+    progress = Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(),
+        TaskProgressColumn(),
+        TextColumn("[yellow]Loss: {task.fields[loss]:.6f}"),
+        TextColumn("[green]Step: {task.fields[step]:.6e}"),
+        console=file_console,
+    )
+
+    start_time = datetime.datetime.now()
+
+    with progress:
+        # Standard training for upsampler
+        log_print(f"[yellow]Training for {cfg.num_epochs} epochs[/yellow]\n")
+
+        for epoch in range(cfg.num_epochs):
+            evaluator.train(train_loader, progress, epoch, start_time)
+            evaluator.evaluate(val_loader, epoch)
+
+        evaluator.save_checkpoint(checkpoint_path)
+
+    file_console.file.close()
+    writer.close()  # Close TensorBoard writer
+
+
+if __name__ == "__main__":
+    main()

Pixal3D/torch_hub/hub/valeoai_NAF_main/evaluation/eval_seg_probing.py

@@ -0,0 +1,393 @@
+import datetime
+import os
+import sys
+from pathlib import Path
+
+import hydra
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as T
+from einops import rearrange
+from hydra.core.global_hydra import GlobalHydra
+from hydra.core.hydra_config import HydraConfig
+from hydra.utils import instantiate
+from omegaconf import OmegaConf
+from rich.console import Console
+from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn
+from torch.utils.tensorboard import SummaryWriter
+from torchmetrics.classification import Accuracy, JaccardIndex
+from tqdm import tqdm
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import random
+
+from utils.training import get_batch, get_dataloaders, load_multiple_backbones
+
+LOG_INTERVAL = 100
+IGNORE = 255
+
+
+class UpsamplerEvaluator:
+
+    def __init__(self, model, backbone, device, cfg, writer, console):
+        self.model, self.backbone, self.device, self.cfg, self.writer, self.console = (
+            model,
+            backbone,
+            device,
+            cfg,
+            writer,
+            console,
+        )
+
+        # Use backbone-specific normalization (like train_jafar.py)
+        self.mean_bck = backbone.config["mean"]
+        self.std_bck = backbone.config["std"]
+
+        # Upsampler normalization (ImageNet standard)
+        self.mean_ups = (0.485, 0.456, 0.406)
+        self.std_ups = (0.229, 0.224, 0.225)
+
+        # Initialize task-specific components
+        self.accuracy_metric = Accuracy(num_classes=cfg.metrics.seg.num_classes, task="multiclass").to(device)
+        self.iou_metric = JaccardIndex(num_classes=cfg.metrics.seg.num_classes, task="multiclass").to(device)
+        self.classifier = nn.Conv2d(backbone.embed_dim, cfg.metrics.seg.num_classes, 1).to(device)
+
+    def set_up_classifier(self, checkpoint_path):
+        """Load classifier weights from a checkpoint."""
+        if Path(checkpoint_path).exists():
+            checkpoint = torch.load(checkpoint_path)
+            self.classifier.load_state_dict(checkpoint)
+            self.console.print(f"Loaded classifier from checkpoint: {checkpoint_path}")
+        else:
+            raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
+
+    def set_optimizer(self, cfg, loader):
+        params = []
+        params_classifier = self.classifier.parameters()
+        params_model = self.model.parameters()
+
+        params = list(params_classifier)
+        optimizer = instantiate(cfg.optimizer, params=params)
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=cfg.num_epochs * len(loader))
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+
+        num_params = sum(p.numel() for p in params if p.requires_grad)
+        self.log_print(f"[bold cyan]Number of optimized parameters: {num_params:,}[/bold cyan]")
+
+    def log_print(self, *args, **kwargs):
+        Console(force_terminal=True).print(*args, **kwargs)
+        self.console.print(*args, **kwargs)
+        if hasattr(self.console, "file") and self.console.file:
+            self.console.file.flush()
+
+    def log_tensorboard(self, step, loss=None, metrics=None):
+        if loss is not None:
+            self.writer.add_scalar("Loss/Step", loss, step)
+        if metrics is not None:
+            for metric_name, metric_value in metrics.items():
+                self.writer.add_scalar(f"Metrics/{metric_name}", metric_value, step)
+
+    def process_batch(self, image_batch, target, batch_idx=0):
+        H, W = target.shape[-2:]
+
+        image_batch_bck = T.functional.normalize(image_batch, mean=self.mean_bck, std=self.std_bck)
+        image_batch_ups = T.functional.normalize(image_batch, mean=self.mean_ups, std=self.std_ups)
+
+        with torch.no_grad():
+            pred = self.backbone(image_batch_bck)
+
+        input_shape = pred.shape[-2:]
+        with torch.no_grad():
+            pred = self.model(image_batch_ups, pred, (H, W)).detach()
+
+        if batch_idx == 0:
+            self.log_print(
+                f"[red]Model input size {input_shape}, output size {pred.shape[-2:]} for expected size {(H, W)}[/red]"
+            )
+        pred = self.classifier(pred)
+
+        if pred.shape[-2:] != (H, W):
+            pred = F.interpolate(pred, size=(H, W), mode="bilinear")
+
+        if target.shape[-2:] != pred.shape[-2:]:
+            target = (
+                F.interpolate(
+                    target.unsqueeze(1),
+                    size=pred.shape[-2:],
+                    mode="nearest-exact",
+                )
+                .squeeze(1)
+                .to(target.dtype)
+            )
+
+        valid_mask = target != IGNORE
+
+        pred = rearrange(pred, "b c h w -> (b h w) c")
+        target = rearrange(target, "b h w -> (b h w)")
+        valid_mask = rearrange(valid_mask, "b h w -> (b h w)")
+
+        pred = pred[valid_mask]
+        target = target[valid_mask]
+        return pred, target
+
+    def train(
+        self,
+        train_dataloader,
+        progress,
+        epoch,
+        start_time,
+    ):
+        self.log_print(f"[yellow]Training model epoch {epoch+1}...[/yellow]")
+        self.backbone.eval()
+        self.model.eval()
+        self.classifier.train()
+
+        epoch_task = progress.add_task(
+            f"Epoch {epoch+1}/{self.cfg.num_epochs}",
+            total=len(train_dataloader),
+            loss=0.0,
+            step=0,
+        )
+        total_loss = 0
+
+        for batch_idx, batch in enumerate(train_dataloader):
+            batch = get_batch(batch, self.device)
+            image_batch = batch["image"]
+            target = batch["label"].to(self.device)
+
+            if random.random() < 0.5:
+                assert len(image_batch.shape) == 4 and len(target.shape) == 3
+                image_batch = torch.flip(image_batch, dims=[3])
+                target = torch.flip(target, dims=[2])
+
+            self.optimizer.zero_grad()
+
+            pred, target = self.process_batch(image_batch, target, batch_idx=batch_idx)
+
+            loss = F.cross_entropy(pred, target)
+
+            loss.backward()
+            self.optimizer.step()
+            total_loss += loss.item()
+
+            avg_loss = total_loss / (batch_idx + 1)
+
+            if (batch_idx + 1) % LOG_INTERVAL == 0 or batch_idx == len(train_dataloader) - 1:
+                elapsed_time = datetime.datetime.now() - start_time
+                elapsed_str = str(elapsed_time).split(".")[0]
+                current_lr = self.optimizer.param_groups[0]["lr"]
+
+                progress.update(epoch_task, advance=LOG_INTERVAL, loss=avg_loss, step=batch_idx + 1)
+
+                self.log_print(
+                    f"[cyan]Iteration {batch_idx + 1}[/cyan] - "
+                    f"Loss: {avg_loss:.6f} - "
+                    f"LR: {current_lr:.5e} - "
+                    f"Elapsed Time: {elapsed_str}"
+                )
+                if self.console and hasattr(self.console, "file"):
+                    self.console.file.flush()
+
+                progress.refresh()
+
+                self.log_tensorboard(epoch * len(train_dataloader) + batch_idx, loss=avg_loss)
+
+            if self.cfg.sanity and batch_idx == 0:
+                break
+
+            self.scheduler.step()
+
+            if self.cfg.sanity:
+                break
+
+        current_lr = self.optimizer.param_groups[0]["lr"]
+        self.log_print(
+            f"[bold cyan]Epoch {epoch+1} Summary:[/bold cyan] " f"Loss = {avg_loss:.6f} - " f"LR = {current_lr:.2e}"
+        )
+
+        return
+
+    def save_checkpoint(self, checkpoint_path):
+        console = self.console
+        torch.save(self.classifier.state_dict(), checkpoint_path)
+        self.log_print(f"[bold green]Training completed. Model saved at: {checkpoint_path}[/bold green]")
+        console.file.flush()
+        return
+
+    @torch.inference_mode()
+    def evaluate(self, dataloader, epoch):
+        self.log_print("[yellow]Evaluating model...[/yellow]")
+        torch.cuda.empty_cache()
+
+        self.backbone.eval()
+        self.model.eval()
+        self.classifier.eval()
+
+        self.accuracy_metric.reset()
+        self.iou_metric.reset()
+
+        for batch_idx, batch in enumerate(tqdm(dataloader, desc="Evaluating")):
+            batch = get_batch(batch, self.device)
+            image_batch = batch["image"]
+            target = batch["label"].to(self.device)
+
+            pred, target = self.process_batch(image_batch, target, batch_idx=batch_idx)
+
+            self.accuracy_metric(pred, target)
+            self.iou_metric(pred, target)
+
+            if self.cfg.sanity and batch_idx == 0:
+                break
+
+        metrics = {}
+        metrics.update(
+            {
+                "accuracy": self.accuracy_metric.compute().item(),
+                "iou": self.iou_metric.compute().item(),
+            }
+        )
+
+        self.log_tensorboard(step=epoch, metrics=metrics)
+
+        self.log_print(f"[bold green]Results: {metrics}[/bold green]")
+        return
+
+    @torch.inference_mode()
+    def simple_inference(self, image_batch):
+        self.backbone.eval()
+        self.model.eval()
+        self.classifier.eval()
+
+        H, W = image_batch.shape[-2:]
+        image_batch_bck = T.functional.normalize(image_batch, mean=self.mean_bck, std=self.std_bck)
+        image_batch_ups = T.functional.normalize(image_batch, mean=self.mean_ups, std=self.std_ups)
+
+        with torch.no_grad():
+            lr_feats = self.backbone(image_batch_bck)
+            hr_feats = self.model(image_batch_ups, lr_feats, (H, W))
+
+        labels = self.classifier(hr_feats).argmax(dim=1, keepdim=True)
+        return hr_feats, labels, lr_feats
+
+
+@hydra.main(config_path="../config", config_name="eval_probing")
+def main(cfg):
+    hydra_cfg = HydraConfig.get()
+    current_run_dir = hydra_cfg.runtime.output_dir
+
+    task = cfg.eval.task
+    model_name = cfg.model.name if hasattr(cfg.model, "name") else "base"
+
+    # Setup Classifier checkpoint in current Hydra directory
+    checkpoint_filename = "linear_probe.pth"
+    checkpoint_path = os.path.join(current_run_dir, checkpoint_filename)
+
+    # Create persistent consoles instead of creating new ones each time
+    terminal_console = Console()
+    tag = ""
+    if cfg.eval.model_ckpt:
+        try:
+            tag = cfg.eval.model_ckpt.split("/")[-3]
+        except:
+            tag = ""
+
+    # Create log file names
+    file_name = f"train_{model_name}_{tag}_{task}.log"
+
+    # All outputs in Hydra's current directory
+    log_file_path = os.path.join(current_run_dir, file_name)
+    tb_dir = os.path.join(current_run_dir, "tb")
+    os.makedirs(tb_dir, exist_ok=True)
+
+    file_console = Console(file=open(log_file_path, "w"))
+
+    # Initialize TensorBoard writer in current Hydra directory
+    writer = SummaryWriter(log_dir=tb_dir)
+
+    def log_print(*args, **kwargs):
+        """Log to both terminal and file with immediate flushing"""
+        terminal_console.print(*args, **kwargs)
+        file_console.print(*args, **kwargs)
+        file_console.file.flush()
+
+    # Start logging
+    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    log_print(f"\n[bold blue]{'='*50}[/bold blue]")
+    log_print(f"[bold blue]Starting at {timestamp}[/bold blue]")
+    log_print(f"[bold green]Configuration:[/bold green]")
+    log_print(OmegaConf.to_yaml(cfg))
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    log_print(f"[bold yellow]Using device: {device}[/bold yellow]")
+
+    log_print(f"\n[bold cyan]Processing {task} task:[/bold cyan]")
+    log_print(f"\n[bold cyan]Image size: {cfg.img_size}[/bold cyan]")
+
+    # Setup Backbone
+    backbones, *_ = load_multiple_backbones(cfg, cfg.backbone, device)
+    backbone = backbones[0]
+    backbone.requires_grad_(False)
+    backbone.eval()
+
+    # Setup Model
+    model = instantiate(cfg.model).to(device)
+    if cfg.eval.model_ckpt:
+        checkpoint = torch.load(cfg.eval.model_ckpt, map_location=device, weights_only=False)
+        if cfg.model.name == "featup":
+            new_ckpts = {
+                k.replace("model.1.", "norm."): v
+                for k, v in checkpoint["state_dict"].items()
+                if "upsampler" in k or "model.1.norm" in k
+            }
+            model.load_state_dict(new_ckpts, strict=True)
+        else:
+            model.load_state_dict(checkpoint, strict=True)
+        log_print(f"[green]Loaded model from checkpoint: {cfg.eval.model_ckpt}[/green]")
+    else:
+        log_print(f"[yellow]No model checkpoint provided, using untrained model[/yellow]")
+
+    model.requires_grad_(True)
+    model.train()
+
+    # Setup Dataloaders - Modified to match train_jafar.py approach
+    train_loader, val_loader = get_dataloaders(cfg, shuffle=False)
+    log_print(f"[bold cyan]Train Dataset size: {len(train_loader.dataset)}[/bold cyan]")
+    log_print(f"[bold cyan]Val Dataset size: {len(val_loader.dataset)}[/bold cyan]")
+
+    # Setup Evaluator
+    evaluator = UpsamplerEvaluator(model, backbone, device, cfg, writer, file_console)
+
+    log_print("[yellow]Training classifier...[/yellow]\n")
+    evaluator.set_optimizer(cfg, loader=train_loader)
+
+    progress = Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(),
+        TaskProgressColumn(),
+        TextColumn("[yellow]Loss: {task.fields[loss]:.6f}"),
+        TextColumn("[green]Step: {task.fields[step]:.6e}"),
+        console=file_console,
+    )
+
+    start_time = datetime.datetime.now()
+
+    with progress:
+        log_print(f"[yellow]Training for {cfg.num_epochs} epochs[/yellow]\n")
+
+        for epoch in range(cfg.num_epochs):
+            evaluator.train(train_loader, progress, epoch, start_time)
+            evaluator.evaluate(val_loader, epoch)
+
+        evaluator.save_checkpoint(checkpoint_path)
+
+    file_console.file.close()
+    writer.close()
+
+
+if __name__ == "__main__":
+    main()

Pixal3D/torch_hub/hub/valeoai_NAF_main/evaluation/eval_video_seg.py

@@ -0,0 +1,816 @@
+"""
+Evaluate JAFAR performance on DAVIS video segmentation. Based on https://github.com/davisvideochallenge/davis2017-evaluation and https://github.com/facebookresearch/dino/blob/main/eval_video_segmentation.py
+
+Code adapted by: Saksham Suri and Matthew Walmer
+
+Sample Usage:
+python eval_davis.py
+python eval_davis.py backbone.name=vit_small_patch14_reg4_dinov2 model=jafar
+"""
+
+import argparse
+import copy
+import datetime
+import glob
+import math
+import os
+import queue
+import sys
+import time
+import warnings
+from collections import defaultdict
+from urllib.request import urlopen
+
+import cv2
+import hydra
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import torch
+import torch.multiprocessing
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from hydra.core.hydra_config import HydraConfig
+from hydra.utils import instantiate
+from omegaconf import DictConfig, OmegaConf
+from PIL import Image
+from rich.console import Console
+from torch import nn
+from torch.nn import functional as F
+from torchvision import transforms
+from tqdm import tqdm
+
+warnings.filterwarnings("ignore", category=RuntimeWarning)
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from utils.training import load_multiple_backbones
+
+torch.multiprocessing.set_sharing_strategy("file_system")
+
+
+# ============================================================================
+# DAVIS Evaluation Classes and Functions (Self-contained)
+# ============================================================================
+
+
+class DAVISDataset(object):
+    """Simplified DAVIS dataset class for evaluation"""
+
+    SUBSET_OPTIONS = ["train", "val", "test-dev", "test-challenge"]
+    TASKS = ["semi-supervised", "unsupervised"]
+    VOID_LABEL = 255
+
+    def __init__(self, root, task="semi-supervised", subset="val", sequences="all", resolution="480p"):
+        self.task = task
+        self.subset = subset
+        self.root = root
+        self.img_path = os.path.join(self.root, "JPEGImages", resolution)
+        self.mask_path = os.path.join(self.root, "Annotations", resolution)
+        self.imagesets_path = os.path.join(self.root, "ImageSets", "2017")
+
+        if sequences == "all":
+            with open(os.path.join(self.imagesets_path, f"{self.subset}.txt"), "r") as f:
+                tmp = f.readlines()
+            sequences_names = [x.strip() for x in tmp]
+        else:
+            sequences_names = sequences if isinstance(sequences, list) else [sequences]
+
+        self.sequences = defaultdict(dict)
+        for seq in sequences_names:
+            images = np.sort(glob.glob(os.path.join(self.img_path, seq, "*.jpg"))).tolist()
+            self.sequences[seq]["images"] = images
+            masks = np.sort(glob.glob(os.path.join(self.mask_path, seq, "*.png"))).tolist()
+            self.sequences[seq]["masks"] = masks
+
+    def get_sequences(self):
+        return list(self.sequences.keys())
+
+    def get_all_masks(self, sequence, separate_objects_masks=False):
+        """Get all ground truth masks for a sequence"""
+        masks = []
+        mask_ids = []
+
+        for mask_path in self.sequences[sequence]["masks"]:
+            mask = np.array(Image.open(mask_path))
+            masks.append(mask)
+            mask_ids.append(os.path.splitext(os.path.basename(mask_path))[0])
+
+        masks = np.array(masks)
+
+        if separate_objects_masks:
+            # Convert to separate object masks
+            num_objects = int(np.max(masks))
+            if num_objects > 0:
+                object_masks = np.zeros((num_objects, *masks.shape))
+                for obj_id in range(1, num_objects + 1):
+                    object_masks[obj_id - 1] = (masks == obj_id).astype(np.uint8)
+                void_masks = np.zeros_like(masks[0])  # No void pixels for simplicity
+                return object_masks, void_masks, mask_ids
+
+        return masks, mask_ids
+
+
+class DAVISResults(object):
+    """Results reader for DAVIS evaluation"""
+
+    def __init__(self, root_dir):
+        self.root_dir = root_dir
+
+    def _read_mask(self, sequence, frame_id):
+        mask_path = os.path.join(self.root_dir, sequence, f"{frame_id}.png")
+        if not os.path.exists(mask_path):
+            raise FileNotFoundError(f"Mask not found: {mask_path}")
+        return np.array(Image.open(mask_path))
+
+    def read_masks(self, sequence, masks_id):
+        """Read all prediction masks for a sequence"""
+        mask_0 = self._read_mask(sequence, masks_id[0])
+        masks = np.zeros((len(masks_id), *mask_0.shape))
+        for ii, m in enumerate(masks_id):
+            masks[ii, ...] = self._read_mask(sequence, m)
+
+        # Convert to separate object masks
+        num_objects = int(np.max(masks))
+        if num_objects == 0:
+            return np.zeros((1, *masks.shape))
+
+        tmp = np.ones((num_objects, *masks.shape))
+        tmp = tmp * np.arange(1, num_objects + 1)[:, None, None, None]
+        masks = (tmp == masks[None, ...]) > 0
+        return masks
+
+
+def davis_db_eval_iou(annotation, segmentation, void_pixels=None):
+    """Compute region similarity as the Jaccard Index"""
+    assert (
+        annotation.shape == segmentation.shape
+    ), f"Annotation({annotation.shape}) and segmentation:{segmentation.shape} dimensions do not match."
+
+    annotation = annotation.astype(bool)
+    segmentation = segmentation.astype(bool)
+
+    if void_pixels is not None:
+        void_pixels = void_pixels.astype(bool)
+    else:
+        void_pixels = np.zeros_like(segmentation)
+
+    # Intersection between all sets
+    inters = np.sum((segmentation & annotation) & np.logical_not(void_pixels), axis=(-2, -1))
+    union = np.sum((segmentation | annotation) & np.logical_not(void_pixels), axis=(-2, -1))
+
+    j = inters / union
+    if j.ndim == 0:
+        j = 1 if np.isclose(union, 0) else j
+    else:
+        j[np.isclose(union, 0)] = 1
+    return j
+
+
+def davis_f_measure(foreground_mask, gt_mask, void_pixels=None, bound_th=0.008):
+    """Compute F-measure for boundary accuracy"""
+    assert foreground_mask.shape == gt_mask.shape
+
+    if void_pixels is not None:
+        assert foreground_mask.shape == void_pixels.shape
+        # Apply void mask
+        foreground_mask = foreground_mask.copy()
+        gt_mask = gt_mask.copy()
+        foreground_mask[void_pixels] = 0
+        gt_mask[void_pixels] = 0
+
+    bound_pix = bound_th if bound_th >= 1 else np.ceil(bound_th * np.linalg.norm(foreground_mask.shape))
+
+    # Get the pixel boundaries of both masks
+    fg_boundary = _seg2bmap(foreground_mask)
+    gt_boundary = _seg2bmap(gt_mask)
+
+    from scipy.ndimage import distance_transform_edt
+
+    # Distance transform
+    fg_dist = distance_transform_edt(1 - fg_boundary)
+    gt_dist = distance_transform_edt(1 - gt_boundary)
+
+    # Precision: how many predicted boundary pixels are close to gt
+    precision = np.sum(fg_boundary * (gt_dist <= bound_pix)) / (np.sum(fg_boundary) + 1e-10)
+
+    # Recall: how many gt boundary pixels are close to predicted
+    recall = np.sum(gt_boundary * (fg_dist <= bound_pix)) / (np.sum(gt_boundary) + 1e-10)
+
+    # F-measure
+    if precision + recall == 0:
+        f_score = 0
+    else:
+        f_score = 2 * precision * recall / (precision + recall)
+
+    return f_score
+
+
+def _seg2bmap(seg, width=None, height=None):
+    """Convert segmentation to boundary map"""
+    seg = seg.astype(bool)
+    seg = seg.astype(np.uint8)
+
+    # Compute edges using Sobel operator
+    kernel_x = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=np.float32)
+    kernel_y = np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=np.float32)
+
+    edges_x = cv2.filter2D(seg.astype(np.float32), -1, kernel_x)
+    edges_y = cv2.filter2D(seg.astype(np.float32), -1, kernel_y)
+    edges = np.sqrt(edges_x**2 + edges_y**2)
+
+    # Threshold to get binary boundary
+    bmap = edges > 0.1
+
+    return bmap.astype(bool)
+
+
+def davis_db_eval_boundary(annotation, segmentation, void_pixels=None, bound_th=0.008):
+    """Compute boundary F-measure"""
+    assert annotation.shape == segmentation.shape
+    if void_pixels is not None:
+        assert annotation.shape == void_pixels.shape
+
+    if annotation.ndim == 3:
+        n_frames = annotation.shape[0]
+        f_res = np.zeros(n_frames)
+        for frame_id in range(n_frames):
+            void_pixels_frame = None if void_pixels is None else void_pixels[frame_id, :, :]
+            f_res[frame_id] = davis_f_measure(
+                segmentation[frame_id, :, :], annotation[frame_id, :, :], void_pixels_frame, bound_th=bound_th
+            )
+    elif annotation.ndim == 2:
+        f_res = davis_f_measure(segmentation, annotation, void_pixels, bound_th=bound_th)
+    else:
+        raise ValueError(f"Annotation and segmentation must be 2D or 3D arrays, got {annotation.ndim}D")
+
+    return f_res
+
+
+def davis_db_statistics(per_frame_values):
+    """Compute mean, recall and decay from per-frame evaluation"""
+    # Strip off nan values
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=RuntimeWarning)
+        M = np.nanmean(per_frame_values)
+        O = np.nanmean(per_frame_values > 0.5)
+
+    N_bins = 4
+    ids = np.round(np.linspace(1, len(per_frame_values), N_bins + 1) + 1e-10) - 1
+    ids = ids.astype(np.uint8)
+
+    D_bins = [per_frame_values[ids[i] : ids[i + 1] + 1] for i in range(0, 4)]
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=RuntimeWarning)
+        D = np.nanmean(D_bins[0]) - np.nanmean(D_bins[3])
+
+    return M, O, D
+
+
+class DAVISEvaluation(object):
+    """Complete DAVIS evaluation class"""
+
+    def __init__(self, davis_root, task="semi-supervised", gt_set="val", sequences="all"):
+        self.davis_root = davis_root
+        self.task = task
+        self.dataset = DAVISDataset(root=davis_root, task=task, subset=gt_set, sequences=sequences)
+
+    @staticmethod
+    def _evaluate_semisupervised(all_gt_masks, all_res_masks, all_void_masks, metric):
+        if all_res_masks.shape[0] > all_gt_masks.shape[0]:
+            print("\nIn your PNG files there is an index higher than the number of objects in the sequence!")
+            return None, None
+        elif all_res_masks.shape[0] < all_gt_masks.shape[0]:
+            zero_padding = np.zeros((all_gt_masks.shape[0] - all_res_masks.shape[0], *all_res_masks.shape[1:]))
+            all_res_masks = np.concatenate([all_res_masks, zero_padding], axis=0)
+
+        j_metrics_res, f_metrics_res = np.zeros(all_gt_masks.shape[:2]), np.zeros(all_gt_masks.shape[:2])
+        for ii in range(all_gt_masks.shape[0]):
+            if "J" in metric:
+                j_metrics_res[ii, :] = davis_db_eval_iou(all_gt_masks[ii, ...], all_res_masks[ii, ...], all_void_masks)
+            if "F" in metric:
+                f_metrics_res[ii, :] = davis_db_eval_boundary(all_gt_masks[ii, ...], all_res_masks[ii, ...], all_void_masks)
+        return j_metrics_res, f_metrics_res
+
+    def evaluate(self, res_path, metric=("J", "F"), debug=False):
+        """Run complete DAVIS evaluation"""
+        metric = metric if isinstance(metric, tuple) or isinstance(metric, list) else [metric]
+        if "T" in metric:
+            raise ValueError("Temporal metric not supported!")
+        if "J" not in metric and "F" not in metric:
+            raise ValueError("Metric possible values are J for IoU or F for Boundary")
+
+        # Containers
+        metrics_res = {}
+        if "J" in metric:
+            metrics_res["J"] = {"M": [], "R": [], "D": [], "M_per_object": {}}
+        if "F" in metric:
+            metrics_res["F"] = {"M": [], "R": [], "D": [], "M_per_object": {}}
+
+        # Sweep all sequences
+        results = DAVISResults(root_dir=res_path)
+        for seq in tqdm(list(self.dataset.get_sequences())):
+            try:
+                all_gt_masks, all_void_masks, all_masks_id = self.dataset.get_all_masks(seq, True)
+                if self.task == "semi-supervised":
+                    all_gt_masks, all_masks_id = all_gt_masks[:, 1:-1, :, :], all_masks_id[1:-1]
+                all_res_masks = results.read_masks(seq, all_masks_id)
+
+                if self.task == "semi-supervised":
+                    j_metrics_res, f_metrics_res = self._evaluate_semisupervised(all_gt_masks, all_res_masks, None, metric)
+                else:
+                    # For unsupervised, use semi-supervised for now (simplified)
+                    j_metrics_res, f_metrics_res = self._evaluate_semisupervised(all_gt_masks, all_res_masks, None, metric)
+
+                if j_metrics_res is None:
+                    continue
+
+                for ii in range(all_gt_masks.shape[0]):
+                    seq_name = f"{seq}_{ii+1}"
+                    if "J" in metric:
+                        JM, JR, JD = davis_db_statistics(j_metrics_res[ii])
+                        metrics_res["J"]["M"].append(JM)
+                        metrics_res["J"]["R"].append(JR)
+                        metrics_res["J"]["D"].append(JD)
+                        metrics_res["J"]["M_per_object"][seq_name] = JM
+                    if "F" in metric:
+                        FM, FR, FD = davis_db_statistics(f_metrics_res[ii])
+                        metrics_res["F"]["M"].append(FM)
+                        metrics_res["F"]["R"].append(FR)
+                        metrics_res["F"]["D"].append(FD)
+                        metrics_res["F"]["M_per_object"][seq_name] = FM
+
+                # Show progress
+                if debug:
+                    print(f"Evaluated sequence: {seq}")
+
+            except Exception as e:
+                print(f"Error evaluating sequence {seq}: {e}")
+                continue
+
+        return metrics_res
+
+
+@torch.no_grad()
+def eval_video_tracking_davis(
+    cfg, backbone, upsampler, transform, frame_list, video_dir, first_seg, seg_ori, color_palette, log_print
+):
+    """
+    Evaluate tracking on a video given first frame & segmentation
+    """
+    output_dir = HydraConfig.get().runtime.output_dir
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Create output directory based on model configuration
+    exp_name = f"davis_vidseg_{cfg.eval.ups_factor}_{cfg.model.name}_{cfg.backbone.name}"
+    output_subdir = os.path.join(output_dir, exp_name)
+    os.makedirs(output_subdir, exist_ok=True)
+
+    video_folder = os.path.join(output_subdir, video_dir.split("/")[-1])
+    os.makedirs(video_folder, exist_ok=True)
+
+    # Setup device and normalization
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Setup normalization transforms for backbone
+    mean_bck = backbone.config["mean"]
+    std_bck = backbone.config["std"]
+    mean_bck_tensor = torch.tensor(mean_bck, device=device).view(1, 3, 1, 1)
+    std_bck_tensor = torch.tensor(std_bck, device=device).view(1, 3, 1, 1)
+
+    # Upsampler normalization (ImageNet standard)
+    mean_ups = (0.485, 0.456, 0.406)
+    std_ups = (0.229, 0.224, 0.225)
+    mean_ups_tensor = torch.tensor(mean_ups, device=device).view(1, 3, 1, 1)
+    std_ups_tensor = torch.tensor(std_ups, device=device).view(1, 3, 1, 1)
+
+    # The queue stores the n preceding frames
+    que = queue.Queue(cfg.eval.n_last_frames)
+
+    # first frame
+    frame1, ori_h, ori_w = read_frame(frame_list[0], backbone.config["ps"])
+
+    # extract first frame feature using JAFAR pipeline
+    frame1_feat = extract_feature(
+        cfg,
+        backbone,
+        upsampler,
+        transform,
+        frame1,
+        mean_bck_tensor,
+        std_bck_tensor,
+        mean_ups_tensor,
+        std_ups_tensor,
+    )
+    frame1_feat = F.interpolate(frame1_feat, size=first_seg.shape[-2:], mode="bilinear", align_corners=False)
+    frame1_feat = rearrange(frame1_feat, "1 c h w -> c (h w)")
+
+    # saving first segmentation
+    out_path = os.path.join(video_folder, "00000.png")
+    imwrite_indexed(out_path, seg_ori, color_palette)
+    mask_neighborhood = None
+
+    for cnt in tqdm(range(1, len(frame_list))):
+        frame_tar = read_frame(frame_list[cnt], backbone.config["ps"])[0]
+
+        # we use the first segmentation and the n previous ones
+        used_frame_feats = [frame1_feat] + [pair[0] for pair in list(que.queue)]
+        used_segs = [first_seg] + [pair[1] for pair in list(que.queue)]
+        frame_tar_avg, feat_tar, mask_neighborhood = label_propagation(
+            cfg,
+            backbone,
+            upsampler,
+            transform,
+            frame_tar,
+            used_frame_feats,
+            used_segs,
+            mask_neighborhood,
+            mean_bck_tensor,
+            std_bck_tensor,
+            mean_ups_tensor,
+            std_ups_tensor,
+        )
+
+        # pop out oldest frame if necessary
+        if que.qsize() == cfg.eval.n_last_frames:
+            que.get()
+        # push current results into queue
+        seg = copy.deepcopy(frame_tar_avg)
+        que.put([feat_tar, seg])
+
+        # upsampling & argmax
+        frame_tar_avg = F.interpolate(
+            frame_tar_avg,
+            size=[v * backbone.config["ps"] // cfg.eval.ups_factor for v in frame_tar_avg.shape[-2:]],
+            mode="bilinear",
+            align_corners=False,
+            recompute_scale_factor=False,
+        )[0]
+        frame_tar_avg = frame_tar_avg.squeeze(0)
+        frame_tar_avg = norm_mask(frame_tar_avg)
+        _, frame_tar_seg = torch.max(frame_tar_avg, dim=0)
+
+        # saving to disk
+        frame_tar_seg = np.array(frame_tar_seg.squeeze().cpu(), dtype=np.uint8)
+        frame_tar_seg = np.array(Image.fromarray(frame_tar_seg).resize((ori_w, ori_h), 0))
+        frame_nm = frame_list[cnt].split("/")[-1].replace(".jpg", ".png")
+        imwrite_indexed(os.path.join(video_folder, frame_nm), frame_tar_seg, color_palette)
+
+
+def restrict_neighborhood(cfg, h, w):
+    """
+    Ultra-fast approach using meshgrid and advanced indexing
+    """
+    size_mask = cfg.eval.size_mask_neighborhood
+
+    # Create coordinate meshgrids
+    query_i, query_j = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij")
+    source_i, source_j = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij")
+
+    # Flatten coordinates
+    query_i_flat = query_i.flatten().unsqueeze(1)  # (h*w, 1)
+    query_j_flat = query_j.flatten().unsqueeze(1)  # (h*w, 1)
+    source_i_flat = source_i.flatten().unsqueeze(0)  # (1, h*w)
+    source_j_flat = source_j.flatten().unsqueeze(0)  # (1, h*w)
+
+    # Calculate distances using broadcasting
+    i_dist = torch.abs(query_i_flat - source_i_flat)  # (h*w, h*w)
+    j_dist = torch.abs(query_j_flat - source_j_flat)  # (h*w, h*w)
+
+    # Create mask
+    mask = ((i_dist <= size_mask) & (j_dist <= size_mask)).float()
+
+    return mask.cuda(non_blocking=True)
+
+
+def norm_mask(mask):
+    c, h, w = mask.size()
+    for cnt in range(c):
+        mask_cnt = mask[cnt, :, :]
+        if mask_cnt.max() > 0:
+            mask_cnt = mask_cnt - mask_cnt.min()
+            mask_cnt = mask_cnt / mask_cnt.max()
+            mask[cnt, :, :] = mask_cnt
+    return mask
+
+
+def label_propagation(
+    cfg,
+    backbone,
+    upsampler,
+    transform,
+    frame_tar,
+    list_frame_feats,
+    list_segs,
+    mask_neighborhood,
+    mean_bck_tensor,
+    std_bck_tensor,
+    mean_ups_tensor,
+    std_ups_tensor,
+):
+    """
+    propagate segs of frames in list_frames to frame_tar
+    """
+    list_segs = [s.cuda() for s in list_segs]
+    segs = torch.cat(list_segs)
+    nmb_context, C, h, w = segs.shape
+
+    ## we only need to extract feature of the target frame
+    feat_tar = extract_feature(
+        cfg,
+        backbone,
+        upsampler,
+        transform,
+        frame_tar,
+        mean_bck_tensor,
+        std_bck_tensor,
+        mean_ups_tensor,
+        std_ups_tensor,
+    )
+    feat_tar = F.interpolate(feat_tar, size=(h, w), mode="bilinear", align_corners=False)
+    feat_tar = rearrange(feat_tar, "1 c h w -> (h w) c")
+    return_feat_tar = feat_tar.T  # dim x h*w
+
+    ncontext = len(list_frame_feats)
+    feat_sources = torch.stack(list_frame_feats)  # nmb_context x dim x h*w
+
+    feat_tar = F.normalize(feat_tar, dim=1, p=2)
+    feat_sources = F.normalize(feat_sources, dim=1, p=2)
+
+    feat_tar = feat_tar.unsqueeze(0).repeat(ncontext, 1, 1)
+    aff = torch.exp(torch.bmm(feat_tar, feat_sources) / 0.1)  # nmb_context x h*w (tar: query) x h*w (source: keys)
+
+    if cfg.eval.size_mask_neighborhood > 0:
+        if mask_neighborhood is None:
+            mask_neighborhood = restrict_neighborhood(cfg, h, w)
+            mask_neighborhood = mask_neighborhood.unsqueeze(0).repeat(ncontext, 1, 1)
+        aff *= mask_neighborhood
+
+    aff = aff.transpose(2, 1).reshape(-1, h * w)  # nmb_context*h*w (source: keys) x h*w (tar: queries)
+    tk_val, _ = torch.topk(aff, dim=0, k=cfg.eval.topk)
+    tk_val_min, _ = torch.min(tk_val, dim=0)
+    aff[aff < tk_val_min] = 0
+
+    aff = aff / torch.sum(aff, keepdim=True, axis=0)
+
+    segs = segs.reshape(nmb_context, C, -1).transpose(2, 1).reshape(-1, C).T  # C x nmb_context*h*w
+    seg_tar = torch.mm(segs, aff)
+    seg_tar = seg_tar.reshape(1, C, h, w)
+    return seg_tar, return_feat_tar, mask_neighborhood
+
+
+@torch.no_grad()
+def extract_feature(
+    cfg,
+    backbone,
+    upsampler,
+    transform,
+    frame,
+    mean_bck_tensor,
+    std_bck_tensor,
+    mean_ups_tensor,
+    std_ups_tensor,
+):
+    """Extract one frame feature using JAFAR pipeline."""
+    # Transform frame and move to device
+    frame_tensor = transform(frame).unsqueeze(0).cuda()
+    frame_tensor = F.interpolate(
+        frame_tensor,
+        size=[v // backbone.config["ps"] * backbone.config["ps"] for v in frame_tensor.shape[-2:]],
+        mode="bilinear",
+        align_corners=False,
+    )
+
+    # Prepare image for backbone (normalize)
+    img_bck = (frame_tensor - mean_bck_tensor) / std_bck_tensor
+    lr_feats = backbone(img_bck)
+
+    # Upsample features
+    img_ups = (frame_tensor - mean_ups_tensor) / std_ups_tensor
+
+    # Resize input for upsampler to match target dimensions
+    hr_size = [v * cfg.eval.ups_factor for v in lr_feats.shape[-2:]]
+    img_ups = F.interpolate(img_ups, size=hr_size, mode="bicubic", align_corners=False)
+
+    hr_feats = upsampler(img_ups, lr_feats, hr_size)
+    return hr_feats
+
+
+def imwrite_indexed(filename, array, color_palette):
+    """Save indexed png for DAVIS."""
+    if np.atleast_3d(array).shape[2] != 1:
+        raise Exception("Saving indexed PNGs requires 2D array.")
+
+    im = Image.fromarray(array)
+    im.putpalette(color_palette.ravel())
+    im.save(filename, format="PNG")
+
+
+def to_one_hot(y_tensor, n_dims=None):
+    """
+    Take integer y (tensor or variable) with n dims &
+    convert it to 1-hot representation with n+1 dims.
+    """
+    if n_dims is None:
+        n_dims = int(y_tensor.max() + 1)
+    _, h, w = y_tensor.size()
+    y_tensor = y_tensor.type(torch.LongTensor).view(-1, 1)
+    n_dims = n_dims if n_dims is not None else int(torch.max(y_tensor)) + 1
+    y_one_hot = torch.zeros(y_tensor.size()[0], n_dims).scatter_(1, y_tensor, 1)
+    y_one_hot = y_one_hot.view(h, w, n_dims)
+    return y_one_hot.permute(2, 0, 1).unsqueeze(0)
+
+
+def read_frame_list(video_dir):
+    frame_list = [img for img in glob.glob(os.path.join(video_dir, "*.jpg"))]
+    frame_list = sorted(frame_list)
+    return frame_list
+
+
+def read_frame(frame_dir, ps=None):
+    img = Image.open(frame_dir)
+    ori_w, ori_h = img.size
+    return img, ori_h, ori_w
+
+
+def read_seg(seg_dir, ps, factor):
+    seg = Image.open(seg_dir)
+    _w, _h = seg.size  # note PIL.Image.Image's size is (w, h)
+    small_seg = np.array(seg.resize((_w // ps * factor, _h // ps * factor), 0))
+    small_seg = torch.from_numpy(small_seg.copy()).contiguous().float().unsqueeze(0)
+    return to_one_hot(small_seg), np.asarray(seg)
+
+
+def run_video_segmentation(cfg, backbone, upsampler, log_print):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    transform = transforms.Compose(
+        [
+            transforms.ToTensor(),
+        ]
+    )
+
+    color_palette = []
+    for line in urlopen("https://raw.githubusercontent.com/Liusifei/UVC/master/libs/data/palette.txt"):
+        color_palette.append([int(i) for i in line.decode("utf-8").split("\n")[0].split(" ")])
+    color_palette = np.asarray(color_palette, dtype=np.uint8).reshape(-1, 3)
+
+    video_list = open(os.path.join(cfg.dataset.dataroot, "ImageSets/2017/val.txt")).readlines()
+
+    start_time = time.time()
+
+    for i, video_name in enumerate(video_list):
+        video_name = video_name.strip()
+        log_print(f"[{i}/{len(video_list)}] Begin to segmentate video {video_name}.")
+
+        video_dir = os.path.join(cfg.dataset.dataroot, "JPEGImages/480p/", video_name)
+        frame_list = read_frame_list(video_dir)
+        seg_path = frame_list[0].replace("JPEGImages", "Annotations").replace("jpg", "png")
+
+        # Read segmentation with appropriate scaling
+        first_seg, seg_ori = read_seg(seg_path, backbone.config["ps"], cfg.eval.ups_factor)
+
+        eval_video_tracking_davis(
+            cfg, backbone, upsampler, transform, frame_list, video_dir, first_seg, seg_ori, color_palette, log_print
+        )
+
+    end_time = time.time()
+    log_print(f"[bold blue]Total time taken: {(end_time - start_time) / 60:.2f} minutes[/bold blue]")
+    log_print(f"[bold blue]Max GPU memory used: {torch.cuda.max_memory_allocated() / 1e9:.2f} GB[/bold blue]")
+
+    # Run comprehensive DAVIS evaluation
+    log_print("[bold cyan]Running comprehensive DAVIS eval...[/bold cyan]")
+    run_davis_evaluation(cfg, log_print)
+
+
+def run_davis_evaluation(cfg, log_print):
+    """Run comprehensive DAVIS evaluation using the integrated evaluation method"""
+    output_dir = HydraConfig.get().runtime.output_dir
+    exp_name = f"davis_vidseg_{cfg.eval.ups_factor}_{cfg.model.name}_{cfg.backbone.name}"
+    results_path = os.path.join(output_dir, exp_name)
+
+    if not os.path.exists(results_path):
+        log_print(f"[bold red]Results path does not exist: {results_path}[/bold red]")
+        return
+
+    # Check if results already exist
+    csv_name_global = "global_results-val.csv"
+    csv_name_per_sequence = "per-sequence_results-val.csv"
+    csv_name_global_path = os.path.join(results_path, csv_name_global)
+    csv_name_per_sequence_path = os.path.join(results_path, csv_name_per_sequence)
+
+    if os.path.exists(csv_name_global_path) and os.path.exists(csv_name_per_sequence_path):
+        log_print("[bold yellow]Using precomputed DAVIS results...[/bold yellow]")
+        table_g = pd.read_csv(csv_name_global_path)
+        table_seq = pd.read_csv(csv_name_per_sequence_path)
+    else:
+        log_print("[bold cyan]Computing DAVIS evaluation metrics...[/bold cyan]")
+
+        # Create dataset and evaluate
+        dataset_eval = DAVISEvaluation(davis_root=cfg.dataset.dataroot, task="semi-supervised", gt_set="val")
+        metrics_res = dataset_eval.evaluate(results_path)
+
+        if not metrics_res or "J" not in metrics_res or "F" not in metrics_res:
+            log_print("[bold red]DAVIS evaluation failed![/bold red]")
+            return
+
+        J, F = metrics_res["J"], metrics_res["F"]
+
+        # Generate dataframe for the general results
+        g_measures = ["J&F-Mean", "J-Mean", "J-Recall", "J-Decay", "F-Mean", "F-Recall", "F-Decay"]
+        final_mean = (np.mean(J["M"]) + np.mean(F["M"])) / 2.0
+        g_res = np.array(
+            [final_mean, np.mean(J["M"]), np.mean(J["R"]), np.mean(J["D"]), np.mean(F["M"]), np.mean(F["R"]), np.mean(F["D"])]
+        )
+        g_res = np.reshape(g_res, [1, len(g_res)])
+        table_g = pd.DataFrame(data=g_res, columns=g_measures)
+
+        with open(csv_name_global_path, "w") as f:
+            table_g.to_csv(f, index=False, float_format="%.3f")
+        log_print(f"[bold green]Global results saved in {csv_name_global_path}[/bold green]")
+
+        # Generate a dataframe for the per sequence results
+        seq_names = list(J["M_per_object"].keys())
+        seq_measures = ["Sequence", "J-Mean", "F-Mean"]
+        J_per_object = [J["M_per_object"][x] for x in seq_names]
+        F_per_object = [F["M_per_object"][x] for x in seq_names]
+        table_seq = pd.DataFrame(data=list(zip(seq_names, J_per_object, F_per_object)), columns=seq_measures)
+
+        with open(csv_name_per_sequence_path, "w") as f:
+            table_seq.to_csv(f, index=False, float_format="%.3f")
+        log_print(f"[bold green]Per-sequence results saved in {csv_name_per_sequence_path}[/bold green]")
+
+    # Print the results
+    log_print("[bold blue]--------------------------- Global DAVIS Results ---------------------------[/bold blue]")
+    log_print(table_g.to_string(index=False))
+    log_print("[bold blue]---------- Per-sequence DAVIS Results ----------[/bold blue]")
+    log_print(table_seq.to_string(index=False))
+
+    # Also save as JSON for compatibility
+    metrics_dict = table_g.iloc[0].to_dict()
+    import json
+
+    with open(os.path.join(output_dir, "davis_metrics.json"), "w") as f:
+        json.dump(metrics_dict, f, indent=2)
+    log_print(f"[bold cyan]DAVIS metrics saved to: {os.path.join(output_dir, 'davis_metrics.json')}[/bold cyan]")
+
+
+@hydra.main(config_path="../config", config_name="eval_video_seg", version_base=None)
+def main(cfg: DictConfig):
+    # Setup dual console logging (terminal and file)
+    terminal_console = Console()
+    current_run_dir = HydraConfig.get().runtime.output_dir
+    log_file_path = os.path.join(current_run_dir, "eval_davis.log")
+    file_console = Console(file=open(log_file_path, "w"))
+
+    def log_print(*args, **kwargs):
+        """Log to both terminal and file with immediate flushing"""
+        terminal_console.print(*args, **kwargs)
+        file_console.print(*args, **kwargs)
+        file_console.file.flush()
+
+    # Start logging
+    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    log_print(f"\n[bold blue]{'='*50}[/bold blue]")
+    log_print(f"[bold blue]Starting DAVIS Video Segmentation Evaluation at {timestamp}[/bold blue]")
+    log_print(f"[bold green]Configuration:[/bold green]")
+    log_print(OmegaConf.to_yaml(cfg))
+
+    device = "cuda"
+    log_print(f"[bold yellow]Using device: {device}[/bold yellow]")
+
+    # Load backbone and upsampler using Hydra configuration
+    backbone_configs = [{"name": cfg.backbone.name}]
+    backbones = load_multiple_backbones(cfg, backbone_configs, device="cuda")[0]
+    backbone = backbones[0].eval().cuda()
+    upsampler = instantiate(cfg.model).eval().cuda()
+
+    # Load model checkpoint if provided
+    if cfg.eval.model_ckpt:
+        checkpoint = torch.load(cfg.eval.model_ckpt, map_location=device, weights_only=False)
+        if cfg.model.name == "featup":
+            new_ckpts = {
+                k.replace("model.1.", "norm."): v
+                for k, v in checkpoint["state_dict"].items()
+                if "upsampler" in k or "model.1.norm" in k
+            }
+            upsampler.load_state_dict(new_ckpts, strict=True)
+        else:
+            upsampler.load_state_dict(checkpoint, strict=True)
+        log_print(f"[green]Loaded model from checkpoint: {cfg.eval.model_ckpt}[/green]")
+    else:
+        log_print(f"[yellow]No model checkpoint provided, using untrained model[/yellow]")
+
+    run_video_segmentation(cfg, backbone, upsampler, log_print)
+
+    # Log completion before closing file
+    log_print(f"[bold blue]Evaluation completed. Log saved to: {log_file_path}[/bold blue]")
+
+    # Close file console
+    file_console.file.close()
+
+
+if __name__ == "__main__":
+    main()

Pixal3D/torch_hub/hub/valeoai_NAF_main/hydra_plugins/resolvers.py

@@ -0,0 +1,41 @@
+import re
+
+from omegaconf import OmegaConf
+
+
+def get_feature(target: str) -> int:
+    """Resolve feature dimensions from backbone name patterns"""
+    model_name = target.lower()
+
+    if "vits" in model_name or "small" in model_name:
+        return 384
+    elif "vitb" in model_name or "base" in model_name or model_name == "radio_v2.5-b":
+        return 768
+    elif "vitl" in model_name or "large" in model_name or model_name == "radio_v2.5-l":
+        return 1024
+    elif "tiny" in model_name:
+        return 192
+    else:
+        print(f"Warning: get_feature() - Unsupported backbone: {model_name}. Returning default 0.")
+        return 0
+
+
+def get_patch_size(target: str) -> int:
+    """Resolve patch size from backbone name patterns using regex"""
+    model_name = target.lower()
+
+    # Return a default value for consistency
+    if "franca" in model_name:
+        return 14
+
+    # Use regex to find patchXX pattern
+    match = re.search(r"patch(\d+)", model_name)
+    if match:
+        return int(match.group(1))
+
+    # Default to 16 if no patch size found
+    return 16
+
+
+OmegaConf.register_new_resolver("get_feature", lambda name: get_feature(name))
+OmegaConf.register_new_resolver("get_patch_size", lambda name: get_patch_size(name))

Pixal3D/torch_hub/hub/valeoai_NAF_main/notebooks/attention_maps.ipynb

@@ -0,0 +1,389 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "fa359c75",
+   "metadata": {},
+   "source": [
+    "# NAF: Neighborhood Attention Map\n",
+    "\n",
+    "This notebook provides tools to inspect and visualize attention weights from the NAF model.\n",
+    "\n",
+    "**Features:**\n",
+    "- Load pre-trained NAF model and test images\n",
+    "- Select query points on the image\n",
+    "- Extract and visualize attention weights for the selected position\n",
+    "- Compare attention patterns across different image regions"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "13d519e9",
+   "metadata": {},
+   "source": [
+    "## Setup and Imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fdfa00f8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Use inline matplotlib for better compatibility\n",
+    "%matplotlib widget"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c5aaceb7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "sys.path.append('..')\n",
+    "\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "import matplotlib.patches as patches\n",
+    "from PIL import Image\n",
+    "import torchvision.transforms as T\n",
+    "\n",
+    "from hydra import compose, initialize\n",
+    "from hydra.core.global_hydra import GlobalHydra\n",
+    "from hydra.utils import instantiate\n",
+    "from IPython.display import clear_output\n",
+    "\n",
+    "from src.model.naf import NAF\n",
+    "from utils.training import load_multiple_backbones\n",
+    "\n",
+    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
+    "print(f\"Using device: {device}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2417fe5b",
+   "metadata": {},
+   "source": [
+    "## 2. Load Model and Configuration"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5139fa03",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Initialize Hydra config\n",
+    "if not GlobalHydra.instance().is_initialized():\n",
+    "    initialize(config_path=\"../config\", version_base=None)\n",
+    "\n",
+    "cfg = compose(config_name=\"base\", overrides=[\"model=naf\", \"img_size=1024\"])\n",
+    "\n",
+    "# Initialize NAF model\n",
+    "model = torch.hub.load(\"valeoai/NAF\", \"naf\", pretrained=True, device=device)\n",
+    "model.eval()\n",
+    "\n",
+    "# Load backbone\n",
+    "backbone_configs = [{\"name\": \"vit_base_patch16_dinov3.lvd1689m\"}]\n",
+    "backbones, names, _ = load_multiple_backbones(cfg, backbone_configs, device)\n",
+    "backbone = backbones[0]\n",
+    "backbone.eval()\n",
+    "\n",
+    "print(f\"Model loaded: {type(model).__name__}\")\n",
+    "print(f\"Backbone loaded: {names[0]}\")\n",
+    "print(f\"Kernel size: {model.upsampler.kernel_size}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b7ae72e4",
+   "metadata": {},
+   "source": [
+    "## 3. Load and Preprocess Image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ae219175",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load image\n",
+    "image_path = \"../asset/dog0.jpg\"\n",
+    "pil_image = Image.open(image_path).convert(\"RGB\")\n",
+    "\n",
+    "# Resize to model input size\n",
+    "H, W = cfg.img_size, cfg.img_size\n",
+    "transform = T.Compose([T.Resize((H, W)), T.ToTensor()])\n",
+    "image = transform(pil_image).unsqueeze(0).to(device)\n",
+    "\n",
+    "# Normalize for backbone\n",
+    "mean_bck = backbone.config[\"mean\"]\n",
+    "std_bck = backbone.config[\"std\"]\n",
+    "normalize_bck = T.Normalize(mean=mean_bck, std=std_bck)\n",
+    "image_bck = normalize_bck(image)\n",
+    "\n",
+    "# Normalize for upsampler (ImageNet normalization)\n",
+    "normalize_ups = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
+    "image_ups = normalize_ups(image)\n",
+    "\n",
+    "# Convert to display format\n",
+    "image_vis = (image.clone().cpu().squeeze(0).permute(1, 2, 0).numpy() * 255).astype(np.uint8)\n",
+    "\n",
+    "print(f\"Loaded image from: {image_path}\")\n",
+    "print(f\"Image size: {image.shape}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0457092d",
+   "metadata": {},
+   "source": [
+    "## 4. Extract Backbone Features"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "784c862e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Extract features from backbone\n",
+    "with torch.no_grad():\n",
+    "    hr_feats = backbone(image_bck)\n",
+    "\n",
+    "# Define output size for encoder (match backbone output)\n",
+    "output_size = (hr_feats.shape[2], hr_feats.shape[3])\n",
+    "\n",
+    "# Compute scaling factors between image and feature space\n",
+    "scale_h = output_size[0] / H\n",
+    "scale_w = output_size[1] / W\n",
+    "\n",
+    "print(f\"Backbone features shape: {hr_feats.shape}\")\n",
+    "print(f\"Feature map size: {output_size}\")\n",
+    "print(f\"Scaling factors - h: {scale_h:.3f}, w: {scale_w:.3f}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d4a80c95",
+   "metadata": {},
+   "source": [
+    "## 5. Helper Functions for Query Position and Attention Extraction"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a62a3b4a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def extract_attention_at_position(query_img_h, query_img_w):\n",
+    "    \"\"\"\n",
+    "    Extract attention weights for a query position.\n",
+    "    \n",
+    "    Args:\n",
+    "        query_img_h: Query position height in image coordinates (0-447)\n",
+    "        query_img_w: Query position width in image coordinates (0-447)\n",
+    "        \n",
+    "    Returns:\n",
+    "        attn_map: 2D attention map of shape (kernel_h, kernel_w)\n",
+    "        pos_feat: Query position in feature space\n",
+    "    \"\"\"\n",
+    "    # Convert image coordinates to feature space\n",
+    "    pos_feat = (int(query_img_h * scale_h), int(query_img_w * scale_w))\n",
+    "    \n",
+    "    # Extract attention weights\n",
+    "    with torch.no_grad():\n",
+    "        # Get RoPE-transformed features (queries)\n",
+    "        _, attn_weights = model(image_ups, hr_feats, output_size=output_size, return_weights=True)\n",
+    "        \n",
+    "    # Extract attention for the query position\n",
+    "    # attn_weights: [B, num_heads, H, W, kernel_h * kernel_w]\n",
+    "    query_h, query_w = pos_feat\n",
+    "    attn_at_query = attn_weights[0, :, query_h, query_w, :]  # [num_heads, kernel_h*kernel_w]\n",
+    "    \n",
+    "    # Average across heads\n",
+    "    attn_mean = attn_at_query.mean(dim=0)  # [kernel_h*kernel_w]\n",
+    "    \n",
+    "    # Reshape to 2D attention map\n",
+    "    kernel_h, kernel_w = model.upsampler.kernel_size\n",
+    "    attn_map = attn_mean.view(kernel_h, kernel_w).cpu().numpy()\n",
+    "    \n",
+    "    return attn_map, pos_feat\n",
+    "\n",
+    "\n",
+    "print(\"Helper functions loaded\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f289e60c",
+   "metadata": {},
+   "source": [
+    "## 6. Visualize Attention for a Single Query Point"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ccd18126",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Interactive attention visualization\n",
+    "fig, axes = plt.subplots(1, 3, figsize=(12*2/3, 4*2/3))\n",
+    "\n",
+    "# Initialize with default position\n",
+    "query_img_h, query_img_w = 650, 270\n",
+    "\n",
+    "# Plot elements that will be updated\n",
+    "img_display = axes[0].imshow(image_vis)\n",
+    "query_point, = axes[0].plot([], [], 'r*', markersize=15, markeredgecolor='white', markeredgewidth=1)\n",
+    "rect_patch = patches.Rectangle((0, 0), 0, 0, linewidth=2, edgecolor='red', facecolor='none', linestyle='--')\n",
+    "axes[0].add_patch(rect_patch)\n",
+    "\n",
+    "axes[0].set_xlabel('Image Width')\n",
+    "axes[0].set_ylabel('Image Height')\n",
+    "axes[0].set_title('Click on image to select query point')\n",
+    "axes[0].set_aspect('equal')\n",
+    "\n",
+    "# Attention heatmap (will be updated)\n",
+    "attn_display = axes[1].imshow(np.zeros((5, 5)), cmap='RdBu_r', aspect='equal')\n",
+    "axes[1].set_xlabel('Kernel Width')\n",
+    "axes[1].set_ylabel('Kernel Height')\n",
+    "axes[1].set_title('Attention Map')\n",
+    "\n",
+    "cbar1 = plt.colorbar(attn_display, ax=axes[1], fraction=0.046, pad=0.04)\n",
+    "cbar1.set_label('Attention Weight')\n",
+    "\n",
+    "# Overlay visualization\n",
+    "overlay_display = axes[2].imshow(np.zeros((10, 10, 3), dtype=np.uint8))\n",
+    "axes[2].set_xlabel('Width')\n",
+    "axes[2].set_ylabel('Height')\n",
+    "axes[2].set_title('Attention Overlay on Neighborhood')\n",
+    "axes[2].set_aspect('equal')\n",
+    "\n",
+    "# Statistics text\n",
+    "stats_text = axes[1].text(0.02, 0.98, '', transform=axes[1].transAxes, \n",
+    "                          verticalalignment='top', fontsize=8, \n",
+    "                          bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))\n",
+    "\n",
+    "def update_attention(query_h, query_w):\n",
+    "    \"\"\"Update the attention visualization for a new query position.\"\"\"\n",
+    "    # Extract attention map\n",
+    "    attn_map, pos_feat = extract_attention_at_position(query_h, query_w)\n",
+    "    kernel_h, kernel_w = model.upsampler.kernel_size\n",
+    "    \n",
+    "    # Update query point\n",
+    "    query_point.set_data([query_w], [query_h])\n",
+    "    \n",
+    "    # Update neighborhood rectangle\n",
+    "    half_h, half_w = kernel_h // 2, kernel_w // 2\n",
+    "    box_x0 = max(0, query_w - int(half_w / scale_w))\n",
+    "    box_y0 = max(0, query_h - int(half_h / scale_h))\n",
+    "    box_x1 = min(W, query_w + int(half_w / scale_w))\n",
+    "    box_y1 = min(H, query_h + int(half_h / scale_h))\n",
+    "    \n",
+    "    rect_patch.set_xy((box_x0, box_y0))\n",
+    "    rect_patch.set_width(box_x1 - box_x0)\n",
+    "    rect_patch.set_height(box_y1 - box_y0)\n",
+    "    \n",
+    "    # Update attention heatmap\n",
+    "    attn_display.set_data(attn_map)\n",
+    "    attn_display.set_clim(vmin=attn_map.min(), vmax=attn_map.max())\n",
+    "    axes[1].set_title(f'Attention Map ({kernel_h}×{kernel_w})')\n",
+    "    \n",
+    "    # Create overlay visualization\n",
+    "    # Extract neighborhood from image\n",
+    "    neighborhood = image_vis[box_y0:box_y1, box_x0:box_x1, :]\n",
+    "    \n",
+    "    # Resize attention map to match neighborhood size\n",
+    "    from scipy.ndimage import zoom\n",
+    "    attn_resized = zoom(attn_map, (neighborhood.shape[0] / attn_map.shape[0], \n",
+    "                                    neighborhood.shape[1] / attn_map.shape[1]), order=1)\n",
+    "    \n",
+    "    # Normalize attention to 0-1 for overlay\n",
+    "    attn_norm = (attn_resized - attn_resized.min()) / (attn_resized.max() - attn_resized.min() + 1e-8)\n",
+    "    \n",
+    "    # Create overlay: blend image with attention heatmap\n",
+    "    overlay = neighborhood.copy().astype(float)\n",
+    "    \n",
+    "    # Apply colormap to attention (red-blue)\n",
+    "    import matplotlib.cm as cm\n",
+    "    attn_colored = cm.RdBu_r(attn_norm)[:, :, :3] * 255  # RGB only\n",
+    "    \n",
+    "    # Blend: 50% image, 50% attention heatmap\n",
+    "    overlay_blended = (0.5 * overlay + 0.5 * attn_colored).astype(np.uint8)\n",
+    "    \n",
+    "    # Update overlay display\n",
+    "    overlay_display.set_data(overlay_blended)\n",
+    "    overlay_display.set_extent([box_x0, box_x1, box_y1, box_y0])\n",
+    "    axes[2].set_xlim(box_x0, box_x1)\n",
+    "    axes[2].set_ylim(box_y1, box_y0)\n",
+    "    axes[2].set_title('Attention Overlay on Neighborhood')\n",
+    "    \n",
+    "    # Update statistics text\n",
+    "    fig.canvas.draw_idle()\n",
+    "    \n",
+    "    # Print to console\n",
+    "    print(f\"Query: ({query_h}, {query_w}) | Feature: {pos_feat} | Center attn: {attn_map[kernel_h//2, kernel_w//2]:.4f}\")\n",
+    "\n",
+    "def onclick(event):\n",
+    "    \"\"\"Handle click events on the image.\"\"\"\n",
+    "    if event.inaxes == axes[0] and event.xdata is not None and event.ydata is not None:\n",
+    "        new_w = int(np.clip(event.xdata, 0, W-1))\n",
+    "        new_h = int(np.clip(event.ydata, 0, H-1))\n",
+    "        update_attention(new_h, new_w)\n",
+    "\n",
+    "fig.canvas.mpl_connect('button_press_event', onclick)\n",
+    "\n",
+    "plt.tight_layout()\n",
+    "\n",
+    "# Initial visualization\n",
+    "update_attention(query_img_h, query_img_w)\n",
+    "plt.show()\n",
+    "\n",
+    "print(\"Click on the left image to explore attention patterns!\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "04e96127",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "naf",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

Pixal3D/torch_hub/hub/valeoai_NAF_main/notebooks/inference.ipynb

@@ -0,0 +1,294 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# NAF: Zero-Shot Feature Upsampling via Neighborhood Attention Filtering.\n",
+    "\n",
+    "In this notebook we present how to upsample **Any Feature from Any VFM at Any Resolution**"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Setup"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import random\n",
+    "import sys\n",
+    "from pathlib import Path\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import torch\n",
+    "import PIL\n",
+    "import torch.nn.functional as F\n",
+    "import torchvision.transforms as T\n",
+    "from hydra import compose, initialize\n",
+    "from hydra.core.global_hydra import GlobalHydra\n",
+    "from hydra.utils import instantiate\n",
+    "from IPython.display import clear_output\n",
+    "\n",
+    "project_root = str(Path().absolute().parent)\n",
+    "sys.path.append(project_root)\n",
+    "\n",
+    "from utils.training import get_batch, get_dataloaders, load_multiple_backbones\n",
+    "from utils.visualization import plot_feats"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Load config"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if not GlobalHydra.instance().is_initialized():\n",
+    "    initialize(config_path=\"../config\", version_base=None)\n",
+    "\n",
+    "overrides = [\"val_dataloader.batch_size=1\", \n",
+    "             \"train_dataloader.batch_size=1\", \n",
+    "             \"model=naf\", \n",
+    "             \"img_size=448\"]\n",
+    "\n",
+    "cfg = compose(config_name=\"base\", overrides=overrides)\n",
+    "clear_output()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Dataloader"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_dataloader, val_dataloader = get_dataloaders(cfg, shuffle=True)\n",
+    "clear_output()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+    "model = torch.hub.load(\"valeoai/NAF\", \"naf\", pretrained=True, device=device)\n",
+    "model.cuda()\n",
+    "\n",
+    "clear_output()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Either iterate over the dataloader\n",
+    "batch = next(iter(val_dataloader))\n",
+    "batch = get_batch(batch, device)\n",
+    "img_batch = batch[\"image\"]\n",
+    "bs = img_batch.shape[0]\n",
+    "IMG_PATH = None\n",
+    "\n",
+    "# Or load a custom image\n",
+    "IMG_PATH = \"../asset/dinov3.png\"\n",
+    "img_batch = PIL.Image.open(IMG_PATH).convert(\"RGB\")\n",
+    "img_batch = val_dataloader.dataset.transform(img_batch).unsqueeze(0).to(device)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Backbone"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "@torch.no_grad()\n",
+    "def upsample_backbone(backbone, img, model, mean_std_bck, mean_std_ups, sizes=[512]):\n",
+    "    torch.cuda.empty_cache()\n",
+    "    model.eval()\n",
+    "    mean_bck, std_bck = mean_std_bck\n",
+    "    mean_ups, std_ups = mean_std_ups\n",
+    "\n",
+    "    img_bck = T.functional.normalize(img, mean=mean_bck, std=std_bck)\n",
+    "    img_ups = T.functional.normalize(img, mean=mean_ups, std=std_ups)\n",
+    "\n",
+    "    hr_feats = backbone(img_bck)\n",
+    "    hr_size = hr_feats.shape[-1]\n",
+    "    preds = []\n",
+    "    for size in sizes:\n",
+    "        pred = model(img_ups, hr_feats, (size, size))\n",
+    "        preds.append(pred)\n",
+    "\n",
+    "    hr_feats_ = torch.nn.functional.interpolate(hr_feats, size, mode=\"nearest-exact\")\n",
+    "    plot_feats(\n",
+    "        img_batch[0],\n",
+    "        hr_feats_[0],\n",
+    "        [p[0] for p in preds],\n",
+    "        legend=[f\"Input\", f\"{hr_size}x{hr_size}\"] + [f\"{size}x{size}\" for size in sizes],\n",
+    "        font_size=20,\n",
+    "    )\n",
+    "    plt.show()\n",
+    "    torch.cuda.empty_cache()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Any Backbone"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "backbone_configs = [\n",
+    "    {\"name\": \"vit_base_patch16_dinov3.lvd1689m\"},\n",
+    "    {\"name\": \"radio_v2.5-b\"},\n",
+    "    {\"name\": \"franca_vitb14\"},\n",
+    "    {\"name\": \"vit_base_patch14_reg4_dinov2\"},\n",
+    "    {\"name\": \"vit_base_patch14_dinov2.lvd142m\"},\n",
+    "]\n",
+    "mapping = {\n",
+    "    \"vit_base_patch16_dinov3.lvd1689m\": \"Dinov3-B\",\n",
+    "    \"radio_v2.5-b\": \"Radio-v2.5-B\",\n",
+    "    \"franca_vitb14\": \"Franca-B14\",\n",
+    "    \"vit_base_patch14_reg4_dinov2\": \"Dinov2-R-B\",\n",
+    "    \"vit_base_patch14_dinov2.lvd142m\": \"Dinov2-B\",\n",
+    "}\n",
+    "backbones, *_ = load_multiple_backbones(cfg, backbone_configs, device)\n",
+    "clear_output()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mean_ups, std_ups = (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)\n",
+    "for backbone in backbones:\n",
+    "    backbone.cuda().eval()\n",
+    "\n",
+    "    mean_bck, std_bck = backbone.config[\"mean\"], backbone.config[\"std\"]\n",
+    "\n",
+    "    print(f\"BACKBONE: {mapping[backbone_configs[backbones.index(backbone)]['name']].upper()}\")\n",
+    "    upsample_backbone(backbone, img_batch, model, (mean_bck, std_bck), (mean_ups, std_ups), sizes=[448])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Any Resolution"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "backbone_configs = [\n",
+    "    {\"name\": \"vit_base_patch16_dinov3.lvd1689m\"},\n",
+    "]\n",
+    "backbones, *_ = load_multiple_backbones(cfg, backbone_configs, device)\n",
+    "clear_output()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "backbone = backbones[0]\n",
+    "backbone.cuda().eval()\n",
+    "\n",
+    "mean_bck, std_bck = backbone.config[\"mean\"], backbone.config[\"std\"]\n",
+    "\n",
+    "print(f\"BACKBONE: {backbone_configs[backbones.index(backbone)]['name']}\")\n",
+    "upsample_backbone(backbone, img_batch, model, (mean_bck, std_bck), (mean_ups, std_ups), sizes=[64, 128, 256, 512, 1024])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "naf",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

Pixal3D/torch_hub/hub/valeoai_NAF_main/src/backbone/__init__.py

@@ -0,0 +1 @@
+from .vit_wrapper import PretrainedViTWrapper

Pixal3D/torch_hub/hub/valeoai_NAF_main/src/backbone/vit_wrapper.py

@@ -0,0 +1,180 @@
+import re
+import types
+from typing import List, Tuple, Union
+
+import timm
+import timm.data
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from timm.models.vision_transformer import VisionTransformer
+from torch import nn
+from torchvision import transforms
+
+# We provide a list of timm model names, more are available on their official repo.
+MODEL_LIST = [
+    # DINO
+    "vit_base_patch16_224.dino",
+    # DINOv2
+    "vit_base_patch14_dinov2.lvd142m",
+    # DINOv2-R
+    "vit_base_patch14_reg4_dinov2",
+    # Franca
+    "franca_vitb14",
+    # DINOv3-ViT
+    "vit_base_patch16_dinov3.lvd1689m",
+    "vit_large_patch16_dinov3.lvd1689m",
+    "vit_7b_patch16_dinov3.lvd1689m",
+    # SigLIP2
+    "vit_base_patch16_siglip_512.v2_webli",
+    # PE Core
+    "vit_pe_core_small_patch16_384.fb",
+    # PE Spatial
+    "vit_pe_spatial_tiny_patch16_512.fb",
+    # RADIO
+    "radio_v2.5-b",
+    # CAPI
+    "capi_vitl14_lvd",
+    # MAE
+    "vit_large_patch16_224.mae",
+]
+
+IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
+
+
+class PretrainedViTWrapper(nn.Module):
+
+    def __init__(
+        self,
+        name,
+        norm: bool = True,
+        dynamic_img_size: bool = True,
+        dynamic_img_pad: bool = False,
+        **kwargs,
+    ):
+        super().__init__()
+        # comment out the following line to test the models not in the list
+        self.name = name
+
+        load_weights = False
+        if "dvt_" == name[:4]:
+            load_weights = True
+            load_tag = "dvt"
+            name = name.replace("dvt_", "")
+        if "fit3d_" == name[:6]:
+            load_weights = True
+            load_tag = "fit3d"
+            name = name.replace("fit3d_", "")
+
+        # Set patch size
+        try:
+            self.patch_size = int(re.search(r"patch(\d+)", name).group(1))
+        except:
+            self.patch_size = 16
+        if "franca" in name or "capi" in name:
+            self.patch_size = 14
+        if "convnext" in name:
+            self.patch_size = 32
+
+        name, self.patch_size
+
+        self.dynamic_img_size = dynamic_img_size
+        self.dynamic_img_pad = dynamic_img_pad
+        self.model, self.config = self.create_model(name, **kwargs)
+        self.config["ps"] = self.patch_size
+        self.embed_dim = self.model.embed_dim
+        self.norm = norm
+
+        if load_weights:
+            ckpt = torch.load(f"/home/lchambon/workspace/JAFAR/ckpts/{load_tag}_{name}.pth", map_location="cpu")
+            if load_tag == "dvt":
+                self.load_state_dict(ckpt["model"], strict=True)
+            elif load_tag == "fit3d":
+                self.model.load_state_dict(ckpt, strict=True)
+
+    def create_model(self, name: str, **kwargs) -> Tuple[VisionTransformer, transforms.Compose]:
+        if "radio" in self.name:
+            model = torch.hub.load(
+                "NVlabs/RADIO",
+                "radio_model",
+                version=name,
+                progress=True,
+                skip_validation=True,
+            )
+            data_config = {
+                "mean": torch.tensor([0.0, 0.0, 0.0]),
+                "std": torch.tensor([1.0, 1.0, 1.0]),
+                "input_size": (3, 512, 512),
+            }
+
+        elif "franca" in self.name:
+            model = torch.hub.load("valeoai/Franca", name, use_rasa_head=True)
+            data_config = {"mean": IMAGENET_DEFAULT_MEAN, "std": IMAGENET_DEFAULT_STD, "input_size": (3, 448, 448)}
+
+        elif "capi" in self.name:
+            model = torch.hub.load("facebookresearch/capi:main", name, force_reload=False)
+
+            data_config = {"mean": IMAGENET_DEFAULT_MEAN, "std": IMAGENET_DEFAULT_STD, "input_size": (3, 448, 448)}
+
+        else:
+            timm_kwargs = dict(
+                pretrained=True,
+                num_classes=0,
+                patch_size=self.patch_size,
+            )
+
+            if "sam" not in self.name and "convnext" not in self.name:
+                timm_kwargs["dynamic_img_size"] = self.dynamic_img_size
+                timm_kwargs["dynamic_img_pad"] = self.dynamic_img_pad
+
+            timm_kwargs.update(kwargs)
+            model = timm.create_model(name, **timm_kwargs)
+            data_config = timm.data.resolve_model_data_config(model=model)
+
+        model = model.eval()
+
+        return model, data_config
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        n: Union[int, List[int], Tuple[int]] = 1,
+        return_prefix_tokens: bool = False,
+    ) -> Union[List[torch.Tensor], Tuple[torch.Tensor, List[torch.Tensor]]]:
+        """Intermediate layer accessor inspired by DINO / DINOv2 interface.
+        Args:
+            x: Input tensor.
+            n: Take last n blocks if int, all if None, select matching indices if sequence
+            reshape: Whether to reshape the output.
+        """
+
+        common_kwargs = dict(
+            norm=self.norm,
+            output_fmt="NCHW",
+            intermediates_only=True,
+        )
+
+        if "sam" not in self.name and return_prefix_tokens:
+            common_kwargs["return_prefix_tokens"] = return_prefix_tokens
+
+        elif "franca" in self.name:
+            B, C, H, W = x.shape
+            feats = self.model.forward_features(x, use_rasa_head=True)
+            out = feats["patch_token_rasa"]
+            out = rearrange(out, "b (h w) c -> b c h w", h=H // self.patch_size, w=W // self.patch_size)
+
+        elif "capi" in self.name:
+            *_, out = self.model(x)
+            out = out.permute(0, 3, 1, 2)
+
+        else:
+            out = self.model.forward_intermediates(x, n, **common_kwargs)
+
+        # "sam" models return feats only, others may return (feats, prefix)
+        if not isinstance(out, list) and not isinstance(out, tuple):
+            out = [out]
+            return out[0]
+        else:
+            assert len(out) == 1, f"Out contains {len(out)} elements, expected 1."
+            return out[0]

Pixal3D/torch_hub/hub/valeoai_NAF_main/src/layers/__init__.py

@@ -0,0 +1,3 @@
+from .attentions import CrossAttention
+from .convolutions import encoder
+from .rope import RoPE

Pixal3D/torch_hub/hub/valeoai_NAF_main/src/layers/attentions.py

@@ -0,0 +1,75 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+try:
+    NATTEN_RECENT = False
+    from natten.functional import na2d_av, na2d_qk
+except:
+    NATTEN_RECENT = True
+    from natten import na2d
+
+from src.layers.rope import RoPE
+
+
+def legacy_attention(q, k, v, kernel_size, dilation, scale=1, return_weights=False):
+    q = rearrange(q, "b h w n d -> b n h w d")
+    k = rearrange(k, "b h w n d -> b n h w d")
+    v = rearrange(v, "b h w n d -> b n h w d")
+    attn_scores = na2d_qk(q, k, kernel_size=kernel_size, dilation=dilation)
+    attn_scores = attn_scores * scale
+
+    attn_weights = attn_scores.softmax(dim=-1)
+    features = na2d_av(attn_weights, v, kernel_size=kernel_size, dilation=dilation)
+    features = rearrange(features, "b n h w d -> b h w n d")
+
+    if return_weights:
+        return features, attn_scores
+    return features
+
+
+class CrossAttention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        kernel_size=(9, 9),
+        **kwargs,
+    ):
+        super().__init__()
+        assert dim % num_heads == 0, "dim must be divisible by num_heads"
+
+        self.num_heads = num_heads
+        self.kernel_size = kernel_size
+
+        self.scale = (dim // num_heads) ** -0.5
+
+    def _resize(self, x, size, dtype):
+        x = F.interpolate(x, size=size, mode="nearest-exact")
+        x = rearrange(x, "b (n d) h w -> b h w n d", n=self.num_heads)
+        return x.to(dtype)
+
+    def forward(self, q, k, v, image=None, return_weights=False, **kwargs):
+        hq, wq = q.shape[-2:]
+        hk, wk = k.shape[-2:]
+        dilation = (hq // hk, wq // wk)
+        self.dilation = dilation
+
+        q = rearrange(q, "b (n d) h w -> b h w n d", n=self.num_heads)
+        k = self._resize(k, size=(hq, wq), dtype=q.dtype)
+        v = self._resize(v, size=(hq, wq), dtype=q.dtype)
+
+        # Use legacy attention pattern
+        if return_weights:
+            assert not NATTEN_RECENT, "Return weights not supported with recent natten versions"
+            out, attn_weights = legacy_attention(q, k, v, self.kernel_size, dilation, scale=self.scale, return_weights=True)
+            return rearrange(out, "b h w n d -> b (n d) h w"), attn_weights
+        else:
+            if NATTEN_RECENT:
+                # Use modern na2d attention
+                # Note: Modern na2d doesn't support position bias directly
+                out = na2d(q, k, v, kernel_size=self.kernel_size, dilation=dilation, stride=1, backend="cutlass-fna")
+            else:
+                out = legacy_attention(q, k, v, self.kernel_size, dilation, scale=self.scale)
+            return rearrange(out, "b h w n d -> b (n d) h w")

Pixal3D/torch_hub/hub/valeoai_NAF_main/src/layers/convolutions.py

@@ -0,0 +1,92 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# Convolutions
+class EncBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size=3,
+        norm_kwargs={},
+        pad_mode="zeros",
+        norm_fn=None,
+        activation_fn=nn.SiLU,
+        use_conv_shortcut=False,
+        bias=True,
+        residual=False,
+    ):
+        super().__init__()
+        self.use_conv_shortcut = use_conv_shortcut
+        self.norm1 = norm_fn(**norm_kwargs)
+        self.conv1 = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            padding=kernel_size // 2,
+            padding_mode=pad_mode,
+            bias=bias,
+        )
+        self.norm2 = norm_fn(**norm_kwargs)
+        self.conv2 = nn.Conv2d(
+            out_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            padding=kernel_size // 2,
+            padding_mode=pad_mode,
+            bias=bias,
+        )
+        self.activation_fn = activation_fn()
+        if in_channels != out_channels:
+            self.shortcut = nn.Conv2d(
+                in_channels,
+                out_channels,
+                kernel_size=1,
+                padding=0,
+                padding_mode=pad_mode,
+                bias=bias,
+            )
+        self.residual = residual
+
+    def forward(self, x):
+        residual = x
+        x = self.norm1(x)
+        x = self.activation_fn(x)
+        x = self.conv1(x)
+        x = self.norm2(x)
+        x = self.activation_fn(x)
+        x = self.conv2(x)
+        if self.use_conv_shortcut or residual.shape != x.shape:
+            residual = self.shortcut(residual)
+        if self.residual:
+            return x + residual
+        return x
+
+
+def encoder(in_dim, hidden_dim, kernel_size=1, ks_res=1, num_layers=2, bias=True, num_groups=8, residual=False):
+    return nn.Sequential(
+        nn.Conv2d(
+            in_dim,
+            hidden_dim,
+            kernel_size=kernel_size,
+            padding=kernel_size // 2,
+            padding_mode="reflect",
+            bias=bias,
+        ),
+        *[
+            EncBlock(
+                hidden_dim,
+                hidden_dim,
+                kernel_size=ks_res,
+                pad_mode="reflect",
+                norm_fn=nn.GroupNorm,
+                norm_kwargs={"num_groups": num_groups, "num_channels": hidden_dim},
+                activation_fn=nn.SiLU,
+                use_conv_shortcut=False,
+                bias=bias,
+                residual=residual,
+            )
+            for _ in range(num_layers)
+        ],
+    )

Pixal3D/torch_hub/hub/valeoai_NAF_main/src/layers/rope.py

@@ -0,0 +1,174 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import math
+from typing import Literal, Optional
+
+import numpy as np
+import torch
+from einops import rearrange
+from torch import Tensor, nn
+
+
+def rope_rotate_half(x: Tensor) -> Tensor:
+    # x:   [ x0  x1  x2  x3  x4  x5]
+    # out: [-x3 -x4 -x5  x0  x1  x2]
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def rope_apply(x: Tensor, sin: Tensor, cos: Tensor) -> Tensor:
+    """
+    - x: feature vector of shape [..., D] (e.g., [x0, x1, ..., x_{D-1}])
+    - sin, cos: shape [..., D], e.g., [sin0, sin1, ..., sin_{D/2-1}, sin0, ..., sin_{D/2-1}]
+    - rope_apply rotates the embedding by pairing the first half with the second half:
+    - out = [ x0*cos0 - x_{D/2}*sin0,
+            x1*cos1 - x_{D/2+1}*sin1,
+            ...,
+            x_{D/2}*cos0 + x0*sin0,
+            x_{D/2+1}*cos1 + x1*sin1,
+            ... ]
+    """
+    return (x * cos) + (rope_rotate_half(x) * sin)
+
+
+# RoPE positional embedding with no mixing of coordinates (axial) and no learnable weights
+# Supports two parametrizations of the rope parameters: either using `base` or `min_period` and `max_period`.
+class RoPE(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        *,
+        num_heads: int,
+        base: float = 100.0,
+        min_period: float = None,
+        max_period: float = None,
+        normalize_coords: Literal["min", "max", "separate"] = "separate",
+        shift_coords: float = None,
+        jitter_coords: float = None,
+        rescale_coords: float = None,
+        dtype: torch.dtype = None,
+        device: torch.device = None,
+    ):
+        super().__init__()
+        assert embed_dim % (4 * num_heads) == 0
+        both_periods = min_period is not None and max_period is not None
+        if (base is None and not both_periods) or (base is not None and both_periods):
+            raise ValueError("Either `base` or `min_period`+`max_period` must be provided.")
+
+        D_head = embed_dim // num_heads
+        self.num_heads = num_heads
+        self.base = base
+        self.min_period = min_period
+        self.max_period = max_period
+        self.D_head = D_head
+        self.normalize_coords = normalize_coords
+        self.shift_coords = shift_coords
+        self.jitter_coords = jitter_coords
+        self.rescale_coords = rescale_coords
+
+        self._cached_coords = None
+        self._cached_hw = None
+
+        # Needs persistent=True because we do teacher.load_state_dict(student.state_dict()) to initialize the teacher
+        self.dtype = dtype  # Don't rely on self.periods.dtype
+        self.register_buffer(
+            "periods",
+            torch.empty(D_head // 4, device=device, dtype=dtype),
+            persistent=True,
+        )
+        self._init_weights()
+
+    def create_coordinate(self, *, H: int, W: int) -> tuple[Tensor, Tensor]:
+        device = self.periods.device
+        dtype = self.dtype
+        dd = {"device": device, "dtype": dtype}
+
+        # Prepare coords in range [-1, +1]
+        if self.normalize_coords == "max":
+            max_HW = max(H, W)
+            coords_h = torch.arange(0.5, H, **dd) / max_HW  # [H]
+            coords_w = torch.arange(0.5, W, **dd) / max_HW  # [W]
+        elif self.normalize_coords == "min":
+            min_HW = min(H, W)
+            coords_h = torch.arange(0.5, H, **dd) / min_HW  # [H]
+            coords_w = torch.arange(0.5, W, **dd) / min_HW  # [W]
+        elif self.normalize_coords == "separate":
+            coords_h = torch.arange(0.5, H, **dd) / H  # [H]
+            coords_w = torch.arange(0.5, W, **dd) / W  # [W]
+        else:
+            raise ValueError(f"Unknown normalize_coords: {self.normalize_coords}")
+        coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij"), dim=-1)  # [H, W, 2]
+        coords = coords.flatten(0, 1)  # [HW, 2]
+        coords = 2.0 * coords - 1.0  # Shift range [0, 1] to [-1, +1]
+
+        # Shift coords by adding a uniform value in [-shift, shift]
+        if self.training and self.shift_coords is not None:
+            shift_hw = torch.empty(2, **dd).uniform_(-self.shift_coords, self.shift_coords)
+            coords += shift_hw[None, :]
+
+        # Jitter coords by multiplying the range [-1, 1] by a log-uniform value in [1/jitter, jitter]
+        if self.training and self.jitter_coords is not None:
+            jitter_max = np.log(self.jitter_coords)
+            jitter_min = -jitter_max
+            jitter_hw = torch.empty(2, **dd).uniform_(jitter_min, jitter_max).exp()
+            coords *= jitter_hw[None, :]
+
+        # Rescale coords by multiplying the range [-1, 1] by a log-uniform value in [1/rescale, rescale]
+        if self.training and self.rescale_coords is not None:
+            rescale_max = np.log(self.rescale_coords)
+            rescale_min = -rescale_max
+            rescale_hw = torch.empty(1, **dd).uniform_(rescale_min, rescale_max).exp()
+            coords *= rescale_hw
+
+        return coords
+
+    def _init_weights(self):
+        device = self.periods.device
+        dtype = self.dtype
+        if self.base is not None:
+            periods = self.base ** (2 * torch.arange(self.D_head // 4, device=device, dtype=dtype) / (self.D_head // 2))
+        else:
+            periods = torch.logspace(math.log10(self.min_period), math.log10(self.max_period), steps=self.D_head // 4)
+        self.periods.data = periods
+
+    def rotate(self, x, coords):
+        # Prepare angles and sin/cos
+        angles = 2 * math.pi * coords[:, :, None] / self.periods[None, None, :]  # [HW, 2, D//4]
+        # Flatten u and v frequencies into a single vector per spatial location: [u1, ..., u_{D//4}, v1, ..., v_{D//4}]
+        angles = angles.flatten(1, 2)  # [HW, D//2]
+        # Repeat the vector to match embedding dimension D: [u1,...,u_{D//4}, v1,...,v_{D//4}, u1,...,u_{D//4}, v1,...,v_{D//4}]
+        angles = angles.tile(2)  # [HW, D]
+
+        cos = torch.cos(angles)  # [HW, D]
+        sin = torch.sin(angles)  # [HW, D]
+
+        # x_u0 ... x_u{D//2-1} x_v0 ... x_v{D//2-1}
+
+        # Apply RoPE: rotates each (u_i, v_i) pair
+        # first half:  x[i]*cos[i] - x[i+D/2]*sin[i]
+        # second half: x[i+D/2]*cos[i] + x[i]*sin[i]
+        return rope_apply(x, sin, cos)
+
+    def forward(self, x, layout: str = "spatial"):
+        h, w = x.shape[-2:]
+        x = rearrange(x, "b (n d) h w -> b n (h w) d", n=self.num_heads)
+
+        if (h, w) != self._cached_hw:
+            self._cached_coords = self.create_coordinate(H=h, W=w)
+            self._cached_hw = (h, w)
+
+        coords = self._cached_coords
+
+        # Rotate
+        x = self.rotate(x, coords)
+
+        # Reshape
+        if layout == "spatial":
+            x = rearrange(x, f"b n (h w) d -> b (n d) h w", h=h, w=w)
+        elif layout == "flatten":
+            x = rearrange(x, f"b n (h w) d -> b h w n d", h=h, w=w)
+
+        return x

Pixal3D/torch_hub/hub/valeoai_NAF_main/src/loss.py

@@ -0,0 +1,45 @@
+import torch
+import torch.nn as nn
+from torch import nn
+from torch.nn import functional as F
+
+
+class MSELoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.mse_loss = torch.nn.MSELoss()
+
+    def forward(self, pred, target, normalize=False):
+        if normalize:
+            # If you must normalize (example: min-max scaling)
+            min_val = torch.min(target, dim=1, keepdim=True).values
+            max_val = torch.max(target, dim=1, keepdim=True).values
+            pred_normalized = (pred - min_val) / (max_val - min_val + 1e-6)
+            target_normalized = (target - min_val) / (max_val - min_val + 1e-6)
+        else:
+            pred_normalized = pred
+            target_normalized = target
+
+        return self.mse_loss(pred_normalized, target_normalized)
+
+
+class Loss(nn.Module):
+
+    def __init__(
+        self,
+        loss_type,
+        dim=384,
+    ):
+        super().__init__()
+        self.dim = dim
+
+        if loss_type == "mse":
+            loss = MSELoss()
+        else:
+            raise NotImplementedError(f"Loss type {loss_type} not implemented")
+
+        self.loss_func = loss
+
+    def __call__(self, pred, target, *args, **kwargs):
+        loss = self.loss_func(pred, target, *args, **kwargs)
+        return {"total": loss}

Pixal3D/torch_hub/hub/valeoai_NAF_main/src/model/__init__.py

@@ -0,0 +1,12 @@
+from .anyup import AnyUpsampler
+from .base import BaseUpsampler
+from .bilinear import Bilinear
+from .featup import FeatUp
+from .ircnn import IRCNN
+from .jafar import JAFAR
+from .jbf import JBF
+from .jbu import JBU
+from .naf import NAF
+from .nearest import Nearest
+from .rednet import REDNet
+from .restormer import Restormer