Upload objectron3d.py

objectron3d.py

@@ -0,0 +1,38 @@
+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