import cv2
import tensorflow as tf
import numpy as np
from crop_hands_skeleton import crop_hands
from cvzone.HandTrackingModule import HandDetector

class_names = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 'del', 'nothing', 'space']



def segment_video(video, fps=5):
    real_fps = video.get(cv2.CAP_PROP_FPS)
    print(f"{real_fps=}")
    if real_fps < fps:
        raise Exception("Video FPS cannot be bigger than desired FPS!")

    n = real_fps / fps

    result = []
    i=0
    num = 0
    while True:
        ret, frame = video.read()
        if ret == False:
            break
        if i % n == 0:
            result.append(frame)
            num += 1
        i += 1

    return result, num

def save_frames(frames, dir):
    detector = HandDetector(maxHands=1, mode=True, detectionCon=0.7, minTrackCon=0.8)
    for i, frame in enumerate(frames):
        print(i)        
        cv2.imwrite(f"{dir}/frame{i}.jpg", crop_hands(frame, detector))


def classify(img, model):
    #img = cv2.resize(img, (224, 224))
    img = tf.keras.utils.img_to_array(img)
    img = np.expand_dims(img, axis = 0)
    return class_names[np.argmax(model.predict(img))]


def read_saved_frames(dir, n):
    result = []
    for i in range(n):
        img = tf.keras.utils.load_img(f"{dir}/frame{i}.jpg", target_size = [224, 224])
        result.append(img)
    return result


if __name__ == "__main__":
    video = cv2.VideoCapture("kamil_asl.mp4")
    model = tf.keras.models.load_model('model_pred/VGG16_sign_char_detection_model')

    frames, num = segment_video(video, 20)
    print(num)
    save_frames(frames, "frames")
    frames = read_saved_frames("frames", num)
    result = []
    for frame in frames:
        result.append(classify(frame, model))
    print(result)