Class distribution benchmark added

vla/benchmarks/class_distribution_using_segmentation/DinoV2Features.py

@@ -0,0 +1,31 @@
+import cv2
+import torch
+import numpy as np
+import torchvision.transforms as T
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# =====================================================================
+# Load DINOv2
+# =====================================================================
+dinov2 = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14').eval().to(device)
+
+# =====================================================================
+# Preprocessing
+# =====================================================================
+preprocessor = T.Compose([
+    T.ToTensor(),
+    T.Resize((504, 504)),
+    T.Normalize(mean=[0.485, 0.456, 0.406],
+                std=[0.229, 0.224, 0.225]),
+])
+
+
+def extract_full_features(img):
+    with torch.no_grad():
+        full_feats = dinov2.get_intermediate_layers(img, n=1)[0]
+        return full_feats
+
+class Dinov2Features_full:
+    def encode(self, img):
+        return extract_full_features(img)

vla/benchmarks/class_distribution_using_segmentation/__init__.py

@@ -0,0 +1 @@
+from .train import benchmark

vla/benchmarks/class_distribution_using_segmentation/builders.py

@@ -0,0 +1,97 @@
+import json
+import os
+import random
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import List, Tuple, Dict, Any, Optional
+from PIL import Image
+import numpy as np
+from tqdm import tqdm
+
+# Helpers
+
+
+def safe_load_json(path: str) -> Optional[dict]:
+    if not os.path.exists(path):
+        return None
+    try:
+        with open(path, "r") as f:
+            return json.load(f)
+    except Exception:
+        return None
+
+
+# Base Builder (Abstract)
+
+
+class DatasetBuilder(ABC):
+    def __init__(self, directory: str, dataset_size: int = 200, train_split: float = 0.25, random_seed=None, full_folder=False):
+        self.directory = directory
+        self.dataset_size = dataset_size
+        self.train_split = train_split
+        self.random_seed = random_seed
+        self.full_folder = full_folder
+        if self.random_seed is not None:
+            random.seed(self.random_seed)
+
+    @abstractmethod
+    def extract_samples(self) -> List[Tuple[List, str]]:
+        """
+        Must return:
+            [([class_distribution_probs], image_path)]
+        """
+        ...
+
+    @abstractmethod
+    def build(self, json_name: str) -> Tuple[str, str, list, list]:
+        """
+        Must return:
+            (train_json_path, test_json_path)
+        """
+        ...
+
+# Class Distribution Builder
+
+class ClassDistributionDatasetBuilder(DatasetBuilder):
+
+    def extract_samples(self) -> List[Tuple[List, str]]:
+        """Gather (class_probs, image_stem_path)."""
+        paths = []
+        basepath = Path(self.directory)
+        images_subfolders = list(basepath.rglob("images"))
+
+        for img_folder_path in images_subfolders:
+            distrs_path = img_folder_path.parent / "distributions"
+            images_folder = list(img_folder_path.rglob("*.png"))
+            for img_path in images_folder:
+                distr_path = (distrs_path / img_path.stem).with_suffix(".npy")
+                if not distr_path.exists():
+                    continue
+                paths.append((img_path, distr_path))
+        return paths
+
+    def build(self, json_name: str) -> Tuple[str, str, list, list]:
+        samples = self.extract_samples()
+        if not samples:
+            raise Exception("No valid samples for type dataset.")
+
+        if self.full_folder:
+            train_split = 0
+        else:
+            train_split = int(self.train_split * len(samples))
+
+        random.shuffle(samples)
+
+        train_out = samples[:train_split]
+        test_out = samples[train_split:]
+
+        if self.random_seed is not None:
+            return "", "", train_out, test_out
+
+        train_file = f"train_{json_name}.json"
+        test_file = f"test_{json_name}.json"
+
+        json.dump(train_out, open(train_file, "w"))
+        json.dump(test_out, open(test_file, "w"))
+
+        return train_file, test_file, [], []

