File size: 1,247 Bytes
9837266 49f68da 9837266 0d49fc7 9837266 |
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 |
import torch
from transformers import ViTForImageClassification, ViTFeatureExtractor
import gradio as gr
from PIL import Image
# Check if GPU is available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load pre-trained ViT model from Hugging Face
model = ViTForImageClassification.from_pretrained('Dhahlan2000/banana_ripeness_level_detection', num_labels=20, ignore_mismatched_sizes=True)
model.to(device)
model.eval()
# Load ViT feature extractor
feature_extractor = ViTFeatureExtractor.from_pretrained('Dhahlan2000/banana_ripeness_level_detection')
# Class labels
predicted_classes = ['Overripe', 'ripe', 'rotten', 'unripe', 'new']
# Function for inference
def classify_fruit(image):
inputs = feature_extractor(images=image, return_tensors="pt").to(device)
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
predicted_class_idx = logits.argmax(-1).item()
return predicted_classes[predicted_class_idx]
# Gradio UI
demo = gr.Interface(
fn=classify_fruit,
inputs=gr.Image(type="pil"),
outputs=gr.Label(),
title="Fruit Ripeness Detection",
description="Upload an image of a fruit to determine whether it's fresh or rotten."
)
demo.launch()
|