Antoine101's picture
Update app.py
5d9886e verified
Raw
History Blame Contribute Delete
1.43 kB
import gradio as gr
from PIL import Image, ImageDraw
from transformers import pipeline
import numpy as np
def plot_results(image, results, threshold=0.7):
image = Image.fromarray(np.uint8(image))
draw = ImageDraw.Draw(image)
for result in results:
score = result["score"]
label = result["label"]
box = list(result["box"].values())
if score > threshold:
x, y, x2, y2 = tuple(box)
draw.rectangle((x, y, x2, y2), outline="red", width=1)
draw.text((x, y), label, fill="white")
draw.text(
(x + 0.5, y - 0.5),
text=str(score),
fill="green" if score > 0.7 else "red",
)
return image
def predict(image):
# make the object detection pipeline
obj_detector = pipeline(
"object-detection", model="Antoine101/detr-resnet-50-dc5-fashionpedia-finetuned"
)
results = obj_detector(image)
print("Results:")
print(results)
return plot_results(image, results)
title = "Are you fashion?"
description = """
DETR model finetuned on "detection-datasets/fashionpedia" for apparels detection.
"""
demo = gr.Interface(
fn=predict,
inputs=gr.Image(label="Input Image", type="pil"),
outputs="image",
examples=[["example1.jpg"], ["example2.jpg"], ["example3.jpg"], ["example4.jpg"]],
title=title,
description=description
)
demo.launch()