Spaces:
Sleeping
Sleeping
File size: 2,751 Bytes
192d452 c7d5210 185a1c6 c7d5210 185a1c6 c7d5210 185a1c6 479feb3 185a1c6 5dcea5c 185a1c6 c7d5210 91d639b 185a1c6 f22056c 673d632 185a1c6 673d632 b71e439 185a1c6 673d632 91d639b 673d632 185a1c6 192d452 185a1c6 68dcd2b 185a1c6 192d452 185a1c6 |
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 81 82 |
import gradio as gr
import numpy as np
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForImageClassification
import requests
labels = [
"None",
"Circle",
"Triangle",
"Square",
"Pentagon",
"Hexagon"
]
# Available models for the dropdown
models = {
"0-ma/swin-geometric-shapes-tiny": "0-ma/swin-geometric-shapes-tiny",
"0-ma/mobilenet-v2-geometric-shapes": "0-ma/mobilenet-v2-geometric-shapes",
"0-ma/focalnet-geometric-shapes-tiny": "0-ma/focalnet-geometric-shapes-tiny" ,
"0-ma/efficientnet-b2-geometric-shapes":"0-ma/efficientnet-b2-geometric-shapes",
"0-ma/beit-geometric-shapes-base":"0-ma/beit-geometric-shapes-base",
"0-ma/mit-b0-geometric-shapes":"0-ma/mit-b0-geometric-shapes",
"0-ma/vit-geometric-shapes-base":"0-ma/vit-geometric-shapes-base",
"0-ma/resnet-geometric-shapes":"0-ma/resnet-geometric-shapes",
"0-ma/vit-geometric-shapes-tiny":"0-ma/vit-geometric-shapes-tiny",
}
# Load the default model
#feature_extractor = AutoImageProcessor.from_pretrained(models["Tiny Model"])
#model = AutoModelForImageClassification.from_pretrained(models["Tiny Model"])
feature_extractors = { model_name : AutoImageProcessor.from_pretrained(models[model_name]) for model_name in models}
classification_models = { model_name : AutoModelForImageClassification.from_pretrained(models[model_name]) for model_name in models}
def predict(image, selected_model):
# Load the selected model
# feature_extractor = AutoImageProcessor.from_pretrained(models[selected_model])
# model = AutoModelForImageClassification.from_pretrained(models[selected_model])
feature_extractor = feature_extractors[selected_model]
model = classification_models[selected_model]
inputs = feature_extractor(images=[image], return_tensors="pt")
logits = model(**inputs)['logits'].cpu().detach().numpy()[0]
logits_positive = logits
logits_positive[logits < 0] = 0
logits_positive = logits_positive/np.sum(logits_positive)
confidences = {}
for i in range(len(labels)):
if logits[i] > 0:
confidences[labels[i]] = float(logits_positive[i])
return confidences
title = "Geometric Shape Classifier"
description = "Select a model to classify geometric shapes."
examples = [
'example/1_None.jpg',
'example/2_Circle.jpg',
'example/3_Triangle.jpg',
'example/4_Square.jpg',
'example/5_Pentagone.jpg',
'example/6_Hexagone.jpg'
]
# Adding a dropdown for model selection
gr.Interface(
fn=predict,
inputs=[gr.Image(type="pil"), gr.Dropdown(list(models.keys()), label="Select Model")],
outputs=gr.Label(),
title=title,
description=description,
examples=examples
).launch()
|