Spaces:
Running
Running
| import mediapipe as mp | |
| from utils import read_n_resize | |
| mp_objectron = mp.solutions.objectron | |
| mp_drawing = mp.solutions.drawing_utils | |
| def detect_in_3d(image, model_name='Chair'): | |
| with mp_objectron.Objectron( | |
| static_image_mode=True, | |
| max_num_objects=5, | |
| min_detection_confidence=0.35, | |
| model_name=model_name | |
| ) as objectron: | |
| results = objectron.process(read_n_resize(image, read=False)) | |
| return results | |
| def draw_3d(results, image): | |
| annotated_image = image.copy() | |
| for detected_object in results.detected_objects: | |
| mp_drawing.draw_landmarks( | |
| annotated_image, | |
| detected_object.landmarks_2d, | |
| mp_objectron.BOX_CONNECTIONS | |
| ) | |
| mp_drawing.draw_axis( | |
| annotated_image, | |
| detected_object.rotation, | |
| detected_object.translation | |
| ) | |
| return annotated_image | |
| def mp_objectron_fn(image, min_detect_conf=0.5): | |
| for model_name in ['Chair', 'Shoe', 'Cup', 'Camera']: | |
| results = detect_in_3d(image, model_name=model_name) | |
| if results.detected_objects: | |
| annotated_image = draw_3d(results, image) | |
| return annotated_image |