vla/benchmarks/class_distribution_using_segmentation/config.py

@@ -0,0 +1,32 @@
+from enum import Enum
+import warnings
+
+class ConfigPaths:
+    path_to_raw_data = None
+    feature_store_path = "./precomputed"
+
+class TrainConfig:
+    learning_rate = 1e-3
+    momentum = 0.9
+    epochs = 10
+    output_activation = "none"  # "none", "softmax", "sigmoid"
+    num_layers = 1
+    layers_sizes = [512]
+
+def construct_configs(path_to_raw_data=None, feature_store_path=None,
+        learning_rate=1e-3, momentum=0.9, epochs=10, output_activation="none", num_layers=1, layers_sizes=[512]):
+    TrainConfig.learning_rate = learning_rate
+    TrainConfig.momentum = momentum
+    TrainConfig.epochs = epochs
+    TrainConfig.output_activation = output_activation
+    TrainConfig.num_layers = num_layers
+    TrainConfig.layers_sizes = layers_sizes
+    if not path_to_raw_data:
+        raise ValueError("Path to raw data must be set")
+    else:
+        ConfigPaths.path_to_raw_data = path_to_raw_data
+    if not feature_store_path:
+        warnings.warn("Feature store path is not set by config. If benchmark will be ran with use_precomputed_features "
+                      "default folder ./precomputed will be used to store features")
+    else:
+        ConfigPaths.feature_store_path = feature_store_path

vla/benchmarks/class_distribution_using_segmentation/dataset.py

@@ -0,0 +1,34 @@
+import torch
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+
+from PIL import Image
+
+# PyTorch Dataset
+
+
+class LOSSegmDataset(Dataset):
+    def __init__(self, items, transform=None, feature_store=None):
+        self.items = items
+        self.transform = transform
+        self.feature_store = feature_store
+
+    def __len__(self):
+        return len(self.items)
+
+    def load_image(self, img_path):
+        image = Image.open(img_path).convert("RGB")
+
+        if self.transform:
+            image = self.transform(image)
+        return image
+
+    def __getitem__(self, idx):
+        img_path, distr_path = self.items[idx]
+        distribution = torch.from_numpy(np.load(distr_path))
+        feature_path = str(img_path).strip(".png").replace("/", "_")
+        if self.feature_store and self.feature_store.exists(feature_path):
+            features = self.feature_store.load(feature_path)
+            return distribution, features
+
+        return distribution, self.load_image(img_path)

vla/benchmarks/class_distribution_using_segmentation/example_config.json

@@ -0,0 +1,10 @@
+{
+  "path_to_raw_data" : "./2026_LOS_SEGM",
+  "feature_store_path" : "./precomputed",
+  "learning_rate" : 1e-3,
+  "momentum" : 0.9,
+  "epochs" : 20,
+  "output_activation" : "none",
+  "num_layers" : 2,
+  "layers_sizes" : [128, 128]
+}

vla/benchmarks/class_distribution_using_segmentation/features.py

@@ -0,0 +1,25 @@
+import os
+import torch
+from pathlib import Path
+from tqdm import tqdm
+
+class FeatureStore:
+    """
+    Saves and loads precomputed features to avoid running DINO inside dataloaders.
+    """
+
+    def __init__(self, root="./precomputed_features"):
+        self.root = Path(root)
+        self.root.mkdir(parents=True, exist_ok=True)
+
+    def feature_path(self, img_stem: str) -> Path:
+        return self.root / f"{Path(img_stem).name}.pt"
+
+    def exists(self, img_stem: str) -> bool:
+        return self.feature_path(img_stem).exists()
+
+    def save(self, img_stem: str, tensor: torch.Tensor):
+        torch.save(tensor.cpu(), self.feature_path(img_stem))
+
+    def load(self, img_stem: str) -> torch.Tensor:
+        return torch.load(self.feature_path(img_stem))

vla/benchmarks/class_distribution_using_segmentation/images_to_distribution.py

