Spaces:
Running
Running
| import time | |
| import requests | |
| from io import BytesIO | |
| from dataclasses import dataclass | |
| import numpy as np | |
| import pandas as pd | |
| from PIL import Image | |
| import yolov5 | |
| from yolov5.utils.plots import Annotator, colors | |
| import gradio as gr | |
| from huggingface_hub import get_token | |
| def load_model(model_path, img_size=640): | |
| """Load model from HuggingFace Hub.""" | |
| 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): | |
| """Load image from 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): | |
| """Run inference on image and return annotated 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 load_badges(n): | |
| """Load badges.""" | |
| 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> | |
| """ | |
| class FlaggedCounter: | |
| """Count flagged images in dataset.""" | |
| dataset_name: str | |
| headers: dict = None | |
| def __post_init__(self): | |
| self.API_URL = ( | |
| f"https://datasets-server.huggingface.co/size?dataset={self.dataset_name}" | |
| ) | |
| self.trials = 10 | |
| if self.headers is None: | |
| self.headers = {"Authorization": f"Bearer {get_token()}"} | |
| def query(self): | |
| """Query API.""" | |
| response = requests.get(self.API_URL, headers=self.headers, timeout=5) | |
| return response.json() | |
| def from_query(self, data): | |
| """Count flagged images via API. Might be slow.""" | |
| for i in range(self.trials): | |
| try: | |
| data = self.query() | |
| if "error" not in data and data["size"]["dataset"]["num_rows"] > 0: | |
| print(f"[{i+1}/{self.trials}] {data}") | |
| return data["size"]["dataset"]["num_rows"] | |
| except Exception: | |
| pass | |
| print(f"[{i+1}/{self.trials}] {data}") | |
| time.sleep(5) | |
| return 0 | |
| def from_csv(self): | |
| """Count flagged images from CSV. Fast but relies on local files.""" | |
| dataset_name = self.dataset_name.split("/")[-1] | |
| df = pd.read_csv(f"./flagged/{dataset_name}/data.csv") | |
| return len(df) | |
| def count(self): | |
| """Count flagged images.""" | |
| try: | |
| return self.from_csv() | |
| except FileNotFoundError: | |
| return self.from_query(self.query()) | |