mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
SnapOtterandGitHub
cb5db59f77
Adds a dedicated remove-gif-background AI tool: removes the background from an animated GIF, WebP, or APNG frame by frame and reassembles a transparent (or composited) animation in WebP, APNG, or GIF, with full per-frame effects. Reuses the background-removal bundle. Verified end-to-end with the real rembg model. Closes #496.
343 lines
12 KiB
Python
343 lines
12 KiB
Python
"""Background removal for animated images (GIF, animated WebP, APNG).
|
|
|
|
Reads every frame with disposal-aware coalescing, runs one warm rembg session
|
|
over all frames, applies the chosen effect per frame, and re-encodes an animated
|
|
transparent (or composited) output via Pillow. Reuses remove_bg.py's model
|
|
registration and edge-refinement helpers so the still and animated paths share
|
|
the same matte behaviour.
|
|
"""
|
|
import io
|
|
import json
|
|
import os
|
|
import sys
|
|
|
|
MAX_REMBG_PX = int(os.environ.get("MAX_REMBG_PX", "2048"))
|
|
_OOM_MARKERS = (
|
|
"out of memory",
|
|
"failed to allocate",
|
|
"cudaerrormemoryallocation",
|
|
"cublas_status_alloc_failed",
|
|
"bad_alloc",
|
|
)
|
|
|
|
|
|
def emit_progress(percent, stage):
|
|
print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)
|
|
|
|
|
|
def _is_oom(err):
|
|
m = str(err).lower()
|
|
return any(s in m for s in _OOM_MARKERS)
|
|
|
|
|
|
def _read_frames(path):
|
|
"""Disposal-aware coalescing. Returns (frames[RGBA], durations[ms], loop).
|
|
|
|
Naive ``seek()`` yields partial/garbled frames on disposal-optimized GIFs and
|
|
animated WebP, so composite each frame onto a running canvas. ``copy()`` is
|
|
mandatory: ``ImageSequence.Iterator`` mutates the same underlying object.
|
|
"""
|
|
from PIL import Image, ImageSequence
|
|
|
|
im = Image.open(path)
|
|
try:
|
|
loop = int(im.info.get("loop", 0) or 0)
|
|
except (TypeError, ValueError):
|
|
loop = 0
|
|
frames, durations = [], []
|
|
canvas = None
|
|
for frame in ImageSequence.Iterator(im):
|
|
raw = frame.info.get("duration", im.info.get("duration", 100))
|
|
try:
|
|
dur = max(20, int(raw))
|
|
except (TypeError, ValueError):
|
|
dur = 100
|
|
disposal = frame.info.get("disposal", 0)
|
|
rgba = frame.convert("RGBA")
|
|
if canvas is None:
|
|
canvas = rgba.copy()
|
|
elif disposal == 2:
|
|
canvas = rgba.copy() # restore-to-background: don't smear the prior frame
|
|
else:
|
|
canvas = Image.alpha_composite(canvas, rgba)
|
|
frames.append(canvas.copy())
|
|
durations.append(dur)
|
|
return frames, durations, loop
|
|
|
|
|
|
def _scale_target(size):
|
|
w, h = size
|
|
longest = max(w, h)
|
|
if longest <= MAX_REMBG_PX:
|
|
return None
|
|
scale = MAX_REMBG_PX / longest
|
|
return (max(1, round(w * scale)), max(1, round(h * scale)))
|
|
|
|
|
|
def _hex_to_rgba(value, default=(255, 255, 255, 255)):
|
|
if not value:
|
|
return default
|
|
s = str(value).lstrip("#")
|
|
try:
|
|
if len(s) == 3:
|
|
s = "".join(c * 2 for c in s)
|
|
if len(s) == 6:
|
|
return (int(s[0:2], 16), int(s[2:4], 16), int(s[4:6], 16), 255)
|
|
if len(s) == 8:
|
|
return (int(s[0:2], 16), int(s[2:4], 16), int(s[4:6], 16), int(s[6:8], 16))
|
|
except ValueError:
|
|
pass
|
|
return default
|
|
|
|
|
|
def _gradient(w, h, settings):
|
|
import math
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
c1 = np.array(_hex_to_rgba(settings.get("gradientColor1", "#000000")), dtype=np.float32)
|
|
c2 = np.array(_hex_to_rgba(settings.get("gradientColor2", "#ffffff")), dtype=np.float32)
|
|
angle = math.radians(float(settings.get("gradientAngle", 0)) % 360.0)
|
|
yy, xx = np.mgrid[0:h, 0:w].astype(np.float32)
|
|
proj = xx * math.cos(angle) + yy * math.sin(angle)
|
|
lo, hi = float(proj.min()), float(proj.max())
|
|
t = (proj - lo) / (hi - lo) if hi > lo else np.zeros_like(proj)
|
|
t = t[:, :, None]
|
|
arr = (c1 * (1.0 - t) + c2 * t).astype(np.uint8)
|
|
return Image.fromarray(arr, "RGBA")
|
|
|
|
|
|
def _cover(img, w, h):
|
|
from PIL import Image
|
|
|
|
iw, ih = img.size
|
|
scale = max(w / iw, h / ih)
|
|
nw, nh = max(1, round(iw * scale)), max(1, round(ih * scale))
|
|
resized = img.resize((nw, nh), Image.LANCZOS)
|
|
left, top = (nw - w) // 2, (nh - h) // 2
|
|
return resized.crop((left, top, left + w, top + h)).convert("RGBA")
|
|
|
|
|
|
def _apply_shadow(base, cutout, settings):
|
|
import numpy as np
|
|
from PIL import Image, ImageFilter
|
|
|
|
opacity = max(0.0, min(1.0, float(settings.get("shadowOpacity", 50)) / 100.0))
|
|
w, h = cutout.size
|
|
alpha = cutout.split()[3]
|
|
shadow = Image.new("RGBA", (w, h), (0, 0, 0, 0))
|
|
shadow.paste(Image.new("RGBA", (w, h), (0, 0, 0, 255)), (0, 0), alpha)
|
|
shadow = shadow.filter(ImageFilter.GaussianBlur(radius=max(2.0, w * 0.02)))
|
|
sa = np.array(shadow, dtype=np.float32)
|
|
sa[:, :, 3] *= opacity
|
|
shadow = Image.fromarray(sa.astype(np.uint8), "RGBA")
|
|
offset = max(2, int(w * 0.015))
|
|
shifted = Image.new("RGBA", (w, h), (0, 0, 0, 0))
|
|
shifted.paste(shadow, (offset, offset), shadow)
|
|
return Image.alpha_composite(base, shifted)
|
|
|
|
|
|
def _apply_effect(cutout, original, settings, bg_img):
|
|
from PIL import Image, ImageFilter
|
|
|
|
bg_type = settings.get("backgroundType", "transparent")
|
|
w, h = cutout.size
|
|
if bg_type == "color":
|
|
base = Image.new("RGBA", (w, h), _hex_to_rgba(settings.get("backgroundColor", "#ffffff")))
|
|
elif bg_type == "gradient":
|
|
base = _gradient(w, h, settings)
|
|
elif bg_type == "blur":
|
|
radius = max(1.0, float(settings.get("blurIntensity", 20)) / 3.0)
|
|
base = original.convert("RGBA").filter(ImageFilter.GaussianBlur(radius=radius))
|
|
elif bg_type == "image" and bg_img is not None:
|
|
base = _cover(bg_img, w, h)
|
|
else:
|
|
base = Image.new("RGBA", (w, h), (0, 0, 0, 0))
|
|
if settings.get("shadowEnabled"):
|
|
base = _apply_shadow(base, cutout, settings)
|
|
return Image.alpha_composite(base, cutout)
|
|
|
|
|
|
def _to_gif_frame(rgba):
|
|
"""RGBA -> (P-mode image, transparent_index) for a transparent animated GIF.
|
|
|
|
GIF alpha is 1-bit: threshold at 50%, quantize RGB to 255 colours, and
|
|
reserve palette index 255 for transparency.
|
|
"""
|
|
from PIL import Image
|
|
|
|
alpha = rgba.split()[3]
|
|
p = rgba.convert("RGB").quantize(colors=255, method=Image.MEDIANCUT)
|
|
transparent_mask = alpha.point(lambda a: 255 if a < 128 else 0)
|
|
p.paste(255, transparent_mask)
|
|
return p, 255
|
|
|
|
|
|
def _encode(frames, durations, loop, fmt, out_path):
|
|
if fmt == "apng":
|
|
frames[0].save(
|
|
out_path, "PNG", save_all=True, append_images=frames[1:],
|
|
duration=durations, loop=loop, disposal=1,
|
|
)
|
|
elif fmt == "gif":
|
|
pal = [_to_gif_frame(f) for f in frames]
|
|
first, tidx = pal[0]
|
|
first.save(
|
|
out_path, "GIF", save_all=True, append_images=[p for p, _ in pal[1:]],
|
|
duration=durations, loop=loop, transparency=tidx, disposal=2, optimize=False,
|
|
)
|
|
else: # webp
|
|
frames[0].save(
|
|
out_path, "WEBP", save_all=True, append_images=frames[1:],
|
|
duration=durations, loop=loop, quality=90, method=6,
|
|
)
|
|
|
|
|
|
def _create_session(model, providers, device):
|
|
from rembg import new_session
|
|
from rembg.sessions import sessions_class
|
|
|
|
from remove_bg import _register_hr_matting_session, _register_matting_session
|
|
|
|
_register_matting_session(sessions_class)
|
|
_register_hr_matting_session(sessions_class)
|
|
try:
|
|
return new_session(model, providers=providers), device
|
|
except Exception:
|
|
if "CUDAExecutionProvider" in providers:
|
|
return new_session(model, providers=["CPUExecutionProvider"]), "cpu"
|
|
raise
|
|
|
|
|
|
def _remove_one(frame_rgba, session, use_alpha, settings, target, orig_size):
|
|
from rembg import remove
|
|
from PIL import Image
|
|
|
|
src = frame_rgba if target is None else frame_rgba.resize(target, Image.LANCZOS)
|
|
buf = io.BytesIO()
|
|
src.save(buf, format="PNG")
|
|
data = buf.getvalue()
|
|
try:
|
|
out = remove(
|
|
data, session=session, alpha_matting=use_alpha,
|
|
alpha_matting_foreground_threshold=240, alpha_matting_background_threshold=10,
|
|
)
|
|
except Exception as e:
|
|
if use_alpha and not _is_oom(e):
|
|
out = remove(data, session=session, alpha_matting=False)
|
|
else:
|
|
raise
|
|
edge_refine = settings.get("edgeRefine", 0)
|
|
if edge_refine and int(edge_refine) > 0:
|
|
from remove_bg import _refine_edges
|
|
|
|
out = _refine_edges(out, int(edge_refine))
|
|
if settings.get("decontaminate"):
|
|
from remove_bg import _decontaminate_edges
|
|
|
|
out = _decontaminate_edges(out)
|
|
cut = Image.open(io.BytesIO(out)).convert("RGBA")
|
|
if target is not None:
|
|
cut = cut.resize(orig_size, Image.LANCZOS)
|
|
return cut
|
|
|
|
|
|
def main():
|
|
input_path = sys.argv[1]
|
|
output_path = sys.argv[2]
|
|
settings = json.loads(sys.argv[3]) if len(sys.argv) > 3 else {}
|
|
|
|
fmt = settings.get("outputFormat", "webp")
|
|
if fmt not in ("webp", "gif", "apng"):
|
|
fmt = "webp"
|
|
cancel_file = settings.get("cancelFile")
|
|
bg_path = settings.get("bgImagePath")
|
|
|
|
# Redirect stdout to stderr so rembg/onnx/pooch output cannot contaminate the
|
|
# JSON result. Restored only to write the final line.
|
|
stdout_fd = os.dup(1)
|
|
os.dup2(2, 1)
|
|
|
|
canceled = False
|
|
try:
|
|
from PIL import Image
|
|
|
|
from gpu import onnx_providers
|
|
from remove_bg import ALLOWED_MODELS
|
|
|
|
model = settings.get("model", "u2net")
|
|
if model not in ALLOWED_MODELS:
|
|
model = "u2net"
|
|
|
|
emit_progress(3, "Reading frames")
|
|
frames, durations, loop = _read_frames(input_path)
|
|
n = len(frames)
|
|
if n == 0:
|
|
raise RuntimeError("no frames found in input")
|
|
orig_size = frames[0].size
|
|
target = _scale_target(orig_size)
|
|
bg_img = Image.open(bg_path).convert("RGBA") if bg_path else None
|
|
|
|
# Offline guard for the chosen model (mirrors remove_bg.py home resolution).
|
|
model_home = os.path.expanduser(
|
|
os.getenv("U2NET_HOME", os.path.join(os.getenv("XDG_DATA_HOME", "~"), ".u2net"))
|
|
)
|
|
if not os.path.exists(os.path.join(model_home, f"{model}.onnx")):
|
|
from offline_guard import ensure_download_allowed
|
|
|
|
ensure_download_allowed(f"Background removal model '{model}'")
|
|
|
|
emit_progress(5, "Loading model")
|
|
providers, device = onnx_providers()
|
|
session, device = _create_session(model, providers, device)
|
|
use_alpha = device != "cpu"
|
|
|
|
# Probe frame 0 to settle model (OOM -> lighter model for the WHOLE
|
|
# animation, never per-frame) and matting viability once. Switching model
|
|
# or matting mid-animation would flicker.
|
|
try:
|
|
_remove_one(frames[0], session, use_alpha, settings, target, orig_size)
|
|
except Exception as e:
|
|
if _is_oom(e) and model.startswith("birefnet"):
|
|
model = "u2net"
|
|
emit_progress(5, "Retrying with a lighter model")
|
|
session, device = _create_session(model, providers, device)
|
|
use_alpha = device != "cpu"
|
|
elif use_alpha:
|
|
use_alpha = False # matting not viable on this device/model
|
|
else:
|
|
raise
|
|
|
|
out_frames = []
|
|
for i, frame in enumerate(frames):
|
|
if cancel_file and os.path.exists(cancel_file):
|
|
canceled = True
|
|
break
|
|
cut = _remove_one(frame, session, use_alpha, settings, target, orig_size)
|
|
out_frames.append(_apply_effect(cut, frame, settings, bg_img))
|
|
emit_progress(int(5 + 90 * (i + 1) / n), f"Frame {i + 1}/{n}")
|
|
|
|
if canceled:
|
|
result = json.dumps({"success": False, "error": "canceled"})
|
|
else:
|
|
emit_progress(97, "Encoding animation")
|
|
_encode(out_frames, durations, loop, fmt, output_path)
|
|
result = json.dumps(
|
|
{"success": True, "model": model, "device": device, "frames": n, "format": fmt}
|
|
)
|
|
except ImportError as e:
|
|
print(f"[gif-remove-bg] Import failed: {e}", file=sys.stderr, flush=True)
|
|
result = json.dumps({"success": False, "error": f"import failed: {e}"})
|
|
except Exception as e: # noqa: BLE001
|
|
result = json.dumps({"success": False, "error": str(e)})
|
|
|
|
os.dup2(stdout_fd, 1)
|
|
os.close(stdout_fd)
|
|
sys.stdout.write(result + "\n")
|
|
sys.stdout.flush()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|