from keras.models import load_model
from keras.preprocessing.image import img_to_array
from keras.preprocessing.image import load_img
from numpy import expand_dims
from matplotlib import pyplot
import cv2
from pathlib import Path 
import glob
import os
import re


fps = 30
input_dir = "predict_frames/"
model_dir = 'save_models3/model_005600.h5'
generate_frames = 200

def atoi(text):
    return int(text) if text.isdigit() else text

def natural_keys(text):
    return [ atoi(c) for c in re.split(r'(\d+)', text) ]
    
def load_image(filename, size=(256,256)):
    # load image with the preferred size
    pixels = load_img(filename, target_size=size)
    # convert to numpy array
    pixels = img_to_array(pixels)
    # scale from [0,255] to [-1,1]
    pixels = (pixels - 127.5) / 127.5
    # reshape to 1 sample
    pixels = expand_dims(pixels, 0)
    return pixels

def loadModel():
    # load model
    global model
    model = load_model(model_dir)

def labeledToPhoto(path, outputPic):
    # load source image
    src_image = load_image(path)
    #print('Loaded', src_image.shape)
    
    # generate image from source
    gen_image = model.predict(src_image)
    # scale from [-1,1] to [0,1]
    gen_image = (gen_image + 1) / 2.0
    # write to disk
    filename = input_dir + outputPic
    img = cv2.convertScaleAbs(gen_image[0], alpha=(255.0))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    cv2.imwrite(filename, img)

def convert(image_folder, output_video):
    images = [img for img in os.listdir(image_folder) if img.endswith(".jpg")]
    frame = cv2.imread(os.path.join(image_folder, images[0]))
    height, width, layers = frame.shape
    
    video = cv2.VideoWriter(filename=output_video,  
        fourcc=cv2.VideoWriter_fourcc(*'DIVX'), 
        fps=fps,  
        frameSize=(width, height))
    images.sort(key=natural_keys)
    for image in images:
        print("image:", image)
        video.write(cv2.imread(os.path.join(image_folder, image)))
        #video.write(image)
    
    cv2.destroyAllWindows()
    video.release()

def predict_dir():
    # first frame = 0.jpg
    files = glob.glob (input_dir + "source/*.jpg")
    i = 1
    for myFile in files:
    #for i in range(generate_frames):
        #labeledToPhoto(input_dir + str(i)+".jpg", str(i+1)+".jpg")
        labeledToPhoto(myFile, str(i)+".jpg")
        i += 1
    convert(input_dir, input_dir+"/video.mp4")
    print("generated video video.mp4")

def predict_one_image():
    # first frame = 0.jpg
    i = 0
    for i in range(generate_frames):
        labeledToPhoto(input_dir + str(i)+".jpg", str(i+1)+".jpg")
    convert(input_dir, input_dir+"/video.mp4")
    print("generated video video.mp4")
if __name__ == '__main__':
    loadModel()
    predict_dir()
    #predict_one_image()