Spaces:
Paused
Paused
| 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_in) | |
| #check fps | |
| if clip.fps > 30: | |
| print("vide rate is over 30, resetting to 30") | |
| clip_resized = clip.resize(width=512) | |
| clip_resized.write_videofile(f"{img_type}_video_resized.mp4", fps=30) | |
| else: | |
| print("video rate is OK") | |
| clip_resized = clip.resize(width=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): | |
| timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S") | |
| 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 | |
| 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(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) | |
| 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_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 = 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 "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() |