import keras
import keras_cv
import numpy as np
import tensorflow as tf

from playgrounds.load_video import frames_from_video_file

def main():
    pretrained_model = keras_cv.models.YOLOV8Detector.from_preset(
        "yolo_v8_m_pascalvoc", bounding_box_format="xywh"
    )
    print('Model loaded.')

    inference_resizing = keras_cv.layers.Resizing(
        640, 640, pad_to_aspect_ratio=True, bounding_box_format="xywh"
    )

    class_ids = [
        "Aeroplane", "Bicycle", "Bird", "Boat", "Bottle", "Bus", "Car", "Cat", "Chair", "Cow", "Dining Table",
        "Dog", "Horse", "Motorbike", "Person", "Potted Plant", "Sheep", "Sofa", "Train", "Tvmonitor", "Total",
    ]
    class_mapping = {i: c for (i, c) in enumerate(class_ids)}

    # raw = tf.io.read_file('assets/IMG_9528.gif')
    # video = tf.io.decode_gif(raw)
    video = frames_from_video_file('assets/dataset/Flying/2kNjmM8BnD0_230.0_238.0.mp4', 3, (640,640))
    image = video[0]
    image = (image*255).astype(np.uint8)
    file = tf.io.encode_png(image)
    tf.io.write_file('out/t.png', file)
    # image = keras.utils.load_img('assets/nick-morales-BwYcH78rcpI-unsplash.jpg')
    # image = np.array(image)

    image_batch = inference_resizing([image])

    y_pred = pretrained_model.predict(image_batch)
    classes = y_pred['classes']
    boxes = y_pred["boxes"]
    print(f'Classes: {classes}')
    print(f'Boxes: {boxes}')