@@ -0,0 +1,37 @@
+# This file is used to convert segmented images to npy files which contains class distribution
+import numpy as np
+from pathlib import Path
+from PIL import Image
+
+paths = []
+unique_colors_list = []
+basepath = Path("./2026_LOS_SEGM")
+segmentation_subfolders = list(basepath.rglob("segmentation"))
+
+stub = 0
+
+for segm_path in segmentation_subfolders:
+    images_folder = list(segm_path.rglob("*.png"))
+    for segm_image_path in images_folder:
+        segm_img = np.array(Image.open(segm_image_path))
+        unique = list(np.unique(segm_img.reshape(-1, segm_img.shape[2]), axis=0))
+        for unique_item in unique:
+            unique_item = list(unique_item)
+            if unique_item not in unique_colors_list:
+                unique_colors_list.append(unique_item)
+
+unique_colors_dict = {str(value): index for index, value in enumerate(unique_colors_list)}
+samples = []
+num_classes = len(unique_colors_dict)
+for segm_path in segmentation_subfolders:
+    distribution_path = segm_path.parent / 'distributions'
+    distribution_path.mkdir(parents=True, exist_ok=True)
+    images_folder = list(segm_path.rglob("*.png"))
+    for segm_image_path in images_folder:
+        segm_img = np.array(Image.open(segm_image_path))
+        class_probs = [0] * num_classes
+        color_list = list(segm_img.reshape(-1, segm_img.shape[2]))
+        for color in color_list:
+            class_probs[unique_colors_dict[str(list(color))]] += 1
+        class_probs = np.array(class_probs) / len(color_list)
+        np.save(f'{str(distribution_path)}/{segm_image_path.stem}.npy', class_probs)

vla/benchmarks/class_distribution_using_segmentation/manager.py

@@ -0,0 +1,57 @@
+import json
+from torch.utils.data import DataLoader
+from builders import ClassDistributionDatasetBuilder
+from dataset import LOSSegmDataset
+import numpy as np
+
+
+class DatasetManager:
+    def __init__(self, directory, dataset_size=200, train_split=0.25,
+                 random_seed=None, full_folder=False):
+        self.random_seed = random_seed
+
+        self.builder = ClassDistributionDatasetBuilder(directory, dataset_size=dataset_size,
+                                          train_split=train_split, random_seed=self.random_seed,
+                                          full_folder=full_folder)
+
+    def create(self, name):
+        if self.random_seed is not None:
+            _, _, train_set, test_set = self.builder.build(name)
+            return train_set, test_set
+        train_file, test_file, [], [] = self.builder.build(name)
+        print(f"Train set: {train_file}")
+        print(f"Test  set: {test_file}")
+        return train_file, test_file
+
+
+def load_json(path):
+    with open(path, "r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+def make_dataloaders(train_json, test_json, transform=None, feature_store=None, batch=32, workers=2, dataset_manager=None,
+                     generalization_dataset_manager=None):
+
+    if dataset_manager is not None:
+        train_items, test_items = dataset_manager.create("")
+    else:
+        train_items = load_json(train_json)
+        test_items = load_json(test_json)
+
+    generalization_items = []
+    if generalization_dataset_manager is not None:
+        _, generalization_items = generalization_dataset_manager.create("")
+
+    nc = len(np.load(train_items[0][1]))
+
+    train_ds = LOSSegmDataset(train_items, transform, feature_store)
+    test_ds = LOSSegmDataset(test_items, transform, feature_store)
+    generalization_ds = LOSSegmDataset(generalization_items, transform, feature_store)
+
+    return (
+        DataLoader(train_ds, batch_size=batch, shuffle=True, num_workers=workers),
+        DataLoader(test_ds, batch_size=batch, shuffle=False, num_workers=workers),
+        DataLoader(test_ds, batch_size=1, shuffle=False),
+        DataLoader(generalization_ds, batch_size=1, shuffle=False),
+        nc
+    )