animate.py

import os
import sys

import cv2
import mediapy
import numpy as np
from eval import interpolator, util
from huggingface_hub import snapshot_download
from image_tools.sizes import resize_and_crop
from moviepy.editor import CompositeVideoClip
from moviepy.editor import VideoFileClip as vfc
from PIL import Image

# My installs
os.system("git clone https://github.com/google-research/frame-interpolation")
sys.path.append("frame-interpolation")


# get key positions at which frame needs to be generated
def list_of_positions(num_contours, num_frames=100):
    positions = []
    for i in range(0, num_frames):
        positions.append(int(num_contours / num_frames * i))
    return positions


def contourfinder(image1, image2, text=None, num_frames=100):
    # Create two blank pages to write into
    # I just hardcoded 1024*1024 as the size, ideally this should be np.shape(image1)
    blank = np.zeros((2048, 2048, 3), dtype="uint8")
    blank2 = np.zeros((2048, 2048, 3), dtype="uint8")

    # Threshold and contours for image 1 and 2
    threshold = cv2.Canny(image=image1, threshold1=100, threshold2=200)
    contours, hierarchies = cv2.findContours(
        threshold, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE
    )

    threshold2 = cv2.Canny(image=image2, threshold1=100, threshold2=200)
    contours2, hierarchies2 = cv2.findContours(
        threshold2, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE
    )

    # Initialize three empty videos
    vid1 = cv2.VideoWriter(
        "vid1.mp4", cv2.VideoWriter_fourcc(*"mp4v"), 24, threshold.shape
    )
    vid2 = cv2.VideoWriter(
        "vid2.mp4", cv2.VideoWriter_fourcc(*"mp4v"), 24, threshold.shape
    )
    text_vid = cv2.VideoWriter(
        "text_vid.mp4", cv2.VideoWriter_fourcc(*"mp4v"), 10, threshold.shape
    )

    # Get positions
    positions = list_of_positions((len(contours)))
    frames = []

    # Loop over contours adding them to blank image then writing to video
    for i in range(0, len(contours)):
        cv2.drawContours(
            blank, contours=contours, contourIdx=i, color=(125, 200, 255), thickness=1
        )

        if i in positions:
            frames.append(blank)

            # Complile to video
            vid1.write(blank)

    vid1.release()
    clip1 = vfc("/home/user/app/vid1.mp4")
    # print(f"Type clip 1{type(clip1)}")

    positions = list_of_positions((len(contours2)))

    for i in range(0, len(contours2)):
        cv2.drawContours(
            blank2, contours=contours2, contourIdx=i, color=(125, 200, 255), thickness=1
        )
        if i in positions:
            frames.append(blank2)

            vid2.write(blank2)

    vid2.release()
    clip3 = vfc("/home/user/app/vid2.mp4")

    # Next is the text vid

    if text != None:
        # Reading an image in default mode
        image = np.zeros(original.shape, dtype="uint8")

        # font
        font = cv2.FONT_HERSHEY_COMPLEX

        # org
        org = (10, 400)

        # fontScale
        fontScale = 3

        # Blue color in BGR
        color = (186, 184, 108)

        # Line thickness of 2 px
        thickness = 4

        def text_frames(text, image, org):
            spacing = 55  # spacing between letters
            blink = image
            cv2.imwrite(f"blink.png", blink)
            for i in range(0, len(text) - 1):

                text_vid.write(blink)

                # Using cv2.putText() method
                image = cv2.putText(
                    image, text[i], org, font, fontScale, color, thickness, cv2.LINE_AA
                )

                # Take care of org spacing
                org = (org[0] + spacing, org[1])
                if text[i].isupper():
                    org = (org[0] + spacing + 1, org[1])
                    print(f"Upper {text[i]}")
                print(org)

                # Displaying the image
                cv2.imwrite(f"text_im{i}.png", image)

                # Complile to video
                text_vid.write(image)
            text_vid.release()

        text_frames(text, image, org)
    return clip1, clip3


def load_model(model_name):
    model = interpolator.Interpolator(snapshot_download(repo_id=model_name), None)
    return model


model_names = [
    "akhaliq/frame-interpolation-film-style",
    "NimaBoscarino/frame-interpolation_film_l1",
    "NimaBoscarino/frame_interpolation_film_vgg",
]

models = {model_name: load_model(model_name) for model_name in model_names}

ffmpeg_path = util.get_ffmpeg_path()
mediapy.set_ffmpeg(ffmpeg_path)


def resize(width, img):
    basewidth = width
    img = Image.open(img)
    wpercent = basewidth / float(img.size[0])
    hsize = int((float(img.size[1]) * float(wpercent)))
    img = img.resize((basewidth, hsize), Image.ANTIALIAS)
    return img


def resize_img(img1, img2):
    img_target_size = Image.open(img1)
    img_to_resize = resize_and_crop(
        img2,
        (
            img_target_size.size[0],
            img_target_size.size[1],
        ),  # set width and height to match img1
        crop_origin="middle",
    )
    img_to_resize.save("resized_img2.png")


def predict(frame1, frame2, times_to_interpolate, model_name):
    model = models[model_name]

    img1 = cv2.imread(frame1)
    img2 = cv2.imread(frame2)

    frame1 = resize(256, frame1)
    frame2 = resize(256, frame2)

    frame1.save("test1.png")
    frame2.save("test2.png")

    resize_img("test1.png", "test2.png")
    input_frames = ["test1.png", "resized_img2.png"]

    frames = list(
        util.interpolate_recursively_from_files(
            input_frames, times_to_interpolate, model
        )
    )

    mediapy.write_video("out.mp4", frames, fps=5)

    print(f"TYPES....{type(img1)},{type(img2)} SHAPES{img1.shape} Img {img1}")
    clip1, clip3 = contourfinder(img1, img2)  # has a third text option

    # Use open CV and moviepy code
    # So we move from open CV video 1 to out.mp4 to open CV video2
    clip1 = clip1
    clip2 = vfc("out.mp4").resize(8).set_start(clip1.duration - 0.5).crossfadein(2)
    clip3 = clip3.set_start((clip1.duration - 0.5) + (clip2.duration)).crossfadein(2)

    new_clip = CompositeVideoClip([clip1, clip2, clip3])
    new_clip.write_videofile("out.mp4")

    return "out.mp4"