Spaces:
Running
Running
| ''' | |
| pip install torch accelerate opencv-python pillow numpy timm kornia prettytable typing scikit-image transformers>=4.39.1 gradio==4.44.1 gradio_imageslider loadimg>=0.1.1 "httpx[socks]" moviepy==1.0.3 | |
| huggingface-cli download \ | |
| --repo-type dataset svjack/video-dataset-Lily-Bikini-organized \ | |
| --local-dir video-dataset-Lily-Bikini-organized | |
| python remove_bg_script.py video-dataset-Lily-Bikini-organized video-dataset-Lily-Bikini-rm-background-organized --copy_others | |
| ''' | |
| from PIL import Image, ImageChops | |
| import torch | |
| from torchvision import transforms | |
| from transformers import AutoModelForImageSegmentation | |
| from moviepy.editor import VideoFileClip, ImageSequenceClip | |
| import numpy as np | |
| from tqdm import tqdm | |
| from uuid import uuid1 | |
| import os | |
| import shutil | |
| import argparse | |
| # Load the model | |
| model = AutoModelForImageSegmentation.from_pretrained('briaai/RMBG-2.0', trust_remote_code=True) | |
| torch.set_float32_matmul_precision('high') # Set precision | |
| model.to('cuda') | |
| model.eval() | |
| # Data settings | |
| image_size = (1024, 1024) | |
| transform_image = transforms.Compose([ | |
| transforms.Resize(image_size), | |
| transforms.ToTensor(), | |
| transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) | |
| ]) | |
| def remove_background(image): | |
| """Remove background from a single image.""" | |
| input_images = transform_image(image).unsqueeze(0).to('cuda') | |
| # Prediction | |
| with torch.no_grad(): | |
| preds = model(input_images)[-1].sigmoid().cpu() | |
| pred = preds[0].squeeze() | |
| # Convert the prediction to a mask | |
| mask = (pred * 255).byte() # Convert to 0-255 range | |
| mask_pil = transforms.ToPILImage()(mask).convert("L") | |
| mask_resized = mask_pil.resize(image.size, Image.LANCZOS) | |
| # Apply the mask to the image | |
| image.putalpha(mask_resized) | |
| return image, mask_resized | |
| def process_video(input_video_path, output_video_path): | |
| """Process a video to remove the background from each frame.""" | |
| # Load the video | |
| video_clip = VideoFileClip(input_video_path) | |
| # Process each frame | |
| frames = [] | |
| for frame in tqdm(video_clip.iter_frames()): | |
| frame_pil = Image.fromarray(frame) | |
| frame_no_bg, mask_resized = remove_background(frame_pil) | |
| path = "{}.png".format(uuid1()) | |
| frame_no_bg.save(path) | |
| frame_no_bg = Image.open(path).convert("RGBA") | |
| os.remove(path) | |
| # Convert mask_resized to RGBA mode | |
| mask_resized_rgba = mask_resized.convert("RGBA") | |
| # Apply the mask using ImageChops.multiply | |
| output = ImageChops.multiply(frame_no_bg, mask_resized_rgba) | |
| output_np = np.array(output) | |
| frames.append(output_np) | |
| # Save the processed frames as a new video | |
| processed_clip = ImageSequenceClip(frames, fps=video_clip.fps) | |
| processed_clip.write_videofile(output_video_path, codec='libx264', ffmpeg_params=['-pix_fmt', 'yuva420p']) | |
| def copy_non_video_files(input_path, output_path): | |
| """Copy non-video files and directories from input path to output path.""" | |
| for item in os.listdir(input_path): | |
| item_path = os.path.join(input_path, item) | |
| if not item.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')): | |
| dest_path = os.path.join(output_path, item) | |
| if os.path.isdir(item_path): | |
| shutil.copytree(item_path, dest_path) | |
| else: | |
| shutil.copy2(item_path, dest_path) | |
| def main(input_path, output_path, copy_others=False): | |
| if not os.path.exists(output_path): | |
| os.makedirs(output_path) | |
| if copy_others: | |
| copy_non_video_files(input_path, output_path) | |
| for video_name in os.listdir(input_path): | |
| if video_name.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')): | |
| input_video_path = os.path.join(input_path, video_name) | |
| output_video_path = os.path.join(output_path, video_name) | |
| process_video(input_video_path, output_video_path) | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser(description="Process videos to remove background.") | |
| parser.add_argument("input_path", type=str, help="Path to the input directory containing videos.") | |
| parser.add_argument("output_path", type=str, help="Path to the output directory for processed videos.") | |
| parser.add_argument("--copy_others", action="store_true", help="Copy non-video files and directories from input to output.") | |
| args = parser.parse_args() | |
| main(args.input_path, args.output_path, args.copy_others) | |