thread optimizatin

app.py

@@ -252,7 +252,6 @@
 
 
 
-
 import gradio as gr
 import cv2
 import numpy as np
@@ -274,11 +273,11 @@ logger = logging.getLogger(__name__)
 
 class FeatureExtractor:
     def __init__(self):
-        # Using ResNet50 for 2048-D feature vectors
-        backbone = models.resnet50(weights="IMAGENET1K_V1")
-        self.model = nn.Sequential(*list(backbone.children())[:-1])
-        self.model.eval()
-
+       
+        self.backbone = models.resnet50(weights="IMAGENET1K_V1")
+        self.backbone.eval()
+        
+        # We transform to standard ImageNet resolution
         self.transform = transforms.Compose([
             transforms.Resize((224, 224)),
             transforms.ToTensor(),
@@ -298,33 +297,24 @@ class FeatureExtractor:
         if len(rgb.shape) == 2:
             rgb = cv2.cvtColor(rgb, cv2.COLOR_GRAY2RGB)
         
-        # We want the layer BEFORE the global pooling to get spatial info
-        # resnet.layer4 is the last block
-        # self.model is nn.Sequential(*list(backbone.children())[:-1])
-        # children()[:-1] = [conv1, bn1, relu, maxpool, layer1, layer2, layer3, layer4]
-        
         input_tensor = self.transform(Image.fromarray(rgb)).unsqueeze(0)
         
-        # Get activations from the last conv layer (Layer 4)
+        # Optimized inference: Use the pre-loaded backbone
         with torch.no_grad():
-            # Run through the layers up to global pooling
-            # Using the original backbone for Easier Access to sub-layers
-            backbone = models.resnet50(weights="IMAGENET1K_V1")
-            backbone.eval()
-            
-            x = backbone.conv1(input_tensor)
-            x = backbone.bn1(x)
-            x = backbone.relu(x)
-            x = backbone.maxpool(x)
-            x = backbone.layer1(x)
-            x = backbone.layer2(x)
-            x = backbone.layer3(x)
-            features_spatial = backbone.layer4(x) # [1, 2048, 7, 7]
+            # Walk through layers to capture spatial activations before global pooling
+            x = self.backbone.conv1(input_tensor)
+            x = self.backbone.bn1(x)
+            x = self.backbone.relu(x)
+            x = self.backbone.maxpool(x)
+            x = self.backbone.layer1(x)
+            x = self.backbone.layer2(x)
+            x = self.backbone.layer3(x)
+            features_spatial = self.backbone.layer4(x) # [1, 2048, 7, 7]
             
-            # Global Average Pooling to get the vector
+            # Global Average Pooling (L2 distance is more effective on normalized vectors)
             feat = torch.mean(features_spatial, dim=[2, 3]).squeeze().cpu().numpy()
             
-            # Create Heatmap: sum across channels to see "hot" regions
+            # Heatmap generation: Sum across channels to highlight activated regions
             amap = torch.sum(features_spatial, dim=1).squeeze().cpu().numpy()
             amap = np.maximum(amap, 0)
             amap /= (np.max(amap) + 1e-8)
@@ -332,11 +322,11 @@ class FeatureExtractor:
             amap = np.uint8(255 * amap)
             heatmap = cv2.applyColorMap(amap, cv2.COLORMAP_JET)
             
-            # Overlay heatmap on original image
-            # Convert BGR heatmap to RGB
+            # Overlay BGR heatmap on RGB image properly
             heatmap_rgb = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
             overlay = cv2.addWeighted(rgb, 0.6, heatmap_rgb, 0.4, 0)
 
+        # Vector normalization for Cosine Similarity
         norm = np.linalg.norm(feat)
         return (feat / norm if norm > 1e-8 else feat), overlay