from labelbox import Client
from matplotlib import pyplot as plt
from IPython.display import clear_output
import numpy as np
import json
import ndjson
import requests
import cv2
from typing import Dict, Any
import os
from typing import List, Dict

# Pick a video project with completed labels
PROJECT_ID = ""

# Add your API key
API_KEY = ""
client = Client(api_key=API_KEY)
project = client.get_project(PROJECT_ID)

export_url = project.export_labels()
exports = requests.get(export_url).json()

# Grab the annotation URL
annotation_url = exports[0]["Label"]["frames"]
headers = {"Authorization": f"Bearer {API_KEY}"}
annotations = ndjson.loads(requests.get(annotation_url, headers=headers).text)
annotations = {annot["frameNumber"]: annot for annot in annotations}

# Store the video URL in a file in the current working directory
video_url = exports[0]["Labeled Data"]
with open("12.mp4", "wb") as file:
    file.write(requests.get(video_url).content)

vidcap = cv2.VideoCapture("12.mp4")
count = 0

while True:
    success, image = vidcap.read()

    if not success:
        break  # Exit the loop when there are no more frames

    image = image[:, :, ::-1]  # Convert BGR to RGB

    # Get the annotation data for the current frame
    annotation = annotations.get(count)

    if annotation:
        # Extract the mask URL from the annotation
        mask_url = annotation["objects"][0]["instanceURI"]

        # Download and process the mask
        mask_image = cv2.cvtColor(cv2.imread(mask_url), cv2.COLOR_BGR2GRAY)

        # Resize the mask image to match the frame's dimensions
        mask_image = cv2.resize(mask_image, (image.shape[1], image.shape[0]))

        # Create a binary mask where the mask region is non-zero
        binary_mask = np.zeros_like(mask_image)
        binary_mask[mask_image > 0] = 255

        # Combine the binary mask with the frame using a bitwise operation
        masked_frame = cv2.bitwise_and(image, image, mask=binary_mask)

        # Display the masked frame
        plt.imshow(masked_frame)
        plt.title(f"Frame {count}")
        plt.show()

    count += 1

vidcap.release()