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  # Import List and Dict types from the typing module

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

#Add your api key
API_KEY = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VySWQiOiJjbG83NTZmMGQwOWNuMDcyYjlxemI0Ym9kIiwib3JnYW5pemF0aW9uSWQiOiJjbG83NTZlenEwOWNtMDcyYmVwOHoyODBjIiwiYXBpS2V5SWQiOiJjbG9wcmx3dTMwM3Z0MDc0czlwZmY0ZHJxIiwic2VjcmV0IjoiNWQ5OTUxMGFiMjI5MWQ5YzNkZTc2MTZhNzE5YTA3MDQiLCJpYXQiOjE2OTk0NDgyNDgsImV4cCI6MjMzMDYwMDI0OH0.nHr_hPI1PVtBapLwya16aaMefzQ3JV9U3pMbTWPNsds"
client = Client(api_key = API_KEY)
project = client.get_project(PROJECT_ID)

export_url = project.export_labels()

print(export_url)

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

type(exports)

exports[0]

#View some specific fields of the export
print("LAbel ID:", exports[0]["DataRow ID"])
print("Created By:", exports[0]["Created By"])

#View the video in your browser by clicking on the link produced
video_url = exports[0]["Labeled Data"]
print(video_url)

#Grab the annotation url
annotation_url = exports[0]["Label"]["frames"]
print(annotation_url) #Gives a link and downloads a json file https://api.labelbox.com/v1/frames/clonitedm0qkl070ncjmr5xgu

#Provide the appropriate authorization to view the labeled frames
headers = {"Authorization": f"Bearer {API_KEY}"}
annotations = ndjson.loads(requests.get(annotation_url, headers = headers).text)

#Make the data easy to lookup by assigning each frame its frame number
annotations = {annot["frameNumber"]: annot for annot in annotations}

#Grab the first frame and print the contents
first_frame = annotations[1]
print(json.dumps(first_frame, indent = 2))

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

vidcap = cv2.VideoCapture("12.mp4")
success, image = vidcap.read()
image = image[:, :, ::-1]
#Note that frameNumber 1 in the annotation is frame index 0
count = 1
while success and count < 20:
    plt.figure(1)
    plt.imshow(image)
    plt.title("frameNumber" + str(count))
    plt.pause(0.25)
    plt.clf()
    success, image = vidcap.read()
    count += 1
    if success and count < 20:
        clear_output(wait = True)
    image = image[:, :, ::-1]

# Helper function to visualize the segmentation masks
def visualize_bbox(image: np.ndarray, tool: Dict[str, Any]) -> np.ndarray:
    start = (int(tool["bbox"]["left"]), int(tool["bbox"]["top"]))
    end = (int(tool["bbox"]["left"] + tool["bbox"]["width"]),
           int(tool["bbox"]["top"] + tool["bbox"]["height"]))
    return cv2.rectangle(image, start, end, (255, 0, 0), 5)

#Overlay Annotations

vidcap = cv2.VideoCapture("12.mp4")
success, image = vidcap.read()
image = image[:, :, ::-1]
# Note that frameNumber 1 in the annotation is frame index 0
count = 1

while success and count < 20:
    annotation = annotations.get(count)
    if annotation is not None:
        for tool in annotation["objects"]:
            if "bbox" in tools:
                image = visualize_bbox(image.astype(np.uint8), annotation.get("objects", []))

    plt.figure(1)
    plt.imshow(image)
    plt.title("frameNumber " + str(count))
    plt.pause(0.25)
    plt.clf()

    success, image = vidcap.read()
    count += 1
    if success and count < 20:
        clear_output(wait=True)
    image = image[:, :, ::-1]