| import gradio as gr |
| import os |
| import torch |
| import torch.nn as nn |
| import torchvision.transforms as transforms |
| import torchvision.models as models |
| from torchvision.models import ResNet18_Weights |
| from PIL import Image |
| from timeit import default_timer as timer |
| from typing import Tuple, Dict |
|
|
| |
| class_names = ['Cabbage', 'Carrot', 'Cucumber', 'Tomato'] |
|
|
| |
| weights = ResNet18_Weights.DEFAULT |
| transform = weights.transforms() |
|
|
| |
| def get_resnet18(): |
| model = models.resnet18(weights=weights) |
| num_ftrs = model.fc.in_features |
| model.fc = nn.Linear(num_ftrs, len(class_names)) |
| return model |
|
|
| model = get_resnet18() |
| model.load_state_dict(torch.load("resnet_model.pth", map_location=torch.device("cpu"), weights_only=False)) |
| model.eval() |
|
|
| |
| def predict(img: Image.Image) -> Tuple[Dict[str, float], float]: |
| start = timer() |
| img_tensor = transform(img).unsqueeze(0) |
|
|
| with torch.inference_mode(): |
| output = model(img_tensor) |
| probs = torch.softmax(output, dim=1)[0] |
|
|
| result = {class_names[i]: float(probs[i]) for i in range(len(class_names))} |
| elapsed = round(timer() - start, 5) |
| return result, elapsed |
|
|
| |
| with gr.Blocks() as demo: |
| gr.HTML(""" |
| <div style="background-image: url('https://images.unsplash.com/photo-1501004318641-b39e6451bec6'); |
| background-size: cover; |
| padding: 25px; |
| border-radius: 10px; |
| color: white; |
| font-family: 'Poppins', sans-serif;"> |
| <h1 style='text-align: center;'>🌾 Vegetable Classifier with ResNet18 🌾</h1> |
| <h3 style='text-align: center;'> |
| Created by <a href='https://github.com/ExplorerGumel' target='_blank' style='color: #f0e68c;'>@ExplorerGumel</a> |
| </h3> |
| </div> |
| """) |
|
|
| with gr.Row(): |
| with gr.Column(): |
| image_input = gr.Image(type="pil", label="📸 Upload Vegetable Image") |
| predict_button = gr.Button("🔍 Predict") |
|
|
| |
| examples = [] |
| if os.path.exists("examples"): |
| examples = [["examples/" + fname] for fname in os.listdir("examples") if fname.lower().endswith((".png", ".jpg", ".jpeg"))] |
| gr.Examples(examples=examples, inputs=[image_input]) |
|
|
| with gr.Column(): |
| output_label = gr.Label(num_top_classes=4, label="🌽 Predictions") |
| output_time = gr.Number(label="⏱️ Prediction Time (seconds)") |
|
|
| predict_button.click(fn=predict, inputs=image_input, outputs=[output_label, output_time]) |
|
|
| |
| demo.launch() |
|
|
|
|
|
|
