File size: 2,787 Bytes
9ae7b7e 35ad53f 9ae7b7e 62e392d 779da14 c1ef9cb 4c89202 35ad53f 62e392d c1ef9cb 35ad53f 62e392d c1ef9cb 62e392d 35ad53f 62e392d c1ef9cb 62e392d c1ef9cb 62e392d c1ef9cb 779da14 62e392d 779da14 62e392d 779da14 35ad53f 62e392d 35ad53f 62e392d 35ad53f 62e392d c1ef9cb c3b4ea5 c1ef9cb 779da14 35ad53f c3b4ea5 779da14 c3b4ea5 62e392d c1ef9cb 62e392d 779da14 35ad53f c3b4ea5 c1ef9cb 62e392d 779da14 62e392d 35ad53f c3b4ea5 | 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 | 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/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])
# 6. Launch the app
demo.launch()
|