Initial clean project import

2026-06-19 18:41:41 +08:00
commit a13b804c7a
1306 changed files with 220568 additions and 0 deletions
@@ -0,0 +1,50 @@
+## MODNet - Custom Portrait Video Matting Demo
+This is a MODNet portrait video matting demo that allows you to process custom videos.
+
+### 1. Requirements
+The basic requirements for this demo are:
+- Ubuntu System
+- Python 3+
+
+
+### 2. Introduction
+We use ~400 unlabeled video clips (divided into ~50,000 frames) downloaded from the internet to perform SOC to adapt MODNet to the video domain. **Nonetheless, due to insufficient labeled training data (~3k labeled foregrounds), our model may still make errors in portrait semantics estimation under challenging scenes.** Besides, this demo does not currently support the OFD trick.
+
+
+For a better experience, please make sure your videos satisfy:
+
+*   the portrait and background are distinguishable, <i>i.e.</i>, are not similar
+*   captured in soft and bright ambient lighting
+*   the contents do not move too fast
+
+### 3. Run Demo
+We recommend creating a new conda virtual environment to run this demo, as follow:
+
+1. Clone the MODNet repository:
+    ```
+    git clone https://github.com/ZHKKKe/MODNet.git
+    cd MODNet
+    ```
+
+2. Download the pre-trained model from this [link](https://drive.google.com/file/d/1Nf1ZxeJZJL8Qx9KadcYYyEmmlKhTADxX/view?usp=sharing) and put it into the folder `MODNet/pretrained/`.
+
+
+3. Create a conda virtual environment named `modnet` (if it doesn't exist) and activate it. Here we use `python=3.6` as an example:
+     ```
+    conda create -n modnet python=3.6
+    source activate modnet
+    ```
+
+4. Install the required python dependencies (please make sure your CUDA version is supported by the PyTorch version installed):
+    ```
+    pip install -r demo/video_matting/custom/requirements.txt
+    ```
+
+5. Execute the main code:
+    ```
+    python -m demo.video_matting.custom.run --video YOUR_VIDEO_PATH
+    ```
+    where `YOUR_VIDEO_PATH` is the specific path of your video.  
+    There are some optional arguments:
+     - `--result-type (default=fg)` : fg - save the alpha matte; fg - save the foreground
+     - `--fps (default=30)` : fps of the result video
@@ -0,0 +1,6 @@
+numpy
+Pillow
+opencv-python
+torch >= 1.0.0
+torchvision
+tqdm
@@ -0,0 +1,114 @@
+import os
+import cv2
+import argparse
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+
+import torch
+import torch.nn as nn
+import torchvision.transforms as transforms
+
+from src.models.modnet import MODNet
+
+
+torch_transforms = transforms.Compose(
+    [
+        transforms.ToTensor(),
+        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+    ]
+)
+
+
+def matting(video, result, alpha_matte=False, fps=30):
+    # video capture
+    vc = cv2.VideoCapture(video)
+
+    if vc.isOpened():
+        rval, frame = vc.read()
+    else:
+        rval = False
+
+    if not rval:
+        print('Failed to read the video: {0}'.format(video))
+        exit()
+
+    num_frame = vc.get(cv2.CAP_PROP_FRAME_COUNT)
+    h, w = frame.shape[:2]
+    if w >= h:
+        rh = 512
+        rw = int(w / h * 512)
+    else:
+        rw = 512
+        rh = int(h / w * 512)
+    rh = rh - rh % 32
+    rw = rw - rw % 32
+
+    # video writer
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    video_writer = cv2.VideoWriter(result, fourcc, fps, (w, h))
+
+    print('Start matting...')
+    with tqdm(range(int(num_frame)))as t:
+        for c in t:
+            frame_np = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            frame_np = cv2.resize(frame_np, (rw, rh), cv2.INTER_AREA)
+
+            frame_PIL = Image.fromarray(frame_np)
+            frame_tensor = torch_transforms(frame_PIL)
+            frame_tensor = frame_tensor[None, :, :, :]
+            if GPU:
+                frame_tensor = frame_tensor.cuda()
+
+            with torch.no_grad():
+                _, _, matte_tensor = modnet(frame_tensor, True)
+
+            matte_tensor = matte_tensor.repeat(1, 3, 1, 1)
+            matte_np = matte_tensor[0].data.cpu().numpy().transpose(1, 2, 0)
+            if alpha_matte:
+                view_np = matte_np * np.full(frame_np.shape, 255.0)
+            else:
+                view_np = matte_np * frame_np + (1 - matte_np) * np.full(frame_np.shape, 255.0)
+            view_np = cv2.cvtColor(view_np.astype(np.uint8), cv2.COLOR_RGB2BGR)
+            view_np = cv2.resize(view_np, (w, h))
+            video_writer.write(view_np)
+
+            rval, frame = vc.read()
+            c += 1
+
+    video_writer.release()
+    print('Save the result video to {0}'.format(result))
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--video', type=str, required=True, help='input video file')
+    parser.add_argument('--result-type', type=str, default='fg', choices=['fg', 'matte'], 
+                        help='matte - save the alpha matte; fg - save the foreground')
+    parser.add_argument('--fps', type=int, default=30, help='fps of the result video')
+
+    print('Get CMD Arguments...')
+    args = parser.parse_args()
+
+    if not os.path.exists(args.video):
+        print('Cannot find the input video: {0}'.format(args.video))
+        exit()
+
+    print('Load pre-trained MODNet...')
+    pretrained_ckpt = './pretrained/modnet_webcam_portrait_matting.ckpt'
+    modnet = MODNet(backbone_pretrained=False)
+    modnet = nn.DataParallel(modnet)
+
+    GPU = True if torch.cuda.device_count() > 0 else False
+    if GPU:
+        print('Use GPU...')
+        modnet = modnet.cuda()
+        modnet.load_state_dict(torch.load(pretrained_ckpt))
+    else:
+        print('Use CPU...')
+        modnet.load_state_dict(torch.load(pretrained_ckpt, map_location=torch.device('cpu')))
+    modnet.eval()
+
+    result = os.path.splitext(args.video)[0] + '_{0}.mp4'.format(args.result_type)
+    alpha_matte = True if args.result_type == 'matte' else False
+    matting(args.video, result, alpha_matte, args.fps)