Spaces:
Running
Running
File size: 2,676 Bytes
192d452 c7d5210 185a1c6 079d565 2c0a33a 079d565 e5cf4e9 079d565 185a1c6 2559312 5dcea5c 2559312 dde0d91 185a1c6 f22056c 673d632 185a1c6 673d632 b71e439 185a1c6 673d632 91d639b 673d632 192d452 dad1377 68dcd2b dad1377 4addd34 dde0d91 4addd34 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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
import gradio as gr
import numpy as np
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForImageClassification
import os
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'
# ]
example_dir = "./example"
example_images = []
for example_image in os.list_dir(example_dir):
example_images+= [os.path.join(example_dir,example_image)]
#labels = [example.split("_")[1].split(".")[0] for example in example_images]
labels = [
'None',
'Circle',
'Triangle',
'Square',
'Pentagone',
'Hexagone'
]
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):
if image is None:
return None
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)
model_dropdown = gr.Dropdown(choices=model_names, label="Select Model", value=model_names[0])
image_input = gr.Image(type="pil")
gr.Examples(
examples=example_images,
inputs=image_input,
label="Click on an example image to test",
)
output = gr.Label(label="Classification Result")
image_input.change(fn=predict, inputs=[image_input, model_dropdown], outputs=output)
model_dropdown.change(fn=predict, inputs=[image_input, model_dropdown], outputs=output)
demo.launch() |