Spaces:
Sleeping
Sleeping
File size: 2,020 Bytes
947eaf5 63487f7 fce1d0a 947eaf5 5b7678f fce1d0a 947eaf5 fce1d0a 947eaf5 fce1d0a b6dc81d fce1d0a 947eaf5 |
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 |
import gradio as gr
from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image
import torch
import numpy as np
import spaces
import logging
# Set up verbose logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)
# Load model and processor from the Hugging Face Hub
MODEL_REPO = "Rausda6/autotrain-uo2t1-gvgzu" # Replace with your actual model repo name
logger.debug(f"Loading model from: {MODEL_REPO}")
model = AutoModelForImageClassification.from_pretrained(MODEL_REPO)
processor = AutoImageProcessor.from_pretrained(MODEL_REPO)
labels = model.config.id2label
@spaces.GPU
def classify_image(img: Image.Image):
logger.debug("Received image for classification.")
try:
inputs = processor(images=img, return_tensors="pt")
logger.debug(f"Processed inputs: {inputs}")
with torch.no_grad():
outputs = model(**inputs)
logger.debug(f"Model outputs: {outputs}")
logits = outputs.logits
probs = torch.nn.functional.softmax(logits, dim=-1)[0]
logger.debug(f"Probabilities: {probs}")
# Build result dictionary with confidence values
probs_dict = {labels[i]: float(probs[i]) for i in range(len(probs))}
# Sort and format nicely
sorted_probs = sorted(probs_dict.items(), key=lambda x: x[1], reverse=True)
top_label, top_score = sorted_probs[0]
logger.debug(f"Top prediction: {top_label} with confidence {top_score:.2%}")
return top_label, dict(sorted_probs)
except Exception as e:
logger.exception("Error during classification")
raise e
# Gradio interface
demo = gr.Interface(
fn=classify_image,
inputs=gr.Image(type="pil"),
outputs=[gr.Label(label="Top Prediction"), gr.Label(num_top_classes=6, label="Class Probabilities")],
title="Image Classification with AutoTrain Model",
description="Upload a JPG image to classify it using the fine-tuned model."
)
demo.launch()
|