Posible infinite loop on some videos when removing background

[nadermx]

In trying to kill two birds with one stone, basically this is the issue. If anyone looks at the code, specifically at the utilities function, https://github.com/nadermx/backgroundremover/blob/main/backgroundremover/utilities.py

There are two while loops 1 and 2 .

The while loop with a sleep function to wait for results to be added to the results_dict. This can cause the program to hang indefinitely if a result is never added to the dictionary or if the loop is unable to exit for some other reason.

A better approach would be to use a synchronization primitive, such as a multiprocessing.Queue or a multiprocessing.Event, to signal when a result has been added to the dictionary.

This would probably increase, or give a route to increase the video processing speed and GPU usage.

I have a primitive updated utilities.by below doing this, haven't had much time to debug specifically yet, but the matte_key seems to work, but the transparent video no.

import os
import math
import torch.multiprocessing as multiprocessing
import subprocess as sp
import time
import ffmpeg
import numpy as np
import torch
from .bg import DEVICE, Net, iter_frames, remove_many
import shlex
import tempfile
import requests
from pathlib import Path

multiprocessing.set_start_method('spawn', force=True)

def worker(worker_nodes,
           worker_index,
           result_dict,
           model_name,
           gpu_batchsize,
           total_frames,
           frames_dict,
           finished_event):
    print(F"WORKER {worker_index} ONLINE")

    output_index = worker_index + 1
    base_index = worker_index * gpu_batchsize
    net = Net(model_name)
    script_net = None
    for fi in (list(range(base_index + i * worker_nodes * gpu_batchsize,
                          min(base_index + i * worker_nodes * gpu_batchsize + gpu_batchsize, total_frames)))
               for i in range(math.ceil(total_frames / worker_nodes / gpu_batchsize))):
        if not fi:
            break

        # are we processing frames faster than the frame ripper is saving them?
        last = fi[-1]
        while last not in frames_dict:
            if finished_event.is_set():
                return
            time.sleep(0.1)

        input_frames = [frames_dict[index] for index in fi]
        if script_net is None:
            script_net = torch.jit.trace(net,
                                         torch.as_tensor(np.stack(input_frames), dtype=torch.float32, device=DEVICE))

        result_dict[output_index] = remove_many(input_frames, script_net)

        # clean up the frame buffer
        for fdex in fi:
            del frames_dict[fdex]
        output_index += worker_nodes

def capture_frames(file_path, frames_dict, prefetched_samples, total_frames, finished_event):
    print(F"WORKER FRAMERIPPER ONLINE")
    for idx, frame in enumerate(iter_frames(file_path)):
        frames_dict[idx] = frame
        while len(frames_dict) > prefetched_samples:
            if finished_event.is_set():
                return
            time.sleep(0.1)
        if idx >= total_frames - 1:
            finished_event.set()
            return

