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 results_folder_content[0], results_folder_content[1]
def get_frames(video_in, img_type):
frames = []
#resize the video
clip = VideoFileClip(video_in)
#check fps
if clip.fps > 30:
print("vide rate is over 30, resetting to 30")
clip_resized = clip.resize(height=512)
clip_resized.write_videofile(f"{img_type}_video_resized.mp4", fps=30)
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
10, # 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):
timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
print(video_in)
matte_video = get_matte(video_in, subject_to_remove)
# Cut the video to the first 3 seconds
#video_cut = f"video_cut.mp4"
#ffmpeg_extract_subclip(video_in, t1=0, t2=3, targetname=video_cut)
video_frames = get_frames(video_in, "source")
frames_list = video_frames[0]
print(video_frames[0])
masks_frames = get_frames(matte_video, "mask")
masks_list = masks_frames[0]
print(masks_frames[0])
# Check the lengths of the two lists
frames_length = len(frames_list)
masks_length = len(masks_list)
# Make the lists the same length if they are different
if frames_length > masks_length:
frames_list = frames_list[:masks_length]
elif masks_length > frames_length:
masks_list = masks_list[:frames_length]
# Now, both lists have the same length
# 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
my_video_directory = f"{project_name}_{timestamp}"
os.makedirs(my_video_directory)
print(f"Created the dir: {my_video_directory}")
# Assuming 'images' is a list of image file paths
for idx, image in enumerate(frames_list):
# 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)
else:
# If the directory already exists, add a timestamp to the new directory name
my_masks_directory = f"{project_name}_masks_{timestamp}"
os.makedirs(my_masks_directory)
print(f"Created the dir: {my_masks_directory}")
# Assuming 'images' is a list of image file paths
for idx, image in enumerate(masks_list):
# 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 = f"results_{timestamp}"
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)}",
#f"--fp16"
]
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, my_video_directory)
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 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("""
ProPainter
[ICCV 2023] ProPainter: Improving Propagation and Transformer for Video Inpainting
code | project page
""")
with gr.Row():
with gr.Tab("Manual"):
with gr.Column():
project_name = gr.Textbox(label="Name your project", info="no spaces nor special characters", 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", info="no spaces nor special characters", 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")
gr.Examples(
examples = [
[
"example_1",
"hf-examples/knight.mp4",
"knight"
],
[
"example_2",
"hf-examples/knight.mp4",
"horse"
],
[
"example_3",
"hf-examples/knight.mp4",
"tail"
]
],
fn = infer_auto,
inputs=[project_name_2, video_in, subject_to_remove],
#outputs=[res_masked, res_files]
)
with gr.Column():
gr.Markdown("""
### Results
""")
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=[res_masked, res_files])
submit_auto_btn.click(fn=infer_auto, inputs=[project_name_2, video_in, subject_to_remove], outputs=[res_masked, res_files])
demo.queue(max_size=12).launch()