File size: 1,485 Bytes
5be0841 |
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 |
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/ripeness_detection', num_labels=20)
model.to(device)
model.eval()
# Load ViT feature extractor
feature_extractor = ViTFeatureExtractor.from_pretrained('Dhahlan2000/ripeness_detection')
# Class labels
predicted_classes = [
'FreshApple', 'FreshBanana', 'FreshBellpepper', 'FreshCarrot', 'FreshCucumber', 'FreshMango', 'FreshOrange',
'FreshPotato', 'FreshStrawberry', 'FreshTomato', 'RottenApple', 'RottenBanana', 'RottenBellpepper', 'RottenCarrot',
'RottenCucumber', 'RottenMango', 'RottenOrange', 'RottenPotato', 'RottenStrawberry', 'RottenTomato']
# 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()
|