347 lines
11 KiB
Python
347 lines
11 KiB
Python
import argparse
|
|
import io
|
|
import os
|
|
import ffmpeg
|
|
import pathlib
|
|
import threading
|
|
import numpy as np
|
|
from queue import Queue
|
|
from shutil import rmtree, move
|
|
from PIL import Image
|
|
from rembg import new_session, remove
|
|
|
|
os.environ.setdefault(
|
|
"PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True,managed_memory:True"
|
|
)
|
|
|
|
# Parse args
|
|
parser = argparse.ArgumentParser(
|
|
description="Applies rembg background removal to the frames of a video"
|
|
)
|
|
parser.add_argument("input", type=str, help="Input video")
|
|
parser.add_argument("-o", type=str, default="export", help="Define output path")
|
|
parser.add_argument(
|
|
"--model",
|
|
type=str,
|
|
default="u2net_human_seg",
|
|
help="rembg model to use (default: u2net_human_seg)",
|
|
)
|
|
parser.add_argument(
|
|
"--workers",
|
|
type=int,
|
|
default=1,
|
|
help="Number of concurrent processing workers (default: 1)",
|
|
)
|
|
parser.add_argument(
|
|
"--smooth",
|
|
type=int,
|
|
default=3,
|
|
help="Temporal mask smoothing window size in frames (default: 3, 0 to disable)",
|
|
)
|
|
parser.add_argument(
|
|
"--buffer-size",
|
|
type=int,
|
|
default=8,
|
|
help="Set maximum number of frames in buffer (default: 8)",
|
|
)
|
|
parser.add_argument(
|
|
"--smooth-workers",
|
|
type=int,
|
|
default=os.cpu_count() or 4,
|
|
help="Number of cpu threads to use for temporal mask smoothing (default: cpu count)",
|
|
)
|
|
parser.add_argument(
|
|
"--output-type",
|
|
type=str,
|
|
choices=["complete", "mask", "mask_seq"],
|
|
default="complete",
|
|
help="What way to output keyed video. (default: complete)",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
|
|
def is_oom_error(exc):
|
|
text = str(exc).lower()
|
|
return any(
|
|
phrase in text
|
|
for phrase in (
|
|
"out of memory",
|
|
"cuda out of memory",
|
|
"failed to allocate",
|
|
"oom",
|
|
"memory error",
|
|
)
|
|
)
|
|
|
|
|
|
def image_to_bytes(image):
|
|
with io.BytesIO() as buffer:
|
|
image.save(buffer, format="TIFF")
|
|
return buffer.getvalue()
|
|
|
|
|
|
def remove_with_fallback(image_bytes, session):
|
|
try:
|
|
rembg_out = remove(image_bytes, session=session)
|
|
if args.output_type == "complete":
|
|
return rembg_out
|
|
elif args.output_type in ("mask", "mask_seq"):
|
|
if isinstance(rembg_out, (bytes, bytearray)):
|
|
mask_bytes = bytes(rembg_out)
|
|
else:
|
|
raise TypeError('mask_bytes is not of type "bytes"')
|
|
alpha = Image.open(io.BytesIO(mask_bytes)).convert("RGBA").getchannel("A")
|
|
return image_to_bytes(alpha)
|
|
else:
|
|
raise Exception("Unknown output type.")
|
|
|
|
except Exception as exc:
|
|
if is_oom_error(exc):
|
|
raise RuntimeError("Insufficient GPU memory!") from exc
|
|
else:
|
|
raise RuntimeError("Background removal failed unexpectedly.")
|
|
|
|
|
|
# Extract video info
|
|
probe = ffmpeg.probe(args.input)
|
|
video_stream = next(
|
|
(stream for stream in probe["streams"] if stream["codec_type"] == "video"), None
|
|
)
|
|
if video_stream is None:
|
|
raise ValueError(f"No video stream found in input file: {args.input}")
|
|
width = int(video_stream["width"])
|
|
height = int(video_stream["height"])
|
|
whstr = str(width) + "x" + str(height)
|
|
framerate = video_stream["avg_frame_rate"]
|
|
|
|
# Setup workspace folders
|
|
frames_dir = os.path.join(str(pathlib.Path(__file__).parent.absolute()), "frames")
|
|
processed_dir = os.path.join(str(pathlib.Path(__file__).parent.absolute()), "processed")
|
|
smoothed_dir = processed_dir + "_smoothed"
|
|
rmtree(frames_dir, ignore_errors=True)
|
|
rmtree(processed_dir, ignore_errors=True)
|
|
rmtree(smoothed_dir, ignore_errors=True)
|
|
|
|
# Extract input video frames
|
|
os.mkdir(frames_dir)
|
|
stream = ffmpeg.input(args.input)
|
|
stream = ffmpeg.output(stream, os.path.join(frames_dir, "%04d.tiff"))
|
|
ffmpeg.run(stream)
|
|
|
|
_SENTINEL = object()
|
|
|
|
try:
|
|
# Process frames with pipelined reader -> processors -> writer
|
|
if not os.path.isdir(processed_dir):
|
|
os.mkdir(processed_dir)
|
|
|
|
files = sorted(os.listdir(frames_dir))
|
|
total_files = len(files)
|
|
|
|
print(f"Loading rembg session (model={args.model})...", flush=True)
|
|
session = new_session(
|
|
args.model, providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
|
|
)
|
|
|
|
read_queue = Queue(maxsize=args.buffer_size)
|
|
write_queue = Queue(maxsize=args.buffer_size)
|
|
errors = []
|
|
active_processors = [
|
|
args.workers
|
|
] # list so processor() can mutate without nonlocal
|
|
active_processors_lock = threading.Lock()
|
|
|
|
def reader():
|
|
try:
|
|
for idx, file in enumerate(files, 1):
|
|
frame_path = os.path.join(frames_dir, file)
|
|
with open(frame_path, "rb") as f:
|
|
data = f.read()
|
|
read_queue.put((idx, file, data))
|
|
os.remove(frame_path)
|
|
except Exception as e:
|
|
errors.append(e)
|
|
finally:
|
|
# One sentinel per worker so each one knows when to stop
|
|
for _ in range(args.workers):
|
|
read_queue.put(_SENTINEL)
|
|
|
|
def processor():
|
|
try:
|
|
while True:
|
|
item = read_queue.get()
|
|
if item is _SENTINEL:
|
|
break
|
|
idx, file, input_data = item
|
|
print(f"Processing frame {idx}/{total_files}: {file}", flush=True)
|
|
output_data = remove_with_fallback(input_data, session=session)
|
|
write_queue.put((idx, file, output_data))
|
|
except Exception as e:
|
|
errors.append(e)
|
|
finally:
|
|
# Signal writer only when the last processor finishes
|
|
with active_processors_lock:
|
|
active_processors[0] -= 1
|
|
if active_processors[0] == 0:
|
|
write_queue.put(_SENTINEL)
|
|
|
|
def writer():
|
|
try:
|
|
while True:
|
|
item = write_queue.get()
|
|
if item is _SENTINEL:
|
|
break
|
|
idx, file, output_data = item
|
|
with open(os.path.join(processed_dir, file), "wb") as f:
|
|
f.write(output_data)
|
|
print(f"Written frame {idx}/{total_files}: {file}", flush=True)
|
|
except Exception as e:
|
|
errors.append(e)
|
|
|
|
reader_thread = threading.Thread(target=reader, daemon=True)
|
|
processor_threads = [
|
|
threading.Thread(target=processor, daemon=True) for _ in range(args.workers)
|
|
]
|
|
writer_thread = threading.Thread(target=writer, daemon=True)
|
|
|
|
reader_thread.start()
|
|
for t in processor_threads:
|
|
t.start()
|
|
writer_thread.start()
|
|
|
|
reader_thread.join()
|
|
for t in processor_threads:
|
|
t.join()
|
|
writer_thread.join()
|
|
|
|
if errors:
|
|
raise errors[0]
|
|
|
|
# Temporal mask smoothing
|
|
if args.smooth > 0:
|
|
files = sorted(os.listdir(processed_dir))
|
|
total = len(files)
|
|
window = args.smooth
|
|
half = window // 2
|
|
print(
|
|
f"Applying temporal mask smoothing (window={window}, "
|
|
f"workers={args.smooth_workers})...",
|
|
flush=True,
|
|
)
|
|
|
|
smoothing_errors = []
|
|
progress_lock = threading.Lock()
|
|
progress_count = [0]
|
|
n_workers = max(1, args.smooth_workers)
|
|
|
|
# Write smoothed frames to a separate directory rather than
|
|
# overwriting processed_dir in place. Overlapping windows mean a
|
|
# frame can be a *read* dependency for several write_idx tasks;
|
|
# writing in place risked one thread reading a file while another
|
|
# was mid-save on it (truncated/corrupt img -> shape errors).
|
|
if not os.path.isdir(smoothed_dir):
|
|
os.mkdir(smoothed_dir)
|
|
|
|
def get_alpha(idx):
|
|
file = files[idx]
|
|
img = Image.open(os.path.join(processed_dir, file)).convert("RGBA")
|
|
return np.array(img)[:, :, 3].astype(np.float32)
|
|
|
|
def smooth_frame(write_idx):
|
|
try:
|
|
start = max(0, write_idx - half)
|
|
end = min(total - 1, write_idx + half)
|
|
alphas = np.stack([get_alpha(j) for j in range(start, end + 1)])
|
|
smoothed_alpha = np.mean(alphas, axis=0).astype(np.uint8)
|
|
|
|
filename = files[write_idx]
|
|
out_img = Image.open(os.path.join(processed_dir, filename)).convert(
|
|
"RGBA"
|
|
)
|
|
out_arr = np.array(out_img)
|
|
out_arr[:, :, 3] = smoothed_alpha
|
|
Image.fromarray(out_arr).save(os.path.join(smoothed_dir, filename))
|
|
|
|
with progress_lock:
|
|
progress_count[0] += 1
|
|
print(
|
|
f"Smoothed frame {progress_count[0]}/{total}: {filename}",
|
|
flush=True,
|
|
)
|
|
except Exception as e:
|
|
smoothing_errors.append(e)
|
|
|
|
smooth_queue = Queue()
|
|
for write_idx in range(total):
|
|
smooth_queue.put(write_idx)
|
|
|
|
def smoothing_worker():
|
|
while True:
|
|
try:
|
|
write_idx = smooth_queue.get_nowait()
|
|
except Exception:
|
|
return
|
|
if smoothing_errors:
|
|
return
|
|
smooth_frame(write_idx)
|
|
|
|
smoothing_threads = [
|
|
threading.Thread(target=smoothing_worker, daemon=True)
|
|
for _ in range(n_workers)
|
|
]
|
|
for t in smoothing_threads:
|
|
t.start()
|
|
for t in smoothing_threads:
|
|
t.join()
|
|
|
|
if smoothing_errors:
|
|
rmtree(smoothed_dir, ignore_errors=True)
|
|
raise smoothing_errors[0]
|
|
|
|
# Swap the smoothed frames in as the new processed_dir contents.
|
|
rmtree(processed_dir)
|
|
os.rename(smoothed_dir, processed_dir)
|
|
|
|
if args.output_type != "mask_seq":
|
|
# Output video
|
|
stream = ffmpeg.input(
|
|
os.path.join(processed_dir, "%04d.tiff"),
|
|
r=framerate,
|
|
f="image2",
|
|
s=whstr,
|
|
)
|
|
if args.output_type == "mask":
|
|
output_file = pathlib.Path(args.o) / ("output.mp4")
|
|
output_file.parent.mkdir(exist_ok=True, parents=True)
|
|
stream = ffmpeg.output(
|
|
stream,
|
|
str(output_file),
|
|
vcodec="libx264",
|
|
pix_fmt="gray",
|
|
crf=0,
|
|
preset="veryslow",
|
|
tune="animation",
|
|
)
|
|
else:
|
|
output_file = pathlib.Path(args.o) / ("output.mov")
|
|
output_file.parent.mkdir(exist_ok=True, parents=True)
|
|
stream = ffmpeg.output(
|
|
stream, str(output_file), vcodec="prores_ks", **{"profile:v": "4"}
|
|
)
|
|
|
|
ffmpeg.run(stream, overwrite_output=True)
|
|
|
|
else:
|
|
img_seq_out_folder = os.path.join(args.o, "output_img_seq")
|
|
rmtree(img_seq_out_folder, ignore_errors=True)
|
|
move(processed_dir, img_seq_out_folder)
|
|
|
|
except KeyboardInterrupt:
|
|
print("\nInterrupted by user")
|
|
|
|
finally:
|
|
print("Removing temporary files...")
|
|
rmtree(processed_dir, ignore_errors=True)
|
|
rmtree(frames_dir, ignore_errors=True)
|
|
rmtree(smoothed_dir, ignore_errors=True)
|