BrunoMelicio
bug fix color
3e1778b
raw
history blame
1.88 kB
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()