# model_2.py
import os
import cv2
import numpy as np
import importlib.util
from PIL import Image
import gradio as gr
from common_detection import perform_detection

MODEL_DIR = 'model_2'
GRAPH_NAME = 'detect.tflite'
LABELMAP_NAME = 'labelmap.txt'

pkg = importlib.util.find_spec('tflite_runtime')
if pkg:
    from tflite_runtime.interpreter import Interpreter
else:
    from tensorflow.lite.python.interpreter import Interpreter

PATH_TO_CKPT = os.path.join(MODEL_DIR, GRAPH_NAME)
PATH_TO_LABELS = os.path.join(MODEL_DIR, LABELMAP_NAME)

with open(PATH_TO_LABELS, 'r') as f:
    labels = [line.strip() for line in f.readlines()]

if labels[0] == '???':
    del(labels[0])

interpreter = Interpreter(model_path=PATH_TO_CKPT)
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
height = input_details[0]['shape'][1]
width = input_details[0]['shape'][2]
floating_model = (input_details[0]['dtype'] == np.float32)

def detect_image(input_image):
    image = np.array(input_image)
    resized_image = cv2.resize(image, (640, 640))
    result_image = perform_detection(resized_image, interpreter, labels, input_details, output_details, height, width, floating_model)
    return Image.fromarray(result_image)

def detect_video(input_video):
    cap = cv2.VideoCapture(input_video)
    frames = []

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        resized_frame = cv2.resize(frame, (640, 640))
        result_frame = perform_detection(resized_frame, interpreter, labels, input_details, output_details, height, width, floating_model)
        frames.append(result_frame)

    cap.release()

    if not frames:
        raise ValueError("No frames were read from the video.")

    height, width, layers = frames[0].shape
    size = (width, height)
    output_video_path = "result_" + os.path.basename(input_video)
    out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), 15, size)

    for frame in frames:
        out.write(frame)

    out.release()

    return output_video_path