Merge pull request #2 from jessepisel/dev

HuggingFace updates

.gitignore

@@ -160,3 +160,6 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# yarrr mac's
+.DS_Store

src/results.py → SectionSeeker/results.py

@@ -1,61 +1,73 @@
 import json
 import numpy as np
 
+
 class ResultBuilder:
     def __init__(self):
         self.results = dict()
-
-    def build(self, 
-              query_image_labels: np.ndarray, 
-              matched_labels: np.ndarray,   
-              confidence_scores: np.ndarray):
+
+    def build(
+        self,
+        query_image_labels: np.ndarray,
+        matched_labels: np.ndarray,
+        confidence_scores: np.ndarray,
+    ):
         """
         Prepare results in expected form
-        
+
         :param query_image_labels: numpy array of N reference image labels with shape [N]
         :param matched_labels: numpy array of labels of matched base images. Given N query images, this should have shape (N, 3)
         :param confidence_scores: numpy array of confidence scores for each matched based image. Given N query images, this should have shape (N, 3)
         """
-            
+
         # validate shapes of inputs
         if len(query_image_labels.shape) != 1:
-            raise ValueError(f'Expected query_image_labels to be 1-dimensional array, got {query_image_labels.shape} instead')
-
-        if matched_labels.shape != (query_image_labels.shape[0],3):
-            raise ValueError(f'Expected matched_labels to have shape {(query_image_labels.shape[0], 3)}, got {matched_labels.shape} instead')
-
-        if confidence_scores.shape != (query_image_labels.shape[0],3):
-            raise ValueError(f'Expected confidence_scores to have shape {(query_image_labels.shape[0], 3)}, got {confidence_scores.shape} instead')
-
+            raise ValueError(
+                f"Expected query_image_labels to be 1-dimensional array, got {query_image_labels.shape} instead"
+            )
+
+        if matched_labels.shape != (query_image_labels.shape[0], 3):
+            raise ValueError(
+                f"Expected matched_labels to have shape {(query_image_labels.shape[0], 3)}, got {matched_labels.shape} instead"
+            )
+
+        if confidence_scores.shape != (query_image_labels.shape[0], 3):
+            raise ValueError(
+                f"Expected confidence_scores to have shape {(query_image_labels.shape[0], 3)}, got {confidence_scores.shape} instead"
+            )
+
         for i, x in enumerate(query_image_labels):
             labels = matched_labels[i]
             confidence = confidence_scores[i]
-
-            result_x = [{'label': labels[j], 'confidence': confidence[j]} for j in range(0,3)]
-
+
+            result_x = [
+                {"label": labels[j], "confidence": confidence[j]} for j in range(0, 3)
+            ]
+
             self.results.update({x: result_x})
-        
+
         return self
-    
-    def to_json(self, path: str = '.') -> None:
+
+    def to_json(self, path: str = ".") -> None:
         """
-        Save results to json file 
-        
+        Save results to json file
+
         :param path: parent directory of result.json file
         """
-        
-        path = f'{path}/results.json'
-        with open(path, 'w+') as f:
+
+        path = f"{path}/results.json"
+        with open(path, "w+") as f:
             json.dump(self.results, f)
-
-    def __call__(self,
-              query_image_labels: np.ndarray, 
-              matched_labels: np.ndarray,   
-              confidence_scores: np.ndarray,
-              path: str = '.') -> None:
+
+    def __call__(
+        self,
+        query_image_labels: np.ndarray,
+        matched_labels: np.ndarray,
+        confidence_scores: np.ndarray,
+        path: str = ".",
+    ) -> None:
         """
         Build result and save results to json file
         """
         self.build(query_image_labels, matched_labels, confidence_scores)
         self.to_json(path)
-

src/search.py → SectionSeeker/search.py

@@ -21,7 +21,7 @@
 class ImageSet:
     """
     Subscriptapble dataset-like class for loading, storing and processing image collections
-    
+
     :param root: Path to project root directory, which contains data/image_corpus/ or data/query catalog
     :param base: Build ImageSet on top of image_corpus if True, else on top of query catalog
     :param build: Build ImageSet from filesystem instead of using saved version
@@ -30,31 +30,33 @@ class ImageSet:
     :param greyscale: Load images in grayscale if True, else use 3-channel RGB
     :param normalize: If True, images will be normalized image-wise when loaded from disk
     """
-    def __init__(self, 
-                 root: str, 
-                 base: bool = True,
-                 build: bool = False, 
-                 transform: Callable = None, 
-                 compatibility_mode: bool = False,
-                 greyscale: bool = False,
-                 normalize: bool = True) -> None:
-
+
+    def __init__(
+        self,
+        root: str,
+        base: bool = True,
+        build: bool = False,
+        transform: Callable = None,
+        compatibility_mode: bool = False,
+        greyscale: bool = False,
+        normalize: bool = True,
+    ) -> None:
+
         self.root = root
         self.compatibility_mode = compatibility_mode
         self.greyscale = greyscale
-        self.colormode = 'L' if greyscale else 'RGB'
+        self.colormode = "L" if greyscale else "RGB"
         self.transform = transform
         self.base = base
         self.normalize = normalize
-        
+
         if build:
             self.embeddings = []
             self.data, self.names = self._build()
             return
-        
+
         self.data = self._load()
-
-
+
     def _build(self) -> Tuple[torch.Tensor, str]:
 
         dirpath = f"{self.root}/data/{'image_corpus' if self.base else 'query'}"
