Spaces:
Running
Running
File size: 2,125 Bytes
e921d65 955daea e921d65 955daea e921d65 955daea |
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 |
import requests
from io import BytesIO
import numpy as np
from PIL import Image
import yolov5
from yolov5.utils.plots import Annotator, colors
import gradio as gr
from huggingface_hub import get_token
import time
def load_model(model_path, img_size=640):
model = yolov5.load(model_path, hf_token=get_token())
model.img_size = img_size # add img_size attribute
return model
def load_image_from_url(url):
if not url: # empty or None
return gr.Image(interactive=True)
try:
response = requests.get(url, timeout=5)
image = Image.open(BytesIO(response.content))
except Exception as e:
raise gr.Error("Unable to load image from URL") from e
return image.convert("RGB")
def inference(model, image):
results = model(image, size=model.img_size)
annotator = Annotator(np.asarray(image))
for *box, _, cls in reversed(results.pred[0]):
# label = f'{model.names[int(cls)]} {conf:.2f}'
# print(f'{cls} {conf:.2f} {box}')
annotator.box_label(box, "", color=colors(cls, True))
return annotator.im
def count_flagged_images(dataset_name, trials=10):
headers = {"Authorization": f"Bearer {get_token()}"}
API_URL = f"https://datasets-server.huggingface.co/size?dataset={dataset_name}"
def query():
response = requests.get(API_URL, headers=headers, timeout=5)
return response.json()
for i in range(trials):
try:
data = query()
if "error" not in data and data["size"]["dataset"]["num_rows"] > 0:
print(f"[{i+1}/{trials}] {data}")
return data["size"]["dataset"]["num_rows"]
except Exception:
pass
print(f"[{i+1}/{trials}] {data}")
time.sleep(5)
return 0
def load_badges(dataset_name, trials=10):
n = count_flagged_images(dataset_name, trials)
return f"""
<p style="display: flex">
<img alt="" src="https://img.shields.io/badge/SEA.AI-beta-blue">
<img alt="" src="https://img.shields.io/badge/%F0%9F%96%BC%EF%B8%8F-{n}-green">
</p>
"""
|