Implemented video file removal (#22)

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 <options>` (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 <input file> -o <output file> <other options>`
+
 ## Build from Repo
 
 You can also build the package yourself from source (or grab the latest version from the releases tab)

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
 

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(

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,

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)
+

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 <[email protected]>" ]
 license = "GPL-3.0-only"