Added custom model trained on Eurosat

Model performans on eurosat is not very good.
It will be improved in the next commit.

EUROSAT_CUSTOM_MODEL.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:68981cd12c767b20d0a98970d41d823bbb0a91ea23debf85a70dc2643ab5d619
+size 33657519

app.py

@@ -14,6 +14,16 @@ from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
 import wget
 import cv2
 matplotlib.use('agg')
+import os
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, random_split
+from torchvision import datasets, transforms, models
+import torch.optim as optim
+import torch.nn.functional as F
+import matplotlib.pyplot as plt
+from torchvision.transforms import v2
 
 # Vanilla Legendre between [0,1]
 def Pn(m, x):    
@@ -28,6 +38,32 @@ def Pn(m, x):
 def L(a,b,m,x):
     return np.sqrt((2*m+1)/(b-a))*Pn(m, 2*(x-b)/(b-a)+1)
 
+eurosat_transform = v2.Compose([
+    v2.Resize((64, 64)),  # Resize images to 64x64
+    #transforms.ToTensor(),  # Convert images to PyTorch tensors,
+    v2.ToDtype(torch.float32),
+    v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # Normalize images
+])
+
+class CNN(nn.Module):
+    def __init__(self, num_classes=10):  # Modify num_classes based on the number of your classes
+        super(CNN, self).__init__()
+        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
+        self.pool = nn.MaxPool2d(2, 2)
+        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
+        self.fc1 = nn.Linear(64 * 16 * 16, 512)
+        self.fc2 = nn.Linear(512, num_classes)
+        self.dropout = nn.Dropout(0.25)
+
+    def forward(self, x):
+        x = self.pool(F.relu(self.conv1(x)))
+        x = self.pool(F.relu(self.conv2(x)))
+        x = x.view(-1, 64 * 16 * 16)
+        x = self.dropout(x)
+        x = F.relu(self.fc1(x))
+        x = self.fc2(x)
+        return x
+
 def run_lime(input_image, 
              model_name: str, 
              top_labels: int,
@@ -43,13 +79,11 @@ def run_lime(input_image,
     print('batch_size', batch_size)
     print('input image', type(input_image), input_image.shape)
 
-    model, weights = fetch_model(model_name)
-    preprocess = weights.transforms(antialias=True)
+    model, weights, preprocess, names = fetch_model(model_name)
 
     input_image_processed = preprocess(torch.from_numpy(input_image.transpose(2,0,1))).unsqueeze(0)
     logits = model(input_image_processed)
     probs = F.softmax(logits, dim=1)
-    names = weights.meta['categories']
 
     top_10_classes = []
     print('probs', type(probs), probs.shape)
@@ -149,9 +183,7 @@ def run_hdmr(input_image,
     print('num_features_hdmr', num_features_hdmr)
     print('input image', type(input_image), input_image.shape)
 
-    model, weights = fetch_model(model_name)
-    preprocess = weights.transforms(antialias=True)
-
+    model, weights, preprocess, names = fetch_model(model_name)
     sam_model = fetch_sam_model(sam_model_name)
     mask_generator = SamAutomaticMaskGenerator(sam_model)
     masks = mask_generator.generate(input_image)
@@ -192,7 +224,7 @@ def run_hdmr(input_image,
         l2_distance = np.linalg.norm(logits_normalized - logits_sample_normalized)
         class_id = probs.argmax().item()
         score = probs[class_id].item()
-        category_name = weights.meta["categories"][class_id]
+        category_name = names[class_id]
         print(f"sample:{sample:2d} cosine: {cosine_distance:.5f} l1: {l1_distance:.5f}  l2: {l2_distance:.5f} {category_name}: {100 * score:.1f}%")
 
         y[:,sample] = [cosine_distance, l1_distance, l2_distance]
@@ -225,18 +257,29 @@ def fetch_sam_model(sam_model_name_checkpoint):
     return sam
 
 def fetch_model_names():
-    return models.list_models(module=torchvision.models)
+    model_names = models.list_models(module=torchvision.models)
+    return ['EUROSAT_CUSTOM_MODEL'] + model_names
 
 def fetch_model(model_name):
     print('Retrieving model ', model_name)
+    if model_name == "EUROSAT_CUSTOM_MODEL":
+        model = CNN() 
+        weights = torch.load('EUROSAT_CUSTOM_MODEL.pth')
+        model.load_state_dict(weights)
+        return (model, weights, eurosat_transform,
+                ['AnnualCrop','Forest','HerbaceousVegetation','Highway',
+                 'Industrial','Pasture','PermanentCrop','Residential',
+                 'River','SeaLake'])
     weights_enum = models.get_model_weights(model_name)
     for w in weights_enum:
         if "IMAGENET1K" in w.name:
             weights = w
             model = models.get_model(model_name, weights=weights)
             print('Model weights loaded', w.name)
-            return model, weights
-    return None, None
+            return (model, weights, 
+                    weights.transforms(antialias=True), 
+                    weights.meta['categories'])
+    return None, None, None, None
 
 with gd.Blocks() as demo:
     with gd.Column():
@@ -255,7 +298,7 @@ with gd.Blocks() as demo:
                                          Select the image classification model to use for LIME.
                                          The list is automatically populated by using torchvision library.
                                          ''',
-                                         value='convnext_tiny',
+                                         value='EUROSAT_CUSTOM_MODEL',
                                          choices=fetch_model_names())
                 sam_model_name = gd.Dropdown(label="SAM model",
                             info='Select the SAM model',
@@ -325,4 +368,3 @@ with gd.Blocks() as demo:
   
 if __name__ == "__main__":
     demo.launch()
-    

requirements.txt

@@ -4,4 +4,5 @@ opencv-python==4.7.0.72
 lime==0.2.0.1
 scikit-image==0.20.0
 torch==2.0.0
-wget==3.2
+wget==3.2
+torchaudio==2.0.1