Add missing pa_desc/data/ files; anchor data/ ignore to root

Initial commit excluded ninjadesc/pa_desc/data/ because the .gitignore
rule "data/" matched any directory named data. Anchored the ignore
to the repo root (/data/, /outputs/) and added the package files.

Author: hyojinie

.gitignore

@@ -11,5 +11,5 @@ venv/
 .vscode/
 .DS_Store
 
-outputs/
-data/
+/outputs/
+/data/

ninjadesc/pa_desc/data/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.

ninjadesc/pa_desc/data/demo.py

@@ -0,0 +1,41 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+import torch
+from torch.utils.data import Dataset
+
+from ninjadesc.lemuria.recon.prepare import read_h5
+
+
+class DemoDataset(Dataset):
+    def __init__(
+        self,
+        root_path: str,
+        # img_dir: str = "lemuria/test_images",
+        h5_dir: str = "pa_desc/h5_test_images",
+    ):
+        super().__init__()
+
+        # self.img_dir = os.path.join(root_path, img_dir)
+        self.h5_dir = os.path.join(root_path, h5_dir)
+
+        # imgs = [os.path.join(self.img_dir, img) for img in os.listdir(self.img_dir)]
+        self.h5s = [os.path.join(self.h5_dir, h5) for h5 in os.listdir(self.h5_dir)]
+
+    def __len__(self):
+        return len(self.h5s)
+
+    def __getitem__(self, idx):
+        feats, rgbs = read_h5(
+            self.h5s[idx],
+            descriptor_type="SOS",
+            max_keypoints=1000,
+            flip=False,
+        )
+
+        return {"feats": torch.Tensor(feats), "rgbs": torch.Tensor(rgbs)}

ninjadesc/pa_desc/data/hpatches.py

@@ -0,0 +1,52 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+from torch.utils.data import Dataset
+
+
+class HPatches(Dataset):
+    """HPatches patch-matching benchmark.
+
+    NOTE: This is a placeholder. The original NinjaDesc paper uses the HPatches
+    benchmark from Balntas et al. (CVPR 2017). To run evaluation, populate this
+    class with a loader for the official HPatches release
+    (https://github.com/hpatches/hpatches-dataset). The expected item format
+    matches `PhotoTour`:
+
+        {"patches": Tensor(2, 1, H, W), "labels": Tensor(1)}
+    """
+
+    def __init__(
+        self,
+        split: str = "a",
+        base_path: str = None,
+        in_memory: bool = False,
+        nb_patches_per_track: int = 2,
+        train: bool = False,
+        transform=None,
+    ):
+        super().__init__()
+        if base_path is None:
+            base_path = os.path.join(
+                os.environ.get("NINJADESC_DATA_ROOT", "./data"),
+                f"HPatches/hpatches_32x32_{split}",
+            )
+        self.base_path = base_path
+        self.split = split
+        self.nb_patches_per_track = nb_patches_per_track
+        self.transform = transform
+        self.name = f"hpatches_{split}"
+
+    def __len__(self) -> int:
+        raise NotImplementedError(
+            "HPatches loader not implemented. Populate ninjadesc/pa_desc/data/hpatches.py "
+            "with the official HPatches benchmark loader."
+        )
+
+    def __getitem__(self, idx):
+        raise NotImplementedError

ninjadesc/pa_desc/data/megadepth.py

@@ -0,0 +1,54 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+import torch
+from torch.utils.data import Dataset
+
+from ninjadesc.lemuria.recon.prepare import read_h5
+
+
+def _default_root() -> str:
+    return os.environ.get("NINJADESC_DATA_ROOT", "./data")
+
+
+class MegaDepthDataset(Dataset):
+    def __init__(
+        self,
+        root_path: str = None,
+        h5_dir: str = "megadepth_h5s_sos_original",
+        splits_dir: str = "megadepth_splits",
+        splits_suffix: str = "_sos_original",
+        mode: str = "train",
+        kpt_type: str = "SOS",
+        num_samples: int = 50000,
+    ):
+        super().__init__()
+
+        if root_path is None:
+            root_path = _default_root()
+
+        h5s_txt_path = os.path.join(root_path, splits_dir, f"{mode}{splits_suffix}.txt")
+        with open(h5s_txt_path, "r") as f:
+            h5s = f.read().splitlines()
+
+        self.h5_dir = os.path.join(root_path, h5_dir)
+        self.h5s = [os.path.join(self.h5_dir, h5) for h5 in h5s]
+        self.h5s = self.h5s[:num_samples]
+        self.kpt_type = kpt_type
+
+    def __len__(self):
+        return len(self.h5s)
+
+    def __getitem__(self, idx):
+        feats, rgbs = read_h5(
+            self.h5s[idx],
+            descriptor_type=self.kpt_type,
+            max_keypoints=1000,
+            flip=False,
+        )
+        return torch.Tensor(feats), torch.Tensor(rgbs)

ninjadesc/pa_desc/data/phototour.py

