Update app.py

app.py

@@ -1,14 +1,18 @@
 import torch
+from torch import Tensor as T
 import torchvision.models as models
 import torchvision.transforms as transforms
 import torchvision.datasets as datasets
 from torchvision.transforms import Compose
+from torch.nn import Module
+from torch.nn.functional import softmax
 import requests
+from PIL import Image
 import random
-import gradio as gr
+from gradio import Blocks, Tabs, TabItem, Row, Column, Image, Dropdown, Button, Label
 
 # Predefined models available in torchvision
-image_prediction_models = {
+IMAGE_PREDICTION_MODELS = {
     'resnet': models.resnet50,
     'alexnet': models.alexnet,
     'vgg': models.vgg16,
@@ -29,10 +33,10 @@ image_prediction_models = {
 
 # Load a pretrained model from torchvision
 class ModelLoader:
-    def __init__(self, model_dict):
+    def __init__(self, model_dict : dict):
         self.model_dict = model_dict
 
-    def load_model(self, model_name):
+    def load_model(self, model_name : str) -> Module :
         model_name_lower = model_name.lower()
         if model_name_lower in self.model_dict:
             model_class = self.model_dict[model_name_lower]
@@ -41,7 +45,7 @@ class ModelLoader:
         else:
             raise ValueError(f"Model {model_name} is not available for image prediction in torchvision.models")
 
-    def get_model_names(self):
+    def get_model_names(self) -> list:
         return [name.capitalize() for name in self.model_dict.keys()]
 
 # Preprocessor: Prepares image for model input
@@ -49,7 +53,7 @@ class Preprocessor:
     def __init__(self):
         self.normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 
-    def preprocess(self, model_name):
+    def preprocess(self, model_name : str) -> Compose:
         input_size = 224
         if model_name == 'inception':
             input_size = 299
@@ -62,31 +66,31 @@ class Preprocessor:
 
 # Postprocessor: Processes model output
 class Postprocessor:
-    def __init__(self, labels):
+    def __init__(self, labels : list):
         self.labels = labels
 
-    def postprocess_default(self, output):
-        probabilities = torch.nn.functional.softmax(output[0], dim=0)
-        top_prob, top_catid = torch.topk(probabilities, 5)
+    def postprocess_default(self, output) -> dict:
+        probabilities = softmax(output[0], dim=0)
+        top_prob , top_catid = torch.topk(probabilities, 5)
         confidences = {self.labels[top_catid[i].item()]: top_prob[i].item() for i in range(top_prob.size(0))}
         return confidences
 
-    def postprocess_inception(self, output):
-        probabilities = torch.nn.functional.softmax(output[1], dim=0)
+    def postprocess_inception(self, output) -> dict:
+        probabilities : T = softmax(output[1], dim=0)
         top_prob, top_catid = torch.topk(probabilities, 5)
         confidences = {self.labels[top_catid[i].item()]: top_prob[i].item() for i in range(top_prob.size(0))}
         return confidences
 
 # ImageClassifier: Classifies images using a selected model
 class ImageClassifier:
-    def __init__(self, model_loader, preprocessor, postprocessor):
+    def __init__(self, model_loader : ModelLoader, preprocessor: Preprocessor, postprocessor : Postprocessor):
         self.model_loader = model_loader
         self.preprocessor = preprocessor
         self.postprocessor = postprocessor
 
-    def classify(self, input_image, selected_model):
-        preprocess_input = self.preprocessor.preprocess(model_name=selected_model)
-        input_tensor = preprocess_input(input_image)
+    def classify(self, input_image : Image, selected_model : str) -> dict:
+        preprocess_input : Compose = self.preprocessor.preprocess(model_name=selected_model)
+        input_tensor : T = preprocess_input(input_image)
         input_batch = input_tensor.unsqueeze(0)
         model = self.model_loader.load_model(selected_model)
 
@@ -96,7 +100,7 @@ class ImageClassifier:
 
         model.eval()
         with torch.no_grad():
-            output = model(input_batch)
+            output : T = model(input_batch)
 
         if selected_model.lower() == 'inception':
             return self.postprocessor.postprocess_inception(output)
@@ -105,83 +109,64 @@ class ImageClassifier:
 
 # CIFAR10ImageProvider: Provides random images from CIFAR-10 dataset
 class CIFAR10ImageProvider:
-    def __init__(self, dataset_root='./data'):
+    def __init__(self, dataset_root='./data', transform = transforms.ToTensor()):
         self.dataset_root = dataset_root
+        self.transform = transform
 
-    def get_random_image(self):
-        cifar10 = datasets.CIFAR10(root=self.dataset_root, train=False, download=True, transform=transforms.ToTensor())
+    def get_random_image(self, resize_dim=(256, 256)) -> Image:
+        cifar10 = datasets.CIFAR10(root=self.dataset_root, train=False, download=True, transform= self.transform)
         random_idx = random.randint(0, len(cifar10) - 1)
         image, _ = cifar10[random_idx]
-        image = transforms.ToPILImage()(image)
+        image= transforms.ToPILImage()(image) #bak buraya
+        image = image.resize(resize_dim, )
         return image
 
-# GradioApp: Sets up the Gradio interface
+# Interface
 class GradioApp:
-    def __init__(self, image_classifier, image_provider, model_list):
+    def __init__(self, image_classifier : ImageClassifier, image_provider : CIFAR10ImageProvider, model_list : list):
         self.image_classifier = image_classifier
         self.image_provider = image_provider
         self.model_list = model_list
 
     def launch(self):
-        with gr.Blocks() as demo:
-            with gr.Tabs():
-                with gr.TabItem("Upload Image"):
-                    with gr.Row():
-                        with gr.Column():
-                            upload_image = gr.Image(type='pil', label="Upload Image")
-                            model_dropdown_upload = gr.Dropdown(self.model_list, label="Select Model")
-                            classify_button_upload = gr.Button("Classify")
-                        with gr.Column():
-                            output_label_upload = gr.Label(num_top_classes=5)
+        with Blocks() as demo:
+            with Tabs():
+                with TabItem("Upload Image"):
+                    with Row():
+                        with Column():
+                            upload_image = Image(type='pil', label="Upload Image")
+                            model_dropdown_upload = Dropdown(self.model_list, label="Select Model")
+                            classify_button_upload = Button("Classify")
+                        with Column():
+                            output_label_upload = Label(num_top_classes=5)
                     classify_button_upload.click(self.image_classifier.classify, inputs=[upload_image, model_dropdown_upload], outputs=output_label_upload)
 
-                with gr.TabItem("Generate Random Image"):
-                    with gr.Row():
-                        with gr.Column():
-                            generate_button = gr.Button("Generate Random Image")
-                            random_image_output = gr.Image(type='pil', label="Random CIFAR-10 Image")
-                        with gr.Column():
-                            model_dropdown_random = gr.Dropdown(self.model_list, label="Select Model")
-                            classify_button_random = gr.Button("Classify")
-                            output_label_random = gr.Label(num_top_classes=5)
+                with TabItem("Generate Random Image"):
+                    with Row():
+                        with Column():
+                            generate_button = Button("Generate Random Image")
+                            random_image_output = Image(type='pil', label="Random CIFAR-10 Image")
+                        with Column():
+                            model_dropdown_random = Dropdown(self.model_list, label="Select Model")
+                            classify_button_random = Button("Classify")
+                            output_label_random = Label(num_top_classes=5)
                     generate_button.click(self.image_provider.get_random_image, inputs=[], outputs=random_image_output)
                     classify_button_random.click(self.image_classifier.classify, inputs=[random_image_output, model_dropdown_random], outputs=output_label_random)
 
         demo.launch()
 
-# Main Execution
+# Main
 if __name__ == "__main__":
-    # Define available models
-    image_prediction_models = {
-        'resnet': models.resnet50,
-        'alexnet': models.alexnet,
-        'vgg': models.vgg16,
-        'squeezenet': models.squeezenet1_0,
-        'densenet': models.densenet161,
-        'inception': models.inception_v3,
-        'googlenet': models.googlenet,
-        'shufflenet': models.shufflenet_v2_x1_0,
-        'mobilenet': models.mobilenet_v2,
-        'resnext': models.resnext50_32x4d,
-        'wide_resnet': models.wide_resnet50_2,
-        'mnasnet': models.mnasnet1_0,
-        'efficientnet': models.efficientnet_b0,
-        'regnet': models.regnet_y_400mf,
-        'vit': models.vit_b_16,
-        'convnext': models.convnext_tiny
-    }
-
-    # Initialize components
-    model_loader = ModelLoader(image_prediction_models)
+    # Initialize
+    model_loader = ModelLoader(IMAGE_PREDICTION_MODELS)
     preprocessor = Preprocessor()
     response = requests.get("https://git.io/JJkYN")
     labels = response.text.split("\n")
     postprocessor = Postprocessor(labels)
     image_classifier = ImageClassifier(model_loader, preprocessor, postprocessor)
     image_provider = CIFAR10ImageProvider()
-
     model_list = model_loader.get_model_names()
 
-    # Launch Gradio app
+    # Launch
     app = GradioApp(image_classifier, image_provider, model_list)
-    app.launch()
+    app.launch()