Spaces:
Runtime error
Runtime error
| from typing import Any, List | |
| import numpy as np | |
| import gradio as gr | |
| from PIL import Image, ImageDraw | |
| import yolo_detect as yod | |
| import cv2 | |
| import fiftyone as fo | |
| from flagging import FlaggingCallback, SimpleCSVLogger | |
| class NewLogger(FlaggingCallback): | |
| def __init__(self): | |
| self.flag_data = None | |
| self.flag_option = None | |
| self.flag_index = None | |
| super().__init__() | |
| def flag( | |
| self, | |
| flag_data: List[Any], | |
| flag_option= None, | |
| flag_index = None, | |
| username = None, | |
| ) -> int: | |
| self.flag_data = flag_data | |
| self.flag_option = flag_option | |
| self.flag_index = flag_index | |
| if flag_option == "Bad": | |
| print(flag_option) | |
| # log_filepath = Path(flagging_dir) / "log.csv" | |
| # csv_data = [] | |
| # for component, sample in zip(self.components, flag_data): | |
| # save_dir = Path(flagging_dir) / utils.strip_invalid_filename_characters( | |
| # component.label or "" | |
| # ) | |
| return 2 | |
| def draw_boxes(input_img, iou_threshold, confidence_threshdol): | |
| # Convert the input to a PIL Image | |
| img = cv2.resize(input_img, (640, 640)) | |
| img = Image.fromarray(img) | |
| draw = ImageDraw.Draw(img) | |
| # Example bounding boxes: (x1, y1, x2, y2) | |
| boxes = yod.identifications(input_img, iou_threshold, confidence_threshdol) | |
| # Draw rectangles on the image | |
| for box in boxes: | |
| bbox = box[:4] | |
| draw.rectangle(bbox, outline="red", width=2) | |
| draw.text((bbox[0], bbox[1] - 16), f"{box[5]} -> {str(round(box[4], 2))}", fill="red") | |
| # Convert back to array | |
| return np.array(img) | |
| # Define the Gradio Interface with a title | |
| title = "Vic and the boyzzzzzz" | |
| # Define a slider | |
| iou_slider = gr.Slider(minimum=0, maximum=1, step=0.01, label="IoU Threshold", value=0.0) | |
| conf_slider = gr.Slider(minimum=0, maximum=1, step=0.01, label="Confidence Threshold", value=0.0) | |
| text_box = gr.Textbox(label="Write a description for bad behavior") | |
| demo = gr.Interface(fn=draw_boxes, inputs=[gr.Image(), iou_slider, conf_slider], outputs=gr.Image(), title=title, flagging_options=["Good", "Bad"], allow_flagging="manual", flagging_callback=SimpleCSVLogger()) | |
| # Launch the app | |
| demo.launch() | |