Update test_detector.py

test_detector.py

@@ -15,14 +15,15 @@ from PIL import Image
 from huggingface_hub import hf_hub_download
 import gc
 
-from base_miner.UCF.config.constants import CONFIGS_DIR, WEIGHTS_DIR
-from base_miner.gating_mechanisms import FaceGate
+from bitmind.image_transforms import ucf_transforms, ConvertToRGB, CenterCrop, CLAHE
+from arena.detectors.UCF.config.constants import CONFIGS_DIR, WEIGHTS_DIR
+from arena.detectors.gating_mechanisms import FaceGate
 
-from base_miner.UCF.detectors import DETECTOR
-from base_miner.deepfake_detectors import DeepfakeDetector
-from base_miner import DETECTOR_REGISTRY, GATE_REGISTRY
+from arena.detectors.UCF.detectors import DETECTOR
+from arena.detectors.deepfake_detectors import DeepfakeDetector
+from arena.detectors import DETECTOR_REGISTRY, GATE_REGISTRY
+from arena.utils.image_transforms import CLAHE
 
-import bittensor as bt
 
 @DETECTOR_REGISTRY.register_module(module_name='test')
 class TestDetector(DeepfakeDetector):
@@ -34,11 +35,11 @@ class TestDetector(DeepfakeDetector):
         model_name (str): Name of the detector instance.
         config (str): Name of the YAML file in deepfake_detectors/config/ to load
                       attributes from.
-        device (str): The type of device ('cpu' or 'cuda').
+        cuda (bool): Whether to enable cuda (GPU).
     """
     
-    def __init__(self, model_name: str = 'UCF', config: str = 'ucf.yaml', device: str = 'cpu'):
-        super().__init__(model_name, config, device)
+    def __init__(self, model_name: str = 'UCF', config: str = 'ucf.yaml', cuda: bool = True):
+        super().__init__(model_name, config, cuda)
     
     def ensure_weights_are_available(self, weight_filename):
         destination_path = Path(WEIGHTS_DIR) / Path(weight_filename)
@@ -51,7 +52,7 @@ class TestDetector(DeepfakeDetector):
 
     def load_train_config(self):
         destination_path = Path(CONFIGS_DIR) / Path(self.train_config)
-    
+
         if not destination_path.exists():
             local_config_path = hf_hub_download(self.hf_repo, self.train_config)
             print(f"Downloaded {self.hf_repo}/{self.train_config} to {local_config_path}")
@@ -83,9 +84,8 @@ class TestDetector(DeepfakeDetector):
         self.init_cudnn()
         self.init_seed()
         self.ensure_weights_are_available(self.weights)
-        self.ensure_weights_are_available(self.train_config['pretrained'].split('/')[-1])
+        self.ensure_weights_are_available(self.backbone_weights)
         model_class = DETECTOR[self.train_config['model_name']]
-        bt.logging.info(f"Loaded config from training run: {self.train_config}")
         self.model = model_class(self.train_config).to(self.device)
         self.model.eval()
         weights_path = Path(WEIGHTS_DIR) / self.weights
@@ -115,17 +115,20 @@ class TestDetector(DeepfakeDetector):
             torch.Tensor: The preprocessed image tensor, ready for model inference.
         """
         # Convert image to RGB format to ensure consistent color handling.
-        image = image.convert('RGB')
-    
-        # Define transformation sequence for image preprocessing.
-        transform = transforms.Compose([
-            transforms.Resize((res, res), interpolation=Image.LANCZOS),  # Resize image to specified resolution.
-            transforms.ToTensor(),  # Convert the image to a PyTorch tensor.
-            transforms.Normalize(mean=self.train_config['mean'], std=self.train_config['std'])  # Normalize the image tensor.
-        ])
+        tforms = [
+            ConvertToRGB(),
+            CenterCrop(),
+            transforms.Resize((256,256)),
+            transforms.ToTensor()
+        ]
+        if self.dataset_type == 'real':
+            tforms.append(transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]))
+
+        transform = transforms.Compose(tforms)
         
         # Apply transformations and add a batch dimension for model inference.
         image_tensor = transform(image).unsqueeze(0)
+
         
         # Move the image tensor to the specified device (e.g., GPU).
         return image_tensor.to(self.device)