File size: 1,197 Bytes
cb531f9 |
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 |
import gradio as gr
from transformers import AutoProcessor, AutoModelForImageClassification
import torch
from PIL import Image
model_name = 'e1010101/vit-384-tongue-image'
processor = AutoProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)
def classify_image(image):
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
# Apply sigmoid for multi-label classification
probs = torch.sigmoid(logits)[0].numpy()
# Get label names
labels = model.config.id2label.values()
# Create a dictionary of labels and probabilities
result = {label: float(prob) for label, prob in zip(labels, probs)}
# Sort results by probability
result = dict(sorted(result.items(), key=lambda item: item[1], reverse=True))
return result
interface = gr.Interface(
fn=classify_image,
inputs=gr.inputs.Image(type="pil"),
outputs=gr.outputs.Label(num_top_classes=None),
title="Multi-Label Image Classification",
description="Upload an image to get classification results."
)
if __name__ == "__main__":
interface.launch() |