de4dn0te/kakisalmi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Multiple output formats

Browse Source
This commit is contained in:
de4dn0te
2026-07-14 09:41:59 +00:00
committed by de4dn0te
parent bdb69239bc
commit 3c48ed082b
2 changed files with 79 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rembg_video.py
+65 -93
View File
@@ -6,7 +6,7 @@ import pathlib
import threading
import numpy as np
from queue import Queue
from shutil import rmtree
from shutil import rmtree, move
from PIL import Image
from rembg import new_session, remove
@@ -19,9 +19,7 @@ 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/output.mov", help="Define output path"
)
parser.add_argument("-o", type=str, default="export", help="Define output path")
parser.add_argument(
"--model",
type=str,
@@ -52,6 +50,13 @@ parser.add_argument(
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()
@@ -75,71 +80,26 @@ def image_to_bytes(image):
return buffer.getvalue()
def remove_with_fallback(image_bytes, session, scales=(1.0, 0.8, 0.6, 0.4)):
original = Image.open(io.BytesIO(image_bytes)).convert("RGBA")
width, height = original.size
last_exc = None
for scale in scales:
try:
if scale == 1.0:
return remove(image_bytes, session=session)
resized = original.resize(
(
max(1, int(width * scale)),
max(1, int(height * scale)),
),
Image.Resampling.LANCZOS,
)
with io.BytesIO() as buffer:
resized.save(buffer, format="TIFF")
scaled_bytes = buffer.getvalue()
scaled_output = remove(scaled_bytes, session=session)
if isinstance(scaled_output, (bytes, bytearray)):
scaled_output_bytes = bytes(scaled_output)
elif isinstance(scaled_output, np.ndarray):
with io.BytesIO() as buffer:
Image.fromarray(scaled_output).save(buffer, format="TIFF")
scaled_output_bytes = buffer.getvalue()
elif isinstance(scaled_output, Image.Image):
with io.BytesIO() as buffer:
scaled_output.save(buffer, format="TIFF")
scaled_output_bytes = 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 RuntimeError(
"Unexpected rembg remove() result type during mask fallback."
)
alpha = (
Image.open(io.BytesIO(scaled_output_bytes))
.convert("RGBA")
.getchannel("A")
)
alpha = alpha.resize((width, height), Image.Resampling.LANCZOS)
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.")
output_full = original.copy()
output_full.putalpha(alpha)
return image_to_bytes(output_full)
except Exception as exc:
last_exc = exc
if not is_oom_error(exc):
raise
if scale == scales[-1]:
raise RuntimeError(
"Insufficient GPU memory: background removal failed even after fallback downscales."
) from exc
print(
f"OOM detected during removal at scale {scale:.2f}; retrying with lower-resolution mask...",
flush=True,
)
if last_exc is not None:
raise RuntimeError(
"Background removal failed unexpectedly during fallback."
) from last_exc
raise RuntimeError("Background removal failed unexpectedly.")
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
@@ -154,12 +114,15 @@ 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
rmtree(frames_dir, ignore_errors=True)
os.mkdir(frames_dir)
stream = ffmpeg.input(args.input)
stream = ffmpeg.output(stream, os.path.join(frames_dir, "%04d.tiff"))
@@ -339,33 +302,42 @@ try:
rmtree(processed_dir)
os.rename(smoothed_dir, processed_dir)
# Output video
output_file = pathlib.Path(args.o)
output_file.parent.mkdir(exist_ok=True, parents=True)
if args.output_type != "mask_seq":
# Output video
output_file = pathlib.Path(args.o) / ("output.mov")
output_file.parent.mkdir(exist_ok=True, parents=True)
stream = ffmpeg.input(
os.path.join(processed_dir, "%04d.tiff"),
r=framerate,
f="image2",
s=whstr,
pix_fmt="yuva420p",
)
stream = ffmpeg.output(stream, args.o, vcodec="prores_ks", **{"profile:v": "4"})
stream = ffmpeg.input(
os.path.join(processed_dir, "%04d.tiff"),
r=framerate,
f="image2",
s=whstr,
pix_fmt="yuva420p",
)
stream = ffmpeg.output(
stream, str(output_file), vcodec="prores_ks", **{"profile:v": "4"}
)
"""
stream = ffmpeg.output(
stream,
args.o,
vcodec="libvpx-vp9",
pix_fmt="yuva420p",
**{
"crf": "23",
"b:v": "0",
"auto-alt-ref": "0",
},
)
"""
ffmpeg.run(stream)
"""
stream = ffmpeg.output(
stream,
output_file,
vcodec="libvpx-vp9",
pix_fmt="yuva420p",
**{
"crf": "23",
"b:v": "0",
"auto-alt-ref": "0",
},
)
"""
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")