Spaces:
Sleeping
Sleeping
File size: 2,785 Bytes
82ec4aa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | 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()
|