Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import mediapipe as mp | |
| import cv2 | |
| import numpy as np | |
| BaseOptions = mp.tasks.BaseOptions | |
| ObjectDetector = mp.tasks.vision.ObjectDetector | |
| ObjectDetectorOptions = mp.tasks.vision.ObjectDetectorOptions | |
| VisionRunningMode = mp.tasks.vision.RunningMode | |
| MARGIN = 10 # pixels | |
| ROW_SIZE = 10 # pixels | |
| FONT_SIZE = 1 | |
| FONT_THICKNESS = 1 | |
| TEXT_COLOR = (255, 0, 0) # red | |
| def visualize(image, detection_result) -> np.ndarray: | |
| for detection in detection_result.detections: | |
| # Draw bounding_box | |
| bbox = detection.bounding_box | |
| start_point = bbox.origin_x, bbox.origin_y | |
| end_point = bbox.origin_x + bbox.width, bbox.origin_y + bbox.height | |
| cv2.rectangle(image, start_point, end_point, TEXT_COLOR, 3) | |
| # Draw label and score | |
| category = detection.categories[0] | |
| category_name = category.category_name | |
| probability = round(category.score, 2) | |
| result_text = category_name + ' (' + str(probability) + ')' | |
| text_location = (MARGIN + bbox.origin_x, | |
| MARGIN + ROW_SIZE + bbox.origin_y) | |
| cv2.putText(image, result_text, text_location, cv2.FONT_HERSHEY_PLAIN, | |
| FONT_SIZE, TEXT_COLOR, FONT_THICKNESS) | |
| return image | |
| def analyze_image(image): | |
| model_path = "efficientdet_lite0.tflite" | |
| options = ObjectDetectorOptions( | |
| base_options=BaseOptions(model_asset_path=model_path), | |
| max_results=5, | |
| running_mode=VisionRunningMode.IMAGE) | |
| mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image) | |
| with ObjectDetector.create_from_options(options) as detector: | |
| detection_result = detector.detect(mp_image) | |
| image_copy = np.copy(mp_image.numpy_view()) | |
| annotated_image = visualize(image_copy, detection_result) | |
| return annotated_image | |
| img_in = gr.Image() | |
| #img_out = gr.Image() | |
| iface = gr.Interface(fn=analyze_image, inputs=img_in, outputs="image") | |
| iface.launch() | |