Update 4 files

- /src/models/clip_vit.py
- /src/models/resnet_50.py
- /src/models/mobilenet_v3.py
- /src/classification_model.py

src/classification_model.py

@@ -4,6 +4,9 @@ from PIL import Image
 from .data.model_data import ModelData
 from .models.mobilenet_v3 import MobilenetV3
 from .models.clip_vit import ClipVit
+from .models.google_vit import GoogleVit
+from .models.resnet_50 import Resnet50
+
 from .data.classification_result import ClassificationResult
 
 class ClassificationModel:
@@ -26,7 +29,9 @@ class ClassificationModel:
     def load_model(self):
         self.models = [
             ModelData('clip-vit-base-patch32', model_class=ClipVit()), 
-            ModelData('mobilenet_v3', model_class=MobilenetV3())
+            ModelData('mobilenet_v3', model_class=MobilenetV3()),
+            ModelData('google-vit-base-patch16-224', model_class=GoogleVit()),
+            ModelData('microsoft/resnet-50', model_class=Resnet50())
             ]
 
     def classify(self, model_name, image) -> List[ClassificationResult]:

src/models/clip_vit.py

@@ -1,12 +1,37 @@
+from PIL import Image
+import torch
+from transformers import CLIPProcessor, CLIPModel
 from typing import List
 from src.interface import ModelInterface
 from src.data.classification_result import ClassificationResult
 
 class ClipVit(ModelInterface):
     def __init__(self):
-        print('init... vlip vit model')
+        print('Initializing CLIP VIT model...')
+        # Load pre-trained CLIP model and processor
+        self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+        self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 
     def classify_image(self, image) -> List[ClassificationResult]:
-        class_name = "Example Result"
-        confidence = 0.85
-        return [ClassificationResult(class_name=class_name, confidence=confidence)]
+        # Preprocess the image using CLIPProcessor
+        inputs = self.processor(text=["volcano", "mountain", "alp", "mount", "valley"], images=image, return_tensors="pt", padding=True)
+
+        # Perform inference
+        outputs = self.model(**inputs)
+        logits_per_image = outputs.logits_per_image  # This is the image-text similarity score
+
+        # Convert logits to probabilities using softmax (using PyTorch)
+        probabilities = torch.nn.functional.softmax(logits_per_image, dim=1).detach().numpy()
+
+        # Get the top 5 predicted classes and their probabilities using torch.argsort
+        top_indices = torch.argsort(torch.from_numpy(probabilities), dim=1, descending=True)[0, :5]
+        top_indices = top_indices.tolist()
+        top_probabilities = probabilities[0, top_indices]
+
+        # Get the class labels from the processor's tokenizer
+        class_name = ["volcano", "mountain", "alp", "mount", "valley"]
+
+        # Create a list of ClassificationResult objects with predicted classes and probabilities
+        result = [ClassificationResult(class_name=str(name), confidence=float(probabilities)) for name, probabilities in zip(class_name, top_probabilities)]
+
+        return result

src/models/mobilenet_v3.py

@@ -1,13 +1,44 @@
 from typing import List
-import random
+import torch
+import timm
 from src.interface import ModelInterface
 from src.data.classification_result import ClassificationResult
+from PIL import Image
+import urllib.request
 
 class MobilenetV3(ModelInterface):
     
     def __init__(self):
         print('init... mobilenet v3 model')
+        self.model = timm.create_model('mobilenetv3_large_100', pretrained=True).eval()
+
+        # Download and read class labels
+        url, filename = ("https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt", "imagenet_classes.txt")
+        urllib.request.urlretrieve(url, filename)
+        with open(filename, "r") as f:
+            self.class_labels = [s.strip() for s in f.readlines()]
 
     def classify_image(self, image) -> List[ClassificationResult]:
-        results = [ClassificationResult(class_name=f'example class ({i+1})', confidence=random.uniform(0, 1.0)) for i in range(5)]
+  
+        # Get model-specific transforms (normalization, resize)
+        data_config = timm.data.resolve_model_data_config(self.model)
+        transforms = timm.data.create_transform(**data_config, is_training=False)
+        input_tensor = transforms(image).unsqueeze(0)
+
+        # Perform inference
+        with torch.no_grad():
+            output = self.model(input_tensor)
+
+        # Get the top 5 predictions
+        probabilities, top5_class_indices = torch.topk(output.softmax(dim=1), k=5)
+
+        # Create ClassificationResult objects with confidence information
+        results = [
+            ClassificationResult(
+                class_name=self.class_labels[top5_class_indices[0][i].item()],
+                confidence=probabilities[0][i].item()
+            )
+            for i in range(5)
+        ]
+
         return results

src/models/resnet_50.py

@@ -0,0 +1,36 @@
+from typing import List
+from src.interface import ModelInterface
+from src.data.classification_result import ClassificationResult
+from transformers import AutoImageProcessor, ResNetForImageClassification
+import torch
+
+class Resnet50(ModelInterface):
+    def __init__(self):
+        print('init... clip vit model')
+        self.processor = AutoImageProcessor.from_pretrained("microsoft/resnet-50")
+        self.model = ResNetForImageClassification.from_pretrained("microsoft/resnet-50")
+
+    def classify_image(self, image) -> List[ClassificationResult]:
+        # Preprocess the image
+        inputs = self.processor(images=image, return_tensors="pt")
+
+        # Perform inference
+        outputs = self.model(**inputs)
+        logits = outputs.logits.detach().numpy()
+
+        # Convert logits to probabilities using softmax (using PyTorch)
+        probabilities = torch.nn.functional.softmax(torch.from_numpy(logits), dim=-1).numpy()
+
+        # Get the top 5 predictions
+        top_5 = torch.argsort(torch.from_numpy(probabilities), axis=-1, descending=True)[0][:5].numpy()
+
+        # Create ClassificationResult objects with confidence information
+        results = [
+            ClassificationResult(
+                class_name=self.model.config.id2label[top_5[i]],
+                confidence=float(probabilities[0][top_5[i]])
+            )
+            for i in range(5)
+        ]
+
+        return results