From bcf2f233e8d5639e2b14dc5a031f4100a81ae4f1 Mon Sep 17 00:00:00 2001 From: Dylan Turner Date: Fri, 11 Nov 2022 16:50:24 -0600 Subject: [PATCH] Implemented video file removal (#22) --- README.md | 10 ++++++++++ bgrm/bgrm.py | 33 +++++++++++++++++++++++++-------- bgrm/cam.py | 17 ++++++++++------- bgrm/settings.py | 26 +++++++++++++++++++++++++- bgrm/vid_file.py | 47 +++++++++++++++++++++++++++++++++++++++++++++++ pyproject.toml | 2 +- 6 files changed, 118 insertions(+), 17 deletions(-) create mode 100644 bgrm/vid_file.py diff --git a/README.md b/README.md index e475fdb..62962f4 100644 --- a/README.md +++ b/README.md @@ -10,8 +10,12 @@ Using OpenCV and a v4l2loopback device (basically a virtual webcam you can write Note, this will work anywhere WebCams are used, not just Teams +Now, the program can also be used to remove backgrounds from video files and save them as video files as well! + ## How to Use +### WebCam Replacement + Dependencies: - python >= 3.8 (3.10 is what's supported officially) - pip @@ -32,6 +36,12 @@ Then, you can run: - Run with `python -m bgrm ` (use `--help` to see all options) - Example: `python -m bgrm -b ~/Pictures/Wallpapers/ni-skyline-wallpaper.png -w 320 -H 240 -s 2.0` +### File Replacement + +You can also remove the background from video files. It works just like the WebCam, but instead of setting the `--camera` cli arg, you call the program like this: + +`python -m bgrm --file_mode -i -o ` + ## Build from Repo You can also build the package yourself from source (or grab the latest version from the releases tab) diff --git a/bgrm/bgrm.py b/bgrm/bgrm.py index 3d231b4..4dfd4fe 100755 --- a/bgrm/bgrm.py +++ b/bgrm/bgrm.py @@ -3,18 +3,35 @@ # Author: Dylan Turner # Description: Main loop for bgrm app +from numpy import shape from bgrm.cam import Cam from bgrm.virt_cam import VirtCam +from bgrm.vid_file import VidFile from bgrm.settings import AppSettings def main(): settings = AppSettings.from_cli() - virt_dev_name = '/dev/video' + str(settings.virt_dev) - with Cam(settings) as cam, VirtCam(virt_dev_name) as virt_cam: - virt_cam.format(cam, settings) - while True: - _frame, no_bg_frame, stacked_frames = cam.frames() - virt_cam.write(no_bg_frame) - if not cam.display(stacked_frames): - break + + if not settings.file_mode: + virt_dev_name = '/dev/video' + str(settings.virt_dev) + with Cam(settings) as cam, VirtCam(virt_dev_name) as virt_cam: + virt_cam.format(cam, settings) + while True: + _frame, no_bg_frame, stacked_frames = cam.frames() + virt_cam.write(no_bg_frame) + + if not cam.display(stacked_frames): + break + else: + with VidFile(settings) as vid_file: + while True: + frame, no_bg_frame, stacked_frames = vid_file.frames() + + if shape(frame) == (): + break + + vid_file.write(no_bg_frame) + + if not vid_file.display(stacked_frames): + break diff --git a/bgrm/cam.py b/bgrm/cam.py index b9fd2f6..d16e841 100644 --- a/bgrm/cam.py +++ b/bgrm/cam.py @@ -5,7 +5,7 @@ VideoCapture, resize, namedWindow, moveWindow, imshow, waitKey, \ destroyAllWindows, imread, BORDER_CONSTANT, copyMakeBorder, resize, imread, \ GaussianBlur -from numpy import shape +from numpy import shape, uint8 from cvzone import stackImages from cvzone.SelfiSegmentationModule import SelfiSegmentation from time import sleep @@ -17,11 +17,11 @@ def __init__(self, settings): # Set up WebCam access and output self._settings = settings - self._vidFeed = VideoCapture(settings.camera) - self._vidFeed.set(3, settings.screen_width) - self._vidFeed.set(4, settings.screen_height) + self._vid_feed = VideoCapture(settings.camera) + self._vid_feed.set(3, settings.screen_width) + self._vid_feed.set(4, settings.screen_height) - _success, baseFrame = self._vidFeed.read() + _success, baseFrame = self._vid_feed.read() if not _success: print('Failed to read from camera!') quit() @@ -42,12 +42,15 @@ def __enter__(self): return self def __exit__(self, exc_type, exc_value, traceback): - self._vidFeed.release() + self._vid_feed.release() destroyAllWindows() # Return regular camera frame and with the bg removal applied, as well as them put together def frames(self): - _success, frame = self._vidFeed.read() + success, frame = self._vid_feed.read() + + if not success: + return (None, None, None) if shape(self._bg_img) == (): no_bg_frame = self._segmentor.removeBG( diff --git a/bgrm/settings.py b/bgrm/settings.py index 1acc597..fba1362 100644 --- a/bgrm/settings.py +++ b/bgrm/settings.py @@ -17,6 +17,10 @@ class AppSettings: bg_img: str blur: bool disable_win: bool + file_mode: bool + input_file: str + output_file: str + fps: int @staticmethod def from_cli(): @@ -32,7 +36,8 @@ def from_cli(): return AppSettings( virt_dev, args.camera, args.width, args.height, args.scale, - args.thresh, fill_col, args.bg, args.blur, args.disable_window + args.thresh, fill_col, args.bg, args.blur, args.disable_window, + args.file_mode, args.input, args.output, args.fps ) # Get video settings from command line @@ -91,6 +96,25 @@ def cli_args(): help = 'List of R, G, B to replace background with', ) + # File versions + parser.add_argument( + '--file_mode', + help = 'Remove the background from video files instead of cameras (overrides --camera)', + action = 'store_true' + ) + parser.add_argument( + '-i', '--input', type = str, default = '', + help = 'Only used with file mode. Input file to remove background', + ) + parser.add_argument( + '-o', '--output', type = str, default = '', + help = 'Only used with file mode. Output file to store removed background feed.' + ) + parser.add_argument( + '-f', '--fps', type = int, default = 30, + help = 'Only used with file mode. Output file fps' + ) + # Mostly useless window options parser.add_argument( '-s', '--scale', type = float, default = 0.5, diff --git a/bgrm/vid_file.py b/bgrm/vid_file.py new file mode 100644 index 0000000..09484c6 --- /dev/null +++ b/bgrm/vid_file.py @@ -0,0 +1,47 @@ +# Author: Dylan Turner +# Description: File-to-file version of Cam class + +from cv2 import VideoCapture, VideoWriter, VideoWriter_fourcc, namedWindow, destroyAllWindows +from cvzone.SelfiSegmentationModule import SelfiSegmentation +from bgrm.cam import Cam + +WIN_TITLE = 'Conversion' + +class VidFile(Cam): + def __init__(self, settings): + self._settings = settings + + # Set up file input + self._vid_feed = VideoCapture(settings.input_file) + _success, base_frame = self._vid_feed.read() + height, width, self.channels = base_frame.shape + self._settings.screen_width = width + self._settings.screen_height = height + settings.screen_width = width + settings.screen_height = height + + if not settings.disable_win: + namedWindow(WIN_TITLE) + + # Set up removal + self._segmentor = SelfiSegmentation() + + if settings.bg_img != '': + self._bg_img = self._create_correctly_sized_bg() + else: + self._bg_img = None + + # Set up file output + self._out_feed = VideoWriter( + settings.output_file, VideoWriter_fourcc('M', 'J', 'P', 'G'), + settings.fps, (width, height) + ) + + def __exit__(self, exc_type, exc_value, traceback): + self._vid_feed.release() + self._out_feed.release() + destroyAllWindows() + + def write(self, frame): + self._out_feed.write(frame) + diff --git a/pyproject.toml b/pyproject.toml index f9caff0..844b651 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "bgrm" -version = "8" +version = "9" description = "Remove backgrounds from video feeds in your web cam applications." authors = [ "Dylan Turner " ] license = "GPL-3.0-only"