Upload remove_bg_script.py

remove_bg_script.py

@@ -0,0 +1,114 @@
+'''
+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)