app.py

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
from gradio import Blocks, Tabs, TabItem, Row, Column, Image, Dropdown, Button, Label

# Predefined models available in torchvision
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
}

# Load a pretrained model from torchvision
class ModelLoader:
    def __init__(self, model_dict : dict):
        self.model_dict = model_dict

    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]
            model = model_class(pretrained=True)
            return model
        else:
            raise ValueError(f"Model {model_name} is not available for image prediction in torchvision.models")

    def get_model_names(self) -> list:
        return [name.capitalize() for name in self.model_dict.keys()]

# Preprocessor: Prepares image for model input
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 : str) -> Compose:
        input_size = 224
        if model_name == 'inception':
            input_size = 299
        return transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(input_size),
            transforms.ToTensor(),
            self.normalize,
        ])

# Postprocessor: Processes model output
class Postprocessor:
    def __init__(self, labels : list):
        self.labels = labels

    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) -> 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 : ModelLoader, preprocessor: Preprocessor, postprocessor : Postprocessor):
        self.model_loader = model_loader
        self.preprocessor = preprocessor
        self.postprocessor = postprocessor

    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)

        if torch.cuda.is_available():
            input_batch = input_batch.to('cuda')
            model.to('cuda')

        model.eval()
        with torch.no_grad():
            output : T = model(input_batch)

        if selected_model.lower() == 'inception':
            return self.postprocessor.postprocess_inception(output)
        else:
            return self.postprocessor.postprocess_default(output)

# CIFAR10ImageProvider: Provides random images from CIFAR-10 dataset
class CIFAR10ImageProvider:
    def __init__(self, dataset_root='./data', transform = transforms.ToTensor()):
        self.dataset_root = dataset_root
        self.transform = transform

    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) #bak buraya
        image = image.resize(resize_dim, )
        return image

# Interface
class GradioApp:
    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 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 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
if __name__ == "__main__":
    # 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
    app = GradioApp(image_classifier, image_provider, model_list)
    app.launch()