@@ -0,0 +1,62 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+import torch
+from torch.utils.data import Dataset
+from torchvision.datasets import PhotoTour as TVPhotoTour
+
+
+class PhotoTour(Dataset):
+    """UBC PhotoTour patch dataset (Liberty / Notredame / Yosemite).
+
+    Thin adapter over torchvision.datasets.PhotoTour that returns paired
+    patches and binary match labels in the format expected by NinjaDesc
+    descriptor training:
+
+        item = {"patches": Tensor(2, 1, 32, 32), "labels": Tensor(1)}
+    """
+
+    def __init__(
+        self,
+        name: str = "liberty",
+        data_root: str = None,
+        nb_patches_per_track: int = 2,
+        train: bool = True,
+        transform=None,
+        download: bool = True,
+    ):
+        super().__init__()
+        if data_root is None:
+            data_root = os.path.join(
+                os.environ.get("NINJADESC_DATA_ROOT", "./data"), "PhotoTour"
+            )
+        os.makedirs(data_root, exist_ok=True)
+        # torchvision returns matched/unmatched triplet indices via the train arg.
+        self._tv = TVPhotoTour(
+            root=data_root, name=name, train=train, transform=transform, download=download
+        )
+        self._train = train
+        self.name = name
+        self.nb_patches_per_track = nb_patches_per_track
+
+    def __len__(self) -> int:
+        return len(self._tv)
+
+    def __getitem__(self, idx):
+        sample = self._tv[idx]
+        if self._train:
+            # torchvision train mode returns (anchor, positive, negative)
+            anchor, positive, _ = sample
+            patches = torch.stack([anchor.float(), positive.float()], dim=0)
+            label = torch.tensor(1, dtype=torch.long)
+        else:
+            # eval mode returns (patch_a, patch_b, match_label)
+            patch_a, patch_b, label = sample
+            patches = torch.stack([patch_a.float(), patch_b.float()], dim=0)
+            label = torch.as_tensor(label, dtype=torch.long)
+        return {"patches": patches.unsqueeze(1), "labels": label}

ninjadesc/pa_desc/data/transforms.py

@@ -0,0 +1,88 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import numpy as np
+import torch
+import torchvision.transforms.functional as TF
+
+
+class Grayscale:
+    def __call__(self, img):
+        if isinstance(img, torch.Tensor):
+            if img.ndim == 3 and img.shape[0] == 3:
+                return TF.rgb_to_grayscale(img, num_output_channels=1)
+            return img
+        if isinstance(img, np.ndarray) and img.ndim == 3 and img.shape[-1] == 3:
+            return img.mean(axis=-1, keepdims=True)
+        return img
+
+
+class Resize:
+    def __init__(self, size: int = 32):
+        self.size = size
+
+    def __call__(self, img):
+        if isinstance(img, torch.Tensor):
+            return TF.resize(img, [self.size, self.size], antialias=True)
+        if isinstance(img, np.ndarray):
+            tensor = torch.as_tensor(img).permute(2, 0, 1) if img.ndim == 3 else torch.as_tensor(img)[None]
+            tensor = TF.resize(tensor, [self.size, self.size], antialias=True)
+            return tensor.squeeze(0).numpy() if img.ndim == 2 else tensor.permute(1, 2, 0).numpy()
+        return img
+
+
+class ToFloat:
+    def __init__(self, normalise: bool = False):
+        self.normalise = normalise
+
+    def __call__(self, img):
+        if isinstance(img, torch.Tensor):
+            img = img.float()
+        else:
+            img = np.asarray(img, dtype=np.float32)
+        if self.normalise:
+            img = img / 255.0
+        return img
+
+
+class ToTensor:
+    def __call__(self, img):
+        if isinstance(img, torch.Tensor):
+            return img
+        arr = np.asarray(img)
+        if arr.ndim == 2:
+            arr = arr[None, ...]
+        elif arr.ndim == 3:
+            arr = arr.transpose(2, 0, 1)
+        return torch.as_tensor(arr).float()
+
+
+class RandomFlipUDSet:
+    def __init__(self, p: float = 0.5):
+        self.p = p
+
+    def __call__(self, sample):
+        if torch.rand(1).item() >= self.p:
+            return sample
+        if isinstance(sample, torch.Tensor):
+            return torch.flip(sample, dims=[-2])
+        return np.flip(sample, axis=-3 if sample.ndim >= 3 else 0).copy()
+
+
+class RandomRotateSet:
+    def __init__(self, angles=(0, 90, 180, 270)):
+        self.angles = list(angles)
+
+    def __call__(self, sample):
+        angle = float(self.angles[torch.randint(0, len(self.angles), (1,)).item()])
+        if isinstance(sample, torch.Tensor):
+            return TF.rotate(sample, angle)
+        tensor = torch.as_tensor(sample)
+        if tensor.ndim == 3 and tensor.shape[-1] in (1, 3):
+            tensor = tensor.permute(2, 0, 1)
+            tensor = TF.rotate(tensor, angle)
+            return tensor.permute(1, 2, 0).numpy()
+        return TF.rotate(tensor.unsqueeze(0), angle).squeeze(0).numpy()