@@ -66,39 +68,41 @@ def _build(self) -> Tuple[torch.Tensor, str]:
             # resize into common shape
             im = im.convert(self.colormode).resize((118, 143))
             if self.normalize:
-                im = cv2.normalize(np.array(im), None, 0.0, 1.0, cv2.NORM_MINMAX, cv2.CV_32FC1)
-            image = np.array(im, dtype=np.float32)    
-            fname = filename.split('/')[-1]
+                im = cv2.normalize(
+                    np.array(im), None, 0.0, 1.0, cv2.NORM_MINMAX, cv2.CV_32FC1
+                )
+            image = np.array(im, dtype=np.float32)
+            fname = filename.split("/")[-1]
             data.append(image)
             names.append(fname)
         return torch.from_numpy(np.array(data)), names
-
-    def _load(self) -> Tuple[torch.Tensor, str]:
-        ...
-
-    def save(self) -> None:
-        ...
-
+
+    def _load(self) -> Tuple[torch.Tensor, str]: ...
+
+    def save(self) -> None: ...
+
     def build_embeddings(self, model: SiameseNetwork, device: torch.cuda.device = None):
-        
+
         if device is None:
             device = detect_device()
-        
+
         with torch.no_grad():
             model.eval()
             for img, name in self:
-                img_input = img.transpose(2,0).transpose(2,1).to(device).unsqueeze(0)
+                img_input = img.transpose(2, 0).transpose(2, 1).to(device).unsqueeze(0)
                 embedding = model.get_embedding(img_input)
                 self.embeddings.append((embedding, name))
-                
+
         return self
-        
+
     def get_embeddings(self) -> List[Tuple[torch.Tensor, str]]:
         if self.embeddings is None:
-            raise RuntimeError('Embedding collection is empty. Run self.build_embeddings() method to build it')
-
+            raise RuntimeError(
+                "Embedding collection is empty. Run self.build_embeddings() method to build it"
+            )
+
         return self.embeddings
-        
+
     def __getitem__(self, index: int) -> Tuple[Any, Any]:
         """
         Args:
@@ -118,40 +122,51 @@ def __getitem__(self, index: int) -> Tuple[Any, Any]:
             img = self.transform(img)
 
         return img, name
-      
-    
+
+
 class SearchTree:
     """
     Wrapper for k-d tree built on image embeddings
-    
+
     :param query_set: instance of base ImageSet with built embedding representation
     """
+
     def __init__(self, query_set: ImageSet) -> None:
         embeddings = query_set.get_embeddings()
-        self.embeddings = np.concatenate([x[0].cpu().numpy() for x in embeddings], axis=0)
+        self.embeddings = np.concatenate(
+            [x[0].cpu().numpy() for x in embeddings], axis=0
+        )
         self.names = np.array([x[1] for x in embeddings])
         self.kdtree = self._build_kdtree()
-        
+
     def _build_kdtree(self) -> KDTree:
-        print('Building KD-Tree from embeddings')
+        print("Building KD-Tree from embeddings")
         return KDTree(self.embeddings)
-
-    def query(self, anchors: ImageSet, k: int = 3) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+
+    def query(
+        self, anchors: ImageSet, k: int = 3
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
         """
         Search for k nearest neighbors of provided anchor embeddings
-        
+
         :param anchors: instance of query (reference) ImageSet with built embedding representation
-        
-        :returns: tuple of reference_labels, distances to matched label embeddings, matched label embeddings, matched_labels 
+
+        :returns: tuple of reference_labels, distances to matched label embeddings, matched label embeddings, matched_labels
         """
-        
+
         reference = anchors.get_embeddings()
-        reference_embeddings = np.concatenate([x[0].cpu().numpy() for x in reference], axis=0)
+        reference_embeddings = np.concatenate(
+            [x[0].cpu().numpy() for x in reference], axis=0
+        )
         reference_labels = np.array([x[1] for x in reference])
-
-        distances, indices = self.kdtree.query(reference_embeddings, k=k, workers=-1)          
-        return reference_labels, distances, self.embeddings[indices], self.names[indices]
-
+
+        distances, indices = self.kdtree.query(reference_embeddings, k=k, workers=-1)
+        return (
+            reference_labels,
+            distances,
+            self.embeddings[indices],
+            self.names[indices],
+        )
+
     def __call__(self, *args, **kwargs) -> Any:
         return self.query(*args, **kwargs)
-

SectionSeeker/snn/config.py

@@ -0,0 +1,26 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class ModelConfig:
+    BACKBONE_MODEL: str = "ResNet50"
+    BACKBONE_MODEL_WEIGHTS: str = "ResNet50_Weights.IMAGENET1K_V2"
+    LATENT_SPACE_DIM: int = 8
+    FC_IN_FEATURES: int = -1
+
+
+defaultConfig = ModelConfig()
+
+vitBaseConfig = ModelConfig(
+    BACKBONE_MODEL="ViT_B_16",
+    BACKBONE_MODEL_WEIGHTS="ViT_B_16_Weights.DEFAULT",
+    LATENT_SPACE_DIM=16,
+    FC_IN_FEATURES=768,
+)
+
+vitBaseConfigPretrained = ModelConfig(
+    BACKBONE_MODEL="ViT_B_16",
+    BACKBONE_MODEL_WEIGHTS="../checkpoints/ViT_B_16_SEISMIC_SGD_28G_M75.pth",
+    LATENT_SPACE_DIM=16,
+    FC_IN_FEATURES=768,
+)