app.py

import os
import datetime
import shutil
import subprocess
import cv2
from PIL import Image
from moviepy.editor import *
from moviepy.video.io.ffmpeg_tools import ffmpeg_extract_subclip
from gradio_client import Client
import gradio as gr

matte_client = Client("https://fffiloni-video-matting-anything.hf.space/")

# execute a CLI command
def execute_command(command: str) -> None:
    subprocess.run(command, check=True)


def infer(video_frames, masks_frames, project_name):


    # Create the directory if it doesn't exist
    my_video_directory = f"{project_name}"
    if not os.path.exists(my_video_directory):
        os.makedirs(my_video_directory)
    else:
        # If the directory already exists, add a timestamp to the new directory name
        timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
        my_video_directory = f"{project_name}_{timestamp}"
        os.makedirs(my_video_directory)

    # Assuming 'images' is a list of image file paths
    for idx, image in enumerate(video_frames):
        # Get the base file name (without path) from the original location
        image_name = os.path.basename(image.name)
    
        # Construct the destination path in the working directory
        destination_path = os.path.join(my_video_directory, image_name)
    
        # Copy the image from the original location to the working directory
        shutil.copy(image.name, destination_path)
    
        # Print the image name and its corresponding save path
        print(f"Image {idx + 1}: {image_name} copied to {destination_path}")


    # Create the directory if it doesn't exist
    my_masks_directory = f"{project_name}_masks"
    if not os.path.exists(my_masks_directory):
        os.makedirs(my_masks_directory)

    # Assuming 'images' is a list of image file paths
    for idx, image in enumerate(masks_frames):
        # Get the base file name (without path) from the original location
        image_name = os.path.basename(image.name)
    
        # Construct the destination path in the working directory
        destination_path = os.path.join(my_masks_directory, image_name)
    
        # Copy the image from the original location to the working directory
        shutil.copy(image.name, destination_path)
    
        # Print the image name and its corresponding save path
        print(f"Image {idx + 1}: {image_name} copied to {destination_path}")

    #video_frames_folder = "inputs/object_removal/bmx-trees"
    #masks_folder = "inputs/object_removal/bmx-trees_mask"

    video_frames_folder = f"{my_video_directory}"
    masks_folder = f"{my_masks_directory}"
    
    # Create the "results" folder if it doesn't exist
    output_folder = "results"
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)


    #bmx_trees_folder = os.path.join(output_folder, "bmx-trees")
    
    command = [
      f"python", 
      f"inference_propainter.py",
      f"--video={video_frames_folder}",
      f"--mask={masks_folder}",
      f"--output={output_folder}"
    ]

    execute_command(command)

    # Get the list of files in the "results" folder
    result_files = os.listdir(output_folder)

    # Print the content of the "results" folder
    print(f"Contents of the {output_folder} folder:")
    for item in result_files:
        print(item)

    # List the content of the "bmx-trees" folder within "results"
    results_folder = os.path.join(output_folder, f"{project_name}")
    results_folder_content = [os.path.join(results_folder, item) for item in os.listdir(results_folder)]

    print(f"Contents of the {results_folder} folder:")
    for item in results_folder_content:
        print(item)

    return "done", results_folder_content[0], results_folder_content[1]


def get_frames(video_in, img_type):
    frames = []
    #resize the video
    # Cut the video to the first 5 seconds
    video_out = f"{img_type}_video_cut.mp4"
    
    ffmpeg_extract_subclip(video_in, t1=0, t2=3, targetname=video_out)
    # Now, proceed with resizing the cut video
    clip = VideoFileClip(video_out)
    
    #check fps
    
    if clip.fps > 24:
        print("vide rate is over 24, resetting to 24")
        clip_resized = clip.resize(height=512)
        clip_resized.write_videofile(f"{img_type}_video_resized.mp4", fps=24)
    else:
        print("video rate is OK")
        clip_resized = clip.resize(height=512)
        clip_resized.write_videofile(f"{img_type}_video_resized.mp4", fps=clip.fps)
    
    
    print("video resized to 512 height")
    
    # Opens the Video file with CV2
    cap= cv2.VideoCapture(f"{img_type}_video_resized.mp4")
    
    fps = cap.get(cv2.CAP_PROP_FPS)
    print("video fps: " + str(fps))
    i=0
    while(cap.isOpened()):
        ret, frame = cap.read()
        if ret == False:
            break
        if img_type == "source":
            filename = f'{i:05d}.jpg'
            cv2.imwrite(filename, frame)
            frames.append(filename)
        elif img_type == "mask":
            filename = f'{i:05d}.png'
            cv2.imwrite(filename, frame)
            frames.append(filename)
        i+=1
    
    cap.release()
    cv2.destroyAllWindows()
    print("broke the video into frames")
    
    return frames, fps

def get_matte(video_in, subject_to_remove):
    print("Trying to call video matting")
    result = matte_client.predict(
		f"{video_in}",	# str (filepath on your computer (or URL) of file) in 'parameter_4' Video component
		3,	# int | float (numeric value between 0 and 10) in 'Cut video at (s)' Slider component
		f"{subject_to_remove}",	# str  in 'Text prompt' Textbox component
		"",	# str  in 'Background prompt' Textbox component
		api_name="/go_matte"
    )
    print(result)

    return result[2]


