JIT-Masker

[arxiv]

Getting started

Install requirements with

pip install -r requirements

Install detectron separately with

pip install detectron2==0.1.3 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu101/torch1.5/index.html

To use U2Net, grab the pretrained models for U-2-Net. Either the full model or the small p-variant works - just put it in /saved_models/u2net/u2net.pth or /saved_models/u2netp/u2netp.pth.

A pretrained JITNet model for this repository can be found here. Put this model at /saved_models/jitnet/jitnet.pth.

The main script to run is infer_video.py. For example, to run JIT-Masker with your laptop webcam, run python infer_video.py -i 0.

At the moment you can choose between using a video, your own camera, or any IP-accessible camera (for example, an Android phone with IP-webcam).

The script terminates whenever the stream terminates, or when the button q is pressed.

Training (Optional)

Run python train.py to train a JITNet model on the Supervisely dataset. The Supervisely dataset is available on their official website and should be downloaded to data/supervisely.

Inference