Ease background matting for videos

.gitignore

@@ -0,0 +1,141 @@
+# Created by .ignore support plugin (hsz.mobi)
+### Python template
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+

0_prepare.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+
+# Load tensorflow models (we need DeepLabV3+ for coarse people segmentation)
+git clone https://github.com/tensorflow/models.git
+
+# Create a docker image with all dependencies
+docker build -t backmatting -f dockerfile .
+
+# Download pretrained models
+wget https://gist.githubusercontent.com/andreyryabtsev/458f7450c630952d1e75e195f94845a0/raw/0b4336ac2a2140ac2313f9966316467e8cd3002a/download.sh
+chmod +x download.sh
+./download.sh

dockerfile

@@ -0,0 +1,24 @@
+# Let image base on ubuntu 16.04
+FROM nvidia/cuda:10.1-cudnn7-devel-ubuntu18.04 
+ENV TZ=Europe/Berlin 
+RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone
+# Install python 3.6
+RUN apt-get update && apt-get install -y \
+                                python3.6-dev\
+                                python3-pip \
+                                python3-tk \
+                                git libgtk2.0-dev
+# Install OpenCV requirements
+RUN apt-get update && apt-get install -y \
+				libopencv-dev \
+				python-opencv
+# Install required python libraries
+ADD requirements.txt . 
+RUN pip3 install --upgrade pip 
+RUN pip3 install --upgrade setuptools 
+
+RUN pip3 install torch==1.4.0+cu100 torchvision==0.5.0+cu100 -f https://download.pytorch.org/whl/torch_stable.html
+RUN pip3 install -r requirements.txt
+# Configure environment variables
+ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64 
+ENV CUDA_HOME=/usr/local/cuda

functions.py

@@ -5,106 +5,106 @@
 
 
 def composite4(fg, bg, a):
-	fg = np.array(fg, np.float32)
-	alpha= np.expand_dims(a / 255,axis=2)
-	im = alpha * fg + (1 - alpha) * bg
-	im = im.astype(np.uint8)
-	return im
+    fg = np.array(fg, np.float32)
+    alpha = np.expand_dims(a / 255, axis=2)
+    im = alpha * fg + (1 - alpha) * bg
+    im = im.astype(np.uint8)
+    return im
 
-def compose_image_withshift(alpha_pred,fg_pred,bg,seg):
 
-    image_sh=torch.zeros(fg_pred.shape).cuda()
+def compose_image_withshift(alpha_pred, fg_pred, bg, seg):
+    image_sh = torch.zeros(fg_pred.shape).cuda()
 
-    for t in range(0,fg_pred.shape[0]):
-        al_tmp=to_image(seg[t,...]).squeeze(2)
-        where = np.array(np.where((al_tmp>0.1).astype(np.float32)))
+    for t in range(0, fg_pred.shape[0]):
+        al_tmp = to_image(seg[t, ...]).squeeze(2)
+        where = np.array(np.where((al_tmp > 0.1).astype(np.float32)))
         x1, y1 = np.amin(where, axis=1)
         x2, y2 = np.amax(where, axis=1)
 
-        #select shift
-        n=np.random.randint(-(y1-10),al_tmp.shape[1]-y2-10)
-        #n positive indicates shift to right
-        alpha_pred_sh=torch.cat((alpha_pred[t,:,:,-n:],alpha_pred[t,:,:,:-n]),dim=2)
-        fg_pred_sh=torch.cat((fg_pred[t,:,:,-n:],fg_pred[t,:,:,:-n]),dim=2)
+        # select shift
+        n = np.random.randint(-(y1 - 10), al_tmp.shape[1] - y2 - 10)
+        # n positive indicates shift to right
+        alpha_pred_sh = torch.cat((alpha_pred[t, :, :, -n:], alpha_pred[t, :, :, :-n]), dim=2)
+        fg_pred_sh = torch.cat((fg_pred[t, :, :, -n:], fg_pred[t, :, :, :-n]), dim=2)
 
-        alpha_pred_sh=(alpha_pred_sh+1)/2
+        alpha_pred_sh = (alpha_pred_sh + 1) / 2
 
-        image_sh[t,...]=fg_pred_sh*alpha_pred_sh + (1-alpha_pred_sh)*bg[t,...]
+        image_sh[t, ...] = fg_pred_sh * alpha_pred_sh + (1 - alpha_pred_sh) * bg[t, ...]
 
     return torch.autograd.Variable(image_sh.cuda())
 
-def get_bbox(mask,R,C):
+
+def get_bbox(mask, R, C):
     where = np.array(np.where(mask))
     x1, y1 = np.amin(where, axis=1)
     x2, y2 = np.amax(where, axis=1)
 
-    bbox_init=[x1,y1,np.maximum(x2-x1,y2-y1),np.maximum(x2-x1,y2-y1)]
-
+    bbox_init = [x1, y1, np.maximum(x2 - x1, y2 - y1), np.maximum(x2 - x1, y2 - y1)]
 
-    bbox=create_bbox(bbox_init,(R,C))
+    bbox = create_bbox(bbox_init, (R, C))
 
     return bbox
 