def infer_auto(project_name, video_in, subject_to_remove):

    print(video_in)
    matte_video = get_matte(video_in, subject_to_remove)

    video_frames = get_frames(video_in, "source")
    print(video_frames[0])
 
    masks_frames = get_frames(matte_video, "mask")
    print(masks_frames[0])
    

    # Create the directory if it doesn't exist
    my_video_directory = f"{project_name}"
    if not os.path.exists(my_video_directory):
        os.makedirs(my_video_directory)
    else:
        # If the directory already exists, add a timestamp to the new directory name
        timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
        my_video_directory = f"{project_name}_{timestamp}"
        os.makedirs(my_video_directory)

    # Assuming 'images' is a list of image file paths
    for idx, image in enumerate(video_frames[0]):
        # Get the base file name (without path) from the original location
        image_name = os.path.basename(image)
    
        # Construct the destination path in the working directory
        destination_path = os.path.join(my_video_directory, image_name)
    
        # Copy the image from the original location to the working directory
        shutil.copy(image, destination_path)
    
        # Print the image name and its corresponding save path
        print(f"Image {idx + 1}: {image_name} copied to {destination_path}")


    # Create the directory if it doesn't exist
    my_masks_directory = f"{project_name}_masks"
    if not os.path.exists(my_masks_directory):
        os.makedirs(my_masks_directory)

    # Assuming 'images' is a list of image file paths
    for idx, image in enumerate(masks_frames[0]):
        # Get the base file name (without path) from the original location
        image_name = os.path.basename(image)
    
        # Construct the destination path in the working directory
        destination_path = os.path.join(my_masks_directory, image_name)
    
        # Copy the image from the original location to the working directory
        shutil.copy(image, destination_path)
    
        # Print the image name and its corresponding save path
        print(f"Image {idx + 1}: {image_name} copied to {destination_path}")

    #video_frames_folder = "inputs/object_removal/bmx-trees"
    #masks_folder = "inputs/object_removal/bmx-trees_mask"

    video_frames_folder = f"{my_video_directory}"
    masks_folder = f"{my_masks_directory}"
    
    # Create the "results" folder if it doesn't exist
    output_folder = "results"
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)


    #bmx_trees_folder = os.path.join(output_folder, "bmx-trees")

    # Convert the float fps to an integer
    needed_fps = int(video_frames[1])

        
    command_auto= [
      f"python", 
      f"inference_propainter.py",
      f"--video={video_frames_folder}",
      f"--mask={masks_folder}",
      f"--output={output_folder}",
      f"--save_fps={int(needed_fps)}"
    ]

    execute_command(command_auto)

    # Get the list of files in the "results" folder
    result_files = os.listdir(output_folder)

    # Print the content of the "results" folder
    print(f"Contents of the {output_folder} folder:")
    for item in result_files:
        print(item)

    # List the content of the "bmx-trees" folder within "results"
    results_folder = os.path.join(output_folder, f"{project_name}")
    results_folder_content = [os.path.join(results_folder, item) for item in os.listdir(results_folder)]

    print(f"Contents of the {results_folder} folder:")
    for item in results_folder_content:
        print(item)

    return "done", results_folder_content[0], results_folder_content[1]

css="""
#col-container{
    margin: 0 auto;
    max-width: 840px;
    text-align: left;
}
"""
    
with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.HTML("""
        <h2 style="text-align: center;">ProPainter</h2>
        <p style="text-align: center;">
            [ICCV 2023] ProPainter: Improving Propagation and Transformer for Video Inpainting <br />
            <a href="https://github.com/sczhou/ProPainter" target="_blank">code</a> | <a href="https://shangchenzhou.com/projects/ProPainter/" target="_blank">project page</a>
        </p>
                """)

        with gr.Row():
            with gr.Tab("Manual"):
                with gr.Column():
                    project_name = gr.Textbox(label="Name your project", value="my-new-project")
                    video_frames = gr.File(label="Video frames", file_types=["image"], file_count="multiple")
                    masks_frames = gr.File(label="Masks frames", file_types=["image"], file_count="multiple")
        
                    submit_btn = gr.Button("Submit")

            with gr.Tab("Auto"):
                with gr.Column():
                    project_name_2 = gr.Textbox(label="Name your project", value="my-new-project")
                    video_in = gr.Video(label="Source video", source="upload", format="mp4")
                    subject_to_remove = gr.Textbox(label="Subject to remove")

                    submit_auto_btn = gr.Button("Submit")

            with gr.Column():
                result = gr.Textbox(label="Result")
                res_masked = gr.Video(label="Masked video")
                res_files = gr.Video(label="Final result")

        
            
    submit_btn.click(fn=infer, inputs=[video_frames, masks_frames, project_name], outputs=[result, res_masked, res_files])
    submit_auto_btn.click(fn=infer_auto, inputs=[project_name_2, video_in, subject_to_remove], outputs=[result, res_masked, res_files])
    
demo.queue(max_size=12).launch()