huangyp/backgroundremover
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

去背景
Some checks failed
Publish Docker Image / Build/Push Docker Image to Docker Hub (push) Has been cancelled
Details

Browse Source
This commit is contained in:
huangyp 2025-03-24 14:41:54 +08:00
parent 635fad0e3d
commit ab94a43088
2 changed files with 11 additions and 177 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
backgroundremover/cmd/server.py
View File

@ -14,21 +14,22 @@ app = Flask(__name__)
def rm():
print('开始处理请求')
file = request.files['file']
model_name = request.form.get('model_name')
id = uuid.uuid4()
work_dir = os.path.abspath("/data/work/" + str(id))
print("work_dir: " + str(work_dir))
os.makedirs(work_dir)
file.save(work_dir + "/input.mp4")
file.save(os.path.join(work_dir, "input.mp4"))
utilities.transparentvideo(work_dir + "/output.mov", work_dir + "/input.mp4",
worker_nodes=1,
gpu_batchsize=2,
model_name="u2net",
frame_limit=-1,
framerate=-1)
utilities.matte_key(
output=os.path.join(work_dir, "output.mp4"),
file_path=os.path.join(work_dir, "input.mp4"),
worker_nodes=1,
gpu_batchsize=6,
model_name=model_name)
return send_file(work_dir + "/output.mov")
return send_file(work_dir + "/output.mp4")
def main():

171
backgroundremover/utilities.py
View File

@ -1,4 +1,3 @@
import os
import math
import torch.multiprocessing as multiprocessing
import subprocess as sp
@ -7,9 +6,6 @@ import ffmpeg
import numpy as np
import torch
from .bg import DEVICE, Net, iter_frames, remove_many
import tempfile
import requests
from pathlib import Path
multiprocessing.set_start_method('spawn', force=True)
@ -73,7 +69,6 @@ def matte_key(output, file_path,
results_dict = manager.dict()
frames_dict = manager.dict()
info = ffmpeg.probe(file_path)
cmd = [
"ffprobe",
@ -149,7 +144,9 @@ def matte_key(output, file_path,
proc.stdin.write(frame.tostring())
frame_counter = frame_counter + 1
if frame_counter >= total_frames:
p.join()
for w in workers:
w.join()
@ -164,167 +161,3 @@ def matte_key(output, file_path,
proc.stdin.close()
proc.wait()
return
def transparentgif(output, file_path,
worker_nodes,
gpu_batchsize,
model_name,
frame_limit=-1,
prefetched_batches=4,
framerate=-1):
temp_dir = tempfile.TemporaryDirectory()
tmpdirname = Path(temp_dir.name)
temp_file = os.path.abspath(os.path.join(tmpdirname, "matte.mp4"))
matte_key(temp_file, file_path,
worker_nodes,
gpu_batchsize,
model_name,
frame_limit,
prefetched_batches,
framerate)
cmd = [
'ffmpeg', '-y', '-i', file_path, '-i', temp_file, '-filter_complex',
'[1][0]scale2ref[mask][main];[main][mask]alphamerge=shortest=1,fps=10,split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse',
'-shortest', output
]
sp.run(cmd)
print("Process finished")
return
def transparentgifwithbackground(output, overlay, file_path,
worker_nodes,
gpu_batchsize,
model_name,
frame_limit=-1,
prefetched_batches=4,
framerate=-1):
temp_dir = tempfile.TemporaryDirectory()
tmpdirname = Path(temp_dir.name)
temp_file = os.path.abspath(os.path.join(tmpdirname, "matte.mp4"))
matte_key(temp_file, file_path,
worker_nodes,
gpu_batchsize,
model_name,
frame_limit,
prefetched_batches,
framerate)
print("Starting alphamerge")
cmd = [
'ffmpeg', '-y', '-i', file_path, '-i', temp_file, '-i', overlay, '-filter_complex',
'[1][0]scale2ref[mask][main];[main][mask]alphamerge=shortest=1[fg];[2][fg]overlay=(main_w-overlay_w)/2:(main_h-overlay_h)/2:format=auto,fps=10,split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse',
'-shortest', output
]
sp.run(cmd)
print("Process finished")
try:
temp_dir.cleanup()
except PermissionError:
pass
return
def transparentvideo(output, file_path,
worker_nodes,
gpu_batchsize,
model_name,
frame_limit=-1,
prefetched_batches=4,
framerate=-1):
temp_dir = tempfile.TemporaryDirectory()
tmpdirname = Path(temp_dir.name)
temp_file = os.path.abspath(os.path.join(tmpdirname, "matte.mp4"))
matte_key(temp_file, file_path,
worker_nodes,
gpu_batchsize,
model_name,
frame_limit,
prefetched_batches,
framerate)
print("Starting alphamerge")
cmd = [
'ffmpeg', '-y', '-i', file_path, '-i', temp_file, '-filter_complex',
'[1][0]scale2ref[mask][main];[main][mask]alphamerge=shortest=1', '-c:v', 'qtrle', '-shortest', output
]
sp.run(cmd)
print("Process finished")
try:
temp_dir.cleanup()
except PermissionError:
pass
return
def transparentvideoovervideo(output, overlay, file_path,
worker_nodes,
gpu_batchsize,
model_name,
frame_limit=-1,
prefetched_batches=4,
framerate=-1):
temp_dir = tempfile.TemporaryDirectory()
tmpdirname = Path(temp_dir.name)
temp_file = os.path.abspath(os.path.join(tmpdirname, "matte.mp4"))
matte_key(temp_file, file_path,
worker_nodes,
gpu_batchsize,
model_name,
frame_limit,
prefetched_batches,
framerate)
print("Starting alphamerge")
cmd = [
'ffmpeg', '-y', '-i', file_path, '-i', temp_file, '-i', overlay, '-filter_complex',
'[1][0]scale2ref[mask][main];[main][mask]alphamerge=shortest=1[vid];[vid][2:v]scale2ref[fg][bg];[bg][fg]overlay=shortest=1[out]', '-map', '[out]', '-shortest', output
]
sp.run(cmd)
print("Process finished")
try:
temp_dir.cleanup()
except PermissionError:
pass
return
def transparentvideooverimage(output, overlay, file_path,
worker_nodes,
gpu_batchsize,
model_name,
frame_limit=-1,
prefetched_batches=4,
framerate=-1):
temp_dir = tempfile.TemporaryDirectory()
tmpdirname = Path(temp_dir.name)
temp_file = os.path.abspath(os.path.join(tmpdirname, "matte.mp4"))
matte_key(temp_file, file_path,
worker_nodes,
gpu_batchsize,
model_name,
frame_limit,
prefetched_batches,
framerate)
print("Scale image")
temp_image = os.path.abspath("%s/new.jpg" % tmpdirname)
cmd = [
'ffmpeg', '-y', '-i', overlay, '-i', file_path, '-filter_complex',
'scale2ref[img][vid];[img]setsar=1;[vid]nullsink', '-q:v', '2', temp_image
]
sp.run(cmd)
print("Starting alphamerge")
cmd = [
'ffmpeg', '-y', '-i', temp_image, '-i', file_path, '-i', temp_file, '-filter_complex',
'[0:v]scale2ref=oh*mdar:ih[bg];[1:v]scale2ref=oh*mdar:ih[fg];[bg][fg]overlay=(W-w)/2:(H-h)/2:shortest=1[out]',
'-map', '[out]', '-shortest', output
]
sp.run(cmd)
print("Process finished")
try:
temp_dir.cleanup()
except PermissionError:
pass
return