wannakorn/Insect_Classification_Using_SwinTransformer

1import gradio as gr2import torch3import torch.nn as nn4import time5from torchvision import models, transforms6from PIL import Image7 8 9model = models.swin_t(pretrained=False)10num_classes = 1111model.head = torch.nn.Linear(model.head.in_features, num_classes)12model.load_state_dict(torch.load("SwinTransformer.pth", map_location=torch.device('cpu')))13model.eval()14 15labels = [16    "Batocera rufomaculata",17    "Dermaptera",18    "Unspecified",19    "Xylotrechus quadripes",20    "Oryctes rhinoceros",21    "Dorysthenes buqueti",22    "Aulacophora indica",23    "Micraspis discolor",24    "Coccinella transversalis",25    "Menochilus sexmaculatus",26    "Mantodea"27]28 29def predict(image):30    transform = transforms.Compose([31        transforms.Resize((256, 256)),32        transforms.ToTensor(),33    ])34    image_tensor = transform(image).unsqueeze(0)35 36    start = time.time()37    with torch.no_grad():38        outputs = model(image_tensor)39        probs = torch.nn.functional.softmax(outputs[0], dim=0)40    end = time.time()41 42    result = {labels[i]: float(probs[i]) for i in range(len(labels))}43    time1 = f"Time: {(end - start):.4f} seconds"44    return result, time145 46with gr.Blocks() as demo:47    gr.Markdown("Insect Classifier")48    gr.Markdown("Upload an image to see the predicted class.")49 50    with gr.Row():51        with gr.Column():52            input_image = gr.Image(type="pil")53            predict_btn = gr.Button("Predict")54 55        with gr.Column():56            prediction_output = gr.Label(label="Prediction Results")57            time_output = gr.Markdown()58 59    predict_btn.click(60        fn=predict,61        inputs=input_image,62        outputs=[prediction_output, time_output]63    )64 65demo.launch()