def matte_key(output, file_path,
              worker_nodes,
              gpu_batchsize,
              model_name,
              frame_limit=-1,
              prefetched_batches=4,
              framerate=-1):
    manager = multiprocessing.Manager()
    finished_event = manager.Event()

    results_dict = manager.dict()
    frames_dict = manager.dict()

    info = ffmpeg.probe(file_path)
    cmd = [
        "ffprobe",
        "-v",
        "error",
        "-select_streams",
        "v:0",
        "-count_packets",
        "-show_entries",
        "stream=nb_read_packets",
        "-of",
        "csv=p=0",
        file_path
    ]
    framerate_output = sp.check_output(cmd, universal_newlines=True)
    total_frames = int(framerate_output)
    if frame_limit != -1:
        total_frames = min(frame_limit, total_frames)

    fr = info["streams"][0]["r_frame_rate"]

    if framerate == -1:
        print(F"FRAME RATE DETECTED: {fr} (if this looks wrong, override the frame rate)")
        framerate = math.ceil(eval(fr))

    print(F"FRAME RATE: {framerate} TOTAL FRAMES: {total_frames}")

    p = multiprocessing.Process(target=capture_frames,
                                args=(file_path, frames_dict, gpu_batchsize * prefetched_batches, total_frames,
                                      finished_event))
    p.start()

    # note I am deliberately not using pool
    # we can't trust it to run all the threads concurrently (or at all)
    workers = [multiprocessing.Process(target=worker,
                                       args=(worker_nodes, wn, results_dict, model_name, gpu_batchsize, total_frames,
                                             frames_dict, finished_event))
               for wn in range(worker_nodes)]
    for w in workers:
        w.start()

    command = None
    proc = None
    frame_counter = 0
    for i in range(math.ceil(total_frames / worker_nodes)):
        for wx in range(worker_nodes):

            hash_index = i * worker_nodes + 1 + wx

            while hash_index not in results_dict:
                if finished_event.is_set():
                    p.join()
                    for w in workers:
                        w.join()
                    return
                time.sleep(0.1)

            frames = results_dict[hash_index]
            # don't block access to it anymore
            del results_dict[hash_index]

            for frame in frames:
                if command is None:
                    command = ['ffmpeg',
                               '-y',
                               '-f', 'rawvideo',
                               '-vcodec', 'rawvideo',
                               '-s', F"{frame.shape[1]}x320",
                               '-pix_fmt', 'gray',
                               '-r', F"{framerate}",
                               '-i', '-',
                               '-an',
                               '-vcodec', 'mpeg4',
                               '-b:v', '2000k',
                               '%s' % output]

                    proc = sp.Popen(command, stdin=sp.PIPE)

                proc.stdin.write(frame.tostring())
                frame_counter = frame_counter + 1

                if frame_counter >= total_frames:
                    finished_event.set()
                    p.join()
                    for w in workers:
                        w.join()
                    proc.stdin.close()
                    proc.wait()
                    print(F"FINISHED ALL FRAMES ({total_frames})!")
                    return

    finished_event.set()
    p.join()
    for w in workers:
        w.join()
    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 = "nice -10 ffmpeg -y -i %s -i %s -filter_complex '[1][0]scale2ref[mask][main];[main][mask]alphamerge=shortest=1,fps=10,split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse' -shortest %s" % (
        file_path, temp_file, output)
    sp.run(shlex.split(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 = "nice -10 ffmpeg -y -i %s -i %s -i %s -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 %s" % (
        file_path, temp_file, overlay, output)
    sp.run(shlex.split(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 = "nice -10 ffmpeg -y -nostats -loglevel 0 -i %s -i %s -filter_complex '[1][0]scale2ref[mask][main];[main][mask]alphamerge=shortest=1' -c:v qtrle -shortest %s" % (
        file_path, temp_file, output)
    process = sp.Popen(cmd, shell=True, stdout=sp.PIPE, stderr=sp.PIPE)
    stdout, stderr = process.communicate()
    print('after call')

    if stderr:
        return "ERROR: %s" % stderr.decode("utf-8")
    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 = "nice -10 ffmpeg -y -i %s -i %s -i %s -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 %s" % (
        file_path, temp_file, overlay, output)
    sp.run(shlex.split(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 = "nice -10 ffmpeg -i %s -i %s -filter_complex 'scale2ref[img][vid];[img]setsar=1;[vid]nullsink' -q:v 2 %s" % (
        overlay, file_path, temp_image)
    sp.run(shlex.split(cmd))
    print("Starting alphamerge")
    cmd = "nice -10 ffmpeg -y -i %s -i %s -i %s -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 %s" % (
        temp_image, file_path, temp_file, output)
    sp.run(shlex.split(cmd))
    print("Process finished")
    try:
        temp_dir.cleanup()
    except PermissionError:
        pass
    return

def download_files_from_github(path, model_name):
    if model_name not in ["u2net", "u2net_human_seg"]:
        print("Invalid model name, please use 'u2net' or 'u2net_human_seg'")
        return
    print(f"downloading model [{model_name}] to {path} ...")
    urls = []
    if model_name == "u2net":
        urls = ['https://github.com/nadermx/backgroundremover/raw/main/models/u2aa',
                'https://github.com/nadermx/backgroundremover/raw/main/models/u2ab',
                'https://github.com/nadermx/backgroundremover/raw/main/models/u2ac',
                'https://github.com/nadermx/backgroundremover/raw/main/models/u2ad']
    elif model_name == "u2net_human_seg":
        urls = ['https://github.com/nadermx/backgroundremover/raw/main/models/u2haa',
                'https://github.com/nadermx/backgroundremover/raw/main/models/u2hab',
                'https://github.com/nadermx/backgroundremover/raw/main/models/u2hac',
                'https://github.com/nadermx/backgroundremover/raw/main/models/u2had']

    try:
        os.makedirs(os.path.expanduser("~/.u2net"), exist_ok=True)
    except Exception as e:
        print(f"Error creating directory: {e}")
        return

    try:

        with open(path, 'wb') as out_file:
            for i, url in enumerate(urls):
                print(f'downloading part {i + 1} of {model_name}')
                part_content = requests.get(url)
                out_file.write(part_content.content)
                print(f'finished downloading part {i + 1} of {model_name}')

    finally:
        print()

[nadermx]

So that leaves me to believe that the file might be right, but now perhaps since it is faster it is not reading the mov correctly during the ffmpeg, since that is where it is erroring out

python -m backgroundremover.cmd.cli -i "examplefiles/shave.mp4" -tg -o "crap.mov"
FRAME RATE DETECTED: 155/6 (if this looks wrong, override the frame rate)
FRAME RATE: 26 TOTAL FRAMES: 161
WORKER 0 ONLINE
WORKER FRAMERIPPER ONLINE
ffmpeg version 4.4.3-0ubuntu1~20.04.sav5 Copyright (c) 2000-2022 the FFmpeg developers
  built with gcc 9 (Ubuntu 9.4.0-1ubuntu1~20.04.1)
  configuration: --prefix=/usr --extra-version='0ubuntu1~20.04.sav5' --toolchain=hardened --libdir=/usr/lib/x86_64-linux-gnu --incdir=/usr/include/x86_64-linux-gnu --arch=amd64 --enable-gpl --disable-stripping --enable-gnutls --enable-ladspa --enable-libaom --enable-libass --enable-libbluray --enable-libbs2b --enable-libcaca --enable-libcdio --enable-libcodec2 --enable-libflite --enable-libfontconfig --enable-libfreetype --enable-libfribidi --enable-libgme --enable-libgsm --enable-libjack --enable-libmp3lame --enable-libmysofa --enable-libopenjpeg --enable-libopenmpt --enable-libopus --enable-libpulse --enable-librabbitmq --enable-librubberband --enable-libshine --enable-libsnappy --enable-libsoxr --enable-libspeex --enable-libsrt --enable-libssh --enable-libtheora --enable-libtwolame --enable-libvidstab --enable-libvorbis --enable-libvpx --enable-libwebp --enable-libx265 --enable-libxml2 --enable-libxvid --enable-libzmq --enable-libzvbi --enable-lv2 --enable-omx --enable-openal --enable-opencl --enable-opengl --enable-sdl2 --enable-pocketsphinx --enable-librsvg --enable-libzimg --enable-libdc1394 --enable-libdrm --enable-libiec61883 --enable-chromaprint --enable-frei0r --enable-libx264 --enable-shared
  libavutil      56. 70.100 / 56. 70.100
  libavcodec     58.134.100 / 58.134.100
  libavformat    58. 76.100 / 58. 76.100
  libavdevice    58. 13.100 / 58. 13.100
  libavfilter     7.110.100 /  7.110.100
  libswscale      5.  9.100 /  5.  9.100
  libswresample   3.  9.100 /  3.  9.100
  libpostproc    55.  9.100 / 55.  9.100
Input #0, rawvideo, from 'pipe:':
  Duration: N/A, start: 0.000000, bitrate: 37872 kb/s
  Stream #0:0: Video: rawvideo (Y800 / 0x30303859), gray, 569x320, 37872 kb/s, 26 tbr, 26 tbn, 26 tbc
Stream mapping:
  Stream #0:0 -> #0:0 (rawvideo (native) -> mpeg4 (native))
[swscaler @ 0x5599304047c0] Warning: data is not aligned! This can lead to a speed loss
Output #0, mp4, to '/tmp/tmprn_fgv05/matte.mp4':
  Metadata:
    encoder         : Lavf58.76.100
  Stream #0:0: Video: mpeg4 (mp4v / 0x7634706D), yuv420p(tv, progressive), 569x320, q=2-31, 2000 kb/s, 26 fps, 13312 tbn
    Metadata:
      encoder         : Lavc58.134.100 mpeg4
    Side data:
      cpb: bitrate max/min/avg: 0/0/2000000 buffer size: 0 vbv_delay: N/A
ffmpeg version 4.4.3-0ubuntu1~20.04.sav5 Copyright (c) 2000-2022 the FFmpeg developers=1.07x    
  built with gcc 9 (Ubuntu 9.4.0-1ubuntu1~20.04.1)
  configuration: --prefix=/usr --extra-version='0ubuntu1~20.04.sav5' --toolchain=hardened --libdir=/usr/lib/x86_64-linux-gnu --incdir=/usr/include/x86_64-linux-gnu --arch=amd64 --enable-gpl --disable-stripping --enable-gnutls --enable-ladspa --enable-libaom --enable-libass --enable-libbluray --enable-libbs2b --enable-libcaca --enable-libcdio --enable-libcodec2 --enable-libflite --enable-libfontconfig --enable-libfreetype --enable-libfribidi --enable-libgme --enable-libgsm --enable-libjack --enable-libmp3lame --enable-libmysofa --enable-libopenjpeg --enable-libopenmpt --enable-libopus --enable-libpulse --enable-librabbitmq --enable-librubberband --enable-libshine --enable-libsnappy --enable-libsoxr --enable-libspeex --enable-libsrt --enable-libssh --enable-libtheora --enable-libtwolame --enable-libvidstab --enable-libvorbis --enable-libvpx --enable-libwebp --enable-libx265 --enable-libxml2 --enable-libxvid --enable-libzmq --enable-libzvbi --enable-lv2 --enable-omx --enable-openal --enable-opencl --enable-opengl --enable-sdl2 --enable-pocketsphinx --enable-librsvg --enable-libzimg --enable-libdc1394 --enable-libdrm --enable-libiec61883 --enable-chromaprint --enable-frei0r --enable-libx264 --enable-shared
  libavutil      56. 70.100 / 56. 70.100
  libavcodec     58.134.100 / 58.134.100
  libavformat    58. 76.100 / 58. 76.100
  libavdevice    58. 13.100 / 58. 13.100
  libavfilter     7.110.100 /  7.110.100
  libswscale      5.  9.100 /  5.  9.100
  libswresample   3.  9.100 /  3.  9.100
  libpostproc    55.  9.100 / 55.  9.100
Input #0, mov,mp4,m4a,3gp,3g2,mj2, from '/home/imageeditor/backgroundremover/examplefiles/shave.mp4':
  Metadata:
    major_brand     : isom
    minor_version   : 512
    compatible_brands: isomiso6iso2avc1mp41
    encoder         : Lavf59.27.100
  Duration: 00:00:06.42, start: 0.000000, bitrate: 474 kb/s
  Stream #0:0(und): Video: h264 (Main) (avc1 / 0x31637661), yuv420p(tv, smpte170m/bt709/bt709), 854x480 [SAR 1:1 DAR 427:240], 321 kb/s, 25.79 fps, 25.83 tbr, 90k tbn, 51.59 tbc (default)
    Metadata:
      handler_name    : VideoHandler
      vendor_id       : [0][0][0][0]
  Stream #0:1(und): Audio: aac (LC) (mp4a / 0x6134706D), 44100 Hz, mono, fltp, 93 kb/s (default)
    Metadata:
      handler_name    : SoundHandler
      vendor_id       : [0][0][0][0]
[mov,mp4,m4a,3gp,3g2,mj2 @ 0x564677c20f80] moov atom not found
/tmp/tmprn_fgv05/matte.mp4: Invalid data found when processing input
Process finished
frame=  158 fps= 21 q=2.0 Lsize=     357kB time=00:00:06.03 bitrate= 484.2kbits/s speed=0.787x    
video:355kB audio:0kB subtitle:0kB other streams:0kB global headers:0kB muxing overhead: 0.430353%

Specificially this part

[mov,mp4,m4a,3gp,3g2,mj2 @ 0x564677c20f80] moov atom not found
/tmp/tmprn_fgv05/matte.mp4: Invalid data found when processing input

And this superuser question seems to be among the same issues.

[nadermx]

It seems this isn't putting the frames in the correct order, so back toe drawing board

[marih0106]

@nadermx How it works?

[nadermx]

What do you mean @Sam2003ds? The project works fine. I have had it, perhaps due to this I think, crash on me on occasion after doing a bunch of video's.

After looking at how I had set up the current frame ripping, there is a chance that if it processes frames faster than it can write it, that the while loop could get stuck. What I pasted here, is an updated utilites.py function.

In this I am trying to place the frames in a queue to mitigate the while loops. The one I posted makes a matte_key file, but not sure if it makes it correctly. Since it fails when I input it into ffmpeg and that might be since it seems it may be placing the frames not back in order causing the .mov to be invail since they missing the meta info that is on there in the way the current repo does it.

[marih0106]

@nadermx what are the codes for?

[nadermx]

The project works fine, this is a bug fix that occurs occasionally. I updated my previous comment to you if you want to read it should explain.

[marih0106]

thanks

[nadermx]

Going to close this, as I think it was related to the input files with special charectors. Will keep testing and reopen if it occurs.