success

runners/preprocessors/rgb_feat/abstract_rgb_feat_preprocessor.py

@@ -0,0 +1,47 @@
+import os
+import json
+import numpy as np
+from abc import abstractmethod, ABC
+
+from runners.preprocessor import Preprocessor
+from utils.omni_util import OmniUtil
+
+class RGBFeatPreprocessor(Preprocessor, ABC):
+
+    def __init__(self, config_path):
+        super().__init__(config_path)
+        self.load_experiment("RGBFeat")
+        self.dataset_list_config = self.preprocess_config["dataset_list"]
+        self.model_config = self.preprocess_config["model"]
+
+    def run(self):
+        """
+        - for each dataset
+        --- get its dataloader
+        --- for each batch, do prediction
+        --- preprocess the collected results
+        --- save processed results
+        """
+        for dataset_config in self.dataset_list_config:
+            dataloader = self.get_dataloader(dataset_config)
+            model = self.get_model(self.model_config)
+            predicted_data = self.prediction(model, dataloader)
+            self.save_processed_data(predicted_data,dataset_config)
+
+    def preprocess(self, predicted_data):
+        pass
+    
+    def save_processed_data(self, processed_data, data_config=None):
+        data_path = os.path.join(str(self.experiment_path), Preprocessor.DATA, data_config["source"], data_config["data_type"])
+        for frame_path in processed_data:
+            rgb_feat = processed_data[frame_path]
+            scene = os.path.basename(os.path.dirname(frame_path))
+            idx = os.path.basename(frame_path)
+            target_scene_path = os.path.join(str(data_path), scene) 
+            if not os.path.exists(target_scene_path):
+                os.makedirs(target_scene_path)
+            rgb_feat_save_path = os.path.join(
+                target_scene_path,OmniUtil.RGB_FEAT_TEMPLATE.format(idx))
+            np.save(rgb_feat_save_path, rgb_feat)
+        print("Processed data saved to: ", data_path)
+

runners/preprocessors/rgb_feat/dinov2_preprocessor.py

@@ -0,0 +1,128 @@
+import os
+import sys
+
+path = os.path.abspath(__file__)
+for i in range(4):
+    path = os.path.dirname(path)
+PROJECT_ROOT = path
+sys.path.append(PROJECT_ROOT)
+
+import re
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+from torchvision import transforms
+from utils.omni_util import OmniUtil
+from runners.preprocessors.rgb_feat.abstract_rgb_feat_preprocessor import RGBFeatPreprocessor
+from modules.rgb_encoder.dinov2_encoder import Dinov2Encoder
+from PIL import Image
+
+from torch.utils.data import Dataset
+
+
+class Dinov2InferenceDataset(Dataset):
+    RGB_TEMPLATE = "rgb_{}.png"
+
+    def __init__(
+        self,
+        source="nbv1",
+        data_type="sample",
+        data_dir="/mnt/h/AI/Datasets",
+        image_size = 480
+    ):
+
+        self.data_dir = data_dir
+        self.data_path = str(os.path.join(self.data_dir, source, data_type))
+        self.scene_list = os.listdir(self.data_path)
+        self.data_list = self.get_datalist()
+        self.transform = transforms.Compose([           
+                        transforms.Resize(image_size),
+                        transforms.CenterCrop(int(image_size//14)*14),              
+                        transforms.ToTensor(),                    
+                        transforms.Normalize(mean=0.5, std=0.2)
+                        ])
+
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        frame_path = self.data_list[index]
+        frame_data = self.load_frame_data(frame_path=frame_path)
+        return frame_data
+
+    def get_datalist(self):
+        for scene in self.scene_list:
+            scene_path = os.path.join(self.data_path, scene)
+            file_list = os.listdir(scene_path)
+            scene_frame_list = []
+            for file in file_list:
+                if file.startswith("camera_params"):
+                    frame_index = re.findall(r"\d+", file)[0]
+                    frame_path = os.path.join(scene_path, frame_index)
+                    scene_frame_list.append(frame_path)
+
+        return scene_frame_list
+
+    def load_frame_data(self, frame_path):
+        rgb = OmniUtil.get_rgb(frame_path)
+        rgb = Image.fromarray(rgb)
+        rgb = self.transform(rgb)
+        ret_dict = {"rgb": rgb, "frame_path": frame_path}
+        return ret_dict
+
+
+class Dinov2Preprocessor(RGBFeatPreprocessor):
+    MODULE_NAME: str = "dinov2"
+
+    def __init__(self, config_path):
+        super().__init__(config_path)
+
+    def get_dataloader(self, dataset_config):
+
+        dataset = Dinov2InferenceDataset(
+            source=dataset_config["source"],
+            data_type=dataset_config["data_type"],
+            data_dir=dataset_config["data_dir"],
+            image_size = dataset_config["image_size"]
+        )
+        print("Test dataset length: ", len(dataset))
+        dataloader = DataLoader(
+            dataset,
+            batch_size=dataset_config["batch_size"],
+            shuffle=False,
+            num_workers=0,
+        )
+        print("Test dataloader length: ", len(dataloader))
+        return dataloader
+
+    def get_model(self, model_config=None):
+        model = Dinov2Encoder(model_config["general"]["model_name"])
+        model.to("cuda")
+        return model
+
+    def prediction(self, model, dataloader):
+        results = {}
+        total = len(dataloader)
+        for idx, batch_data in enumerate(dataloader):
+            rgb = batch_data["rgb"].to("cuda")
+            with torch.no_grad():
+                rgb_feat = model.encode_rgb(rgb)
+            frame_paths = batch_data["frame_path"]
+            for i, frame_path in enumerate(frame_paths):
+                results[frame_path] = rgb_feat[i].cpu().numpy()
+            print(f"Processed {idx}/{total} batches")
+                
+        return results
+
+    def visualize_feature(self, rgb_feat, model_name, save_path=None):
+        model = Dinov2Encoder(model_name)
+        model.visualize_features(rgb_feat,save_path)
+
+
+if __name__ == "__main__":
+    rgb_preproc = Dinov2Preprocessor(config_path="configs/server_rgb_feat_preprocess_config.yaml")
+    #ßrgb_preproc.run()
+    rgb_feat = np.load("experiments/rgb_feat_preprocessor_test/data/nbv1/sample/scene_0/rgb_feat_0405.npy")
+    
+    rgb_preproc.visualize_feature(rgb_feat, "dinov2_vits14", './visualize.png')