Spaces:
Sleeping
Sleeping
| 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 | |
| # 1. Class names (order must match your training) | |
| class_names = ['Cabbage', 'Carrot', 'Cucumber', 'Tomato'] | |
| # 2. Load default preprocessing from pretrained weights | |
| weights = ResNet18_Weights.DEFAULT | |
| transform = weights.transforms() | |
| # 3. Define and load the model | |
| 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() | |
| # 4. Prediction function | |
| 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 | |
| # 5. Build the Gradio interface | |
| 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/hassan1303' target='_blank' style='color: #f0e68c;'>@hassan1303</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]) | |
| # 6. Launch the app | |
| demo.launch() | |