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