Initial clean project import
This commit is contained in:
@@ -0,0 +1,52 @@
|
||||
## MODNet - WebCam-Based Portrait Video Matting Demo
|
||||
This is a MODNet portrait video matting demo based on WebCam. It will call your local WebCam and display the matting results in real time. The demo can run under CPU or GPU.
|
||||
|
||||
### 1. Requirements
|
||||
The basic requirements for this demo are:
|
||||
- Ubuntu System
|
||||
- WebCam
|
||||
- Python 3+
|
||||
|
||||
**NOTE**: If your device does not satisfy the above conditions, please try our [online Colab demo](https://colab.research.google.com/drive/1Pt3KDSc2q7WxFvekCnCLD8P0gBEbxm6J?usp=sharing).
|
||||
|
||||
|
||||
### 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, which will be provided soon.
|
||||
|
||||
For a better experience, please:
|
||||
|
||||
* make sure the portrait and background are distinguishable, <i>i.e.</i>, are not similar
|
||||
* run in soft and bright ambient lighting
|
||||
* do not be too close or too far from the WebCam
|
||||
* 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/webcam/requirements.txt
|
||||
```
|
||||
|
||||
5. Execute the main code:
|
||||
```
|
||||
python -m demo.video_matting.webcam.run
|
||||
```
|
||||
|
||||
### 4. Acknowledgement
|
||||
We thank [@tkianai](https://github.com/tkianai) and [@mazhar004](https://github.com/mazhar004) for their contributions to making this demo available for CPU use.
|
||||
@@ -0,0 +1,5 @@
|
||||
numpy
|
||||
Pillow
|
||||
opencv-python
|
||||
torch >= 1.0.0
|
||||
torchvision
|
||||
@@ -0,0 +1,67 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
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)),
|
||||
]
|
||||
)
|
||||
|
||||
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()
|
||||
|
||||
print('Init WebCam...')
|
||||
cap = cv2.VideoCapture(0)
|
||||
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
|
||||
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
|
||||
|
||||
print('Start matting...')
|
||||
while(True):
|
||||
_, frame_np = cap.read()
|
||||
frame_np = cv2.cvtColor(frame_np, cv2.COLOR_BGR2RGB)
|
||||
frame_np = cv2.resize(frame_np, (910, 512), cv2.INTER_AREA)
|
||||
frame_np = frame_np[:, 120:792, :]
|
||||
frame_np = cv2.flip(frame_np, 1)
|
||||
|
||||
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)
|
||||
fg_np = matte_np * frame_np + (1 - matte_np) * np.full(frame_np.shape, 255.0)
|
||||
view_np = np.uint8(np.concatenate((frame_np, fg_np), axis=1))
|
||||
view_np = cv2.cvtColor(view_np, cv2.COLOR_RGB2BGR)
|
||||
|
||||
cv2.imshow('MODNet - WebCam [Press \'Q\' To Exit]', view_np)
|
||||
if cv2.waitKey(1) & 0xFF == ord('q'):
|
||||
break
|
||||
|
||||
print('Exit...')
|
||||
Reference in New Issue
Block a user