Video inference of Mediapipe's Objectron. Currently only supports, chairs, shoes, cups and cameras.

objectron_video.py

import mediapipe as mp
import cv2

mp_objectron = mp.solutions.objectron
mp_drawing = mp.solutions.drawing_utils

if __name__ == '__main__':
    objectron = mp_objectron.Objectron(static_image_mode=True,
                                       max_num_objects=5,
                                       min_detection_confidence=0.5,
                                       model_name='Shoe')
    # Edit model_name variable to swap between models: 
    # posible choices : 'Shoe', 'Chair', 'Cup', 'Camera'

    video = r"yourvideo.mp4"
    vcap = cv2.VideoCapture(video)

    while True:
        isSuccess, frame = vcap.read()
        if isSuccess:
            results = objectron.process(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
            if not results.detected_objects:
                print(f'No box landmarks detected.')
                continue
            annotated_image = frame.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)
            cv2.imshow("Output", annotated_image)
            k = cv2.waitKey(10)
            if k == 27 or k == ord('q'):
                break
        else:
            break
    vcap.release()  # Release the frames
    cv2.destroyAllWindows()  # Destroy all windows