innat/Google-MediaPipe

1import mediapipe as mp2from utils import read_n_resize3 4mp_objectron = mp.solutions.objectron5mp_drawing = mp.solutions.drawing_utils6 7def detect_in_3d(image, model_name='Chair'):8    with mp_objectron.Objectron(9        static_image_mode=True,10        max_num_objects=5,11        min_detection_confidence=0.35,12        model_name=model_name13    ) as objectron:14        results = objectron.process(read_n_resize(image, read=False))15        return results16 17def draw_3d(results, image):18    annotated_image = image.copy()19    for detected_object in results.detected_objects:20        mp_drawing.draw_landmarks(21            annotated_image, 22            detected_object.landmarks_2d, 23            mp_objectron.BOX_CONNECTIONS24        )25        mp_drawing.draw_axis(26            annotated_image, 27            detected_object.rotation, 28            detected_object.translation29        )30    return annotated_image31 32 33def mp_objectron_fn(image, min_detect_conf=0.5):34    for model_name in ['Chair', 'Shoe', 'Cup', 'Camera']:35        results = detect_in_3d(image, model_name=model_name)36        if results.detected_objects:37            annotated_image = draw_3d(results, image)38            return annotated_image