Spaces:
Running
Running
File size: 2,475 Bytes
192d452 c7d5210 185a1c6 2c0a33a 054f74a e5cf4e9 054f74a 185a1c6 2559312 5dcea5c 2559312 185a1c6 f22056c 673d632 185a1c6 673d632 b71e439 185a1c6 673d632 91d639b 673d632 192d452 dad1377 68dcd2b dad1377 657b87f dad1377 054f74a 4f2203a dad1377 |
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 |
import gradio as gr
import numpy as np
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForImageClassification
model_names = [
"0-ma/swin-geometric-shapes-tiny",
"0-ma/mobilenet-v2-geometric-shapes",
"0-ma/focalnet-geometric-shapes-tiny",
"0-ma/efficientnet-b2-geometric-shapes",
"0-ma/beit-geometric-shapes-base",
"0-ma/mit-b0-geometric-shapes",
"0-ma/vit-geometric-shapes-base",
"0-ma/resnet-geometric-shapes",
"0-ma/vit-geometric-shapes-tiny",
]
example_images = [
'example/1_None.jpg',
'example/2_Circle.jpg',
'example/3_Triangle.jpg',
'example/4_Square.jpg',
'example/5_Pentagone.jpg',
'example/6_Hexagone.jpg'
]
labels = [example.split("_")[1].split(".")[0] for example in example_images]
feature_extractors = {model_name: AutoImageProcessor.from_pretrained(model_name) for model_name in model_names}
classification_models = {model_name: AutoModelForImageClassification.from_pretrained(model_name) for model_name in model_names}
def predict(image, 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 and upload an image to classify geometric shapes."
with gr.Blocks() as demo:
gr.Markdown(f"# {title}")
gr.Markdown(description)
with gr.Row():
model_dropdown = gr.Dropdown(choices=model_names, label="Selected Model", value=model_names[0])
gr.Examples(
examples=example_images,
inputs=image_input,
outputs=output,
fn=lambda img: predict(img, model_dropdown.value),
cache_examples=True,
)
image_input = gr.Image(type="pil")
output = gr.Label()
image_input.change(fn=lambda img: predict(img, model_dropdown.value), inputs=[image_input], outputs=output)
model_dropdown.change(fn=lambda img, model: predict(img, model), inputs=[image_input, model_dropdown], outputs=output)
demo.launch() |