-def crop_images(crop_list,reso,bbox):
 
-    for i in range(0,len(crop_list)):
-        img=crop_list[i]
-        if img.ndim>=3:
-            img_crop=img[bbox[0]:bbox[0]+bbox[2],bbox[1]:bbox[1]+bbox[3],...]; img_crop=cv2.resize(img_crop,reso)
+def crop_images(crop_list, reso, bbox):
+    for i in range(len(crop_list)):
+        img = crop_list[i]
+        if img.ndim >= 3:
+            img_crop = img[bbox[0]:bbox[0] + bbox[2], bbox[1]:bbox[1] + bbox[3], ...]
+            img_crop = cv2.resize(img_crop, reso)
         else:
-            img_crop=img[bbox[0]:bbox[0]+bbox[2],bbox[1]:bbox[1]+bbox[3]]; img_crop=cv2.resize(img_crop,reso)
-        crop_list[i]=img_crop
+            img_crop = img[bbox[0]:bbox[0] + bbox[2], bbox[1]:bbox[1] + bbox[3]]
+            img_crop = cv2.resize(img_crop, reso)
+        crop_list[i] = img_crop
 
     return crop_list
 
-def create_bbox(bbox_init,sh):
 
-    w=np.maximum(bbox_init[2],bbox_init[3])
+def create_bbox(bbox_init, sh):
+    w = np.maximum(bbox_init[2], bbox_init[3])
 
-    x1=bbox_init[0]-0.1*w
-    y1=bbox_init[1]-0.1*w
+    x1 = bbox_init[0] - 0.1 * w
+    y1 = bbox_init[1] - 0.1 * w
 
-    x2=bbox_init[0]+1.1*w
-    y2=bbox_init[1]+1.1*w
+    x2 = bbox_init[0] + 1.1 * w
+    y2 = bbox_init[1] + 1.1 * w
 
-    if x1<0: x1=0
-    if y1<0: y1=0
-    if x2>=sh[0]: x2=sh[0]-1
-    if y2>=sh[1]: y2=sh[1]-1
+    if x1 < 0: x1 = 0
+    if y1 < 0: y1 = 0
+    if x2 >= sh[0]: x2 = sh[0] - 1
+    if y2 >= sh[1]: y2 = sh[1] - 1
 
-    bbox=np.around([x1,y1,x2-x1,y2-y1]).astype('int')
+    bbox = np.around([x1, y1, x2 - x1, y2 - y1]).astype('int')
 
     return bbox
 
-def uncrop(alpha,bbox,R=720,C=1280):
-
 
-    alpha=cv2.resize(alpha,(bbox[3],bbox[2]))
+def uncrop(alpha, bbox, R=720, C=1280):
+    alpha = cv2.resize(alpha, (bbox[3], bbox[2]))
 
-    if alpha.ndim==2:
-        alpha_uncrop=np.zeros((R,C))
-        alpha_uncrop[bbox[0]:bbox[0]+bbox[2],bbox[1]:bbox[1]+bbox[3]]=alpha
+    if alpha.ndim == 2:
+        alpha_uncrop = np.zeros((R, C))
+        alpha_uncrop[bbox[0]:bbox[0] + bbox[2], bbox[1]:bbox[1] + bbox[3]] = alpha
     else:
-        alpha_uncrop=np.zeros((R,C,3))
-        alpha_uncrop[bbox[0]:bbox[0]+bbox[2],bbox[1]:bbox[1]+bbox[3],:]=alpha
-
+        alpha_uncrop = np.zeros((R, C, 3))
+        alpha_uncrop[bbox[0]:bbox[0] + bbox[2], bbox[1]:bbox[1] + bbox[3], :] = alpha
 
     return alpha_uncrop.astype(np.uint8)
 
 
 def to_image(rec0):
-    rec0=((rec0.data).cpu()).numpy()
-    rec0=(rec0+1)/2
-    rec0=rec0.transpose((1,2,0))
-    rec0[rec0>1]=1
-    rec0[rec0<0]=0
+    rec0 = ((rec0.data).cpu()).numpy()
+    rec0 = (rec0 + 1) / 2
+    rec0 = rec0.transpose((1, 2, 0))
+    rec0[rec0 > 1] = 1
+    rec0[rec0 < 0] = 0
     return rec0
 
-def write_tb_log(image,tag,log_writer,i):
+
+def write_tb_log(image, tag, log_writer, i):
     # image1
-    output_to_show = image.cpu().data[0:4,...]
-    output_to_show = (output_to_show + 1)/2.0
-    grid = torchvision.utils.make_grid(output_to_show,nrow=4)
+    output_to_show = image.cpu().data[0:4, ...]
+    output_to_show = (output_to_show + 1) / 2.0
+    grid = torchvision.utils.make_grid(output_to_show, nrow=4)
 
     log_writer.add_image(tag, grid, i + 1)
-

requirements.txt

@@ -1,3 +1,4 @@
+tensorflow-gpu==1.14.0
 numpy==1.17.0
 opencv-python==3.4.5.20
 pandas