From 2e54da45561377d0ff5b890a1ddafbcee63a2955 Mon Sep 17 00:00:00 2001 From: lucasb-eyer Date: Wed, 11 Apr 2018 15:16:29 +0200 Subject: [PATCH] v2 (PyTorch) implementation. --- README.md | 37 +++++++++- models/body_80_200.py | 56 -------------- models/full_80_240.py | 56 -------------- models/head_50_50.py | 50 ------------- models/head_80_50.py | 53 ------------- models/v1.py | 18 +++++ models/v1_body_80_200.py | 61 +++++++++++++++ models/v1_full_80_240.py | 61 +++++++++++++++ models/v1_head_50_50.py | 55 ++++++++++++++ models/v1_head_80_50.py | 56 ++++++++++++++ models/v2_head_80_50.py | 155 +++++++++++++++++++++++++++++++++++++++ scripts/common.py | 14 +++- scripts/predict.py | 67 ++++++----------- 13 files changed, 474 insertions(+), 265 deletions(-) delete mode 100644 models/body_80_200.py delete mode 100644 models/full_80_240.py delete mode 100644 models/head_50_50.py delete mode 100644 models/head_80_50.py create mode 100644 models/v1.py create mode 100644 models/v1_body_80_200.py create mode 100644 models/v1_full_80_240.py create mode 100644 models/v1_head_50_50.py create mode 100644 models/v1_head_80_50.py create mode 100644 models/v2_head_80_50.py diff --git a/README.md b/README.md index 2847095..e76b90c 100644 --- a/README.md +++ b/README.md @@ -4,13 +4,37 @@ An implementation of [BiternionNets](http://www.vision.rwth-aachen.de/publicatio Install instructions -------------------- -Most ROS installations are using Ubuntu, this is the recommended way: +Most ROS installations are using Ubuntu, this is the recommended way. + +There are two versions of the code in this repo: +- The original version (`v1`) is based on my `DeepFried2` toolbox + and `Theano`, for which most code and this README are written. + The problem is that `Theano` and hence my toolbox are discontinued. +- The new version (`v2`) uses the newer `PyTorch` library. + Unfortunately, I don't have the time to port everything, so it is + a little less documented, especially the training part. +Which version is used only depends on which model you choose as a parameter +when launching `predict.launch`. + +### Setting up a virtualenv + +This step is common to both versions: ``` $ sudo apt-get install python-virtualenv libopenblas-dev liblapack-dev gfortran $ virtualenv --system-site-packages pyenv $ . pyenv/bin/activate $ pip install --upgrade numpy +``` + +You'll always have to run `. pyenv/bin/activate` when you open a new terminal +and expect to use BiternionNets. + +Now on to the version-specific dependencies: + +### Installing the original Theano-based version + +``` $ pip install Theano==0.9.0 $ pip install git+https://github.com/lucasb-eyer/DeepFried2.git $ pip install git+https://github.com/lucasb-eyer/lbtoolbox.git@18ef22987088c9eb3153abf37dc30abe027b2708 @@ -26,6 +50,17 @@ It may take a while since it is pre-compiling quite some things: $ python -c 'import DeepFried2' ``` +### Installing the new PyTorch-based version + +Follow the [instructions on the official PyTorch website](https://pytorch.org), +but make sure to choose the `pip`-based version (and have the venv activated). + +You can test if it worked by running the following and seeing no error: + +``` +$ python -c 'import torch' +``` + Training a model ---------------- diff --git a/models/body_80_200.py b/models/body_80_200.py deleted file mode 100644 index 9cf3f10..0000000 --- a/models/body_80_200.py +++ /dev/null @@ -1,56 +0,0 @@ -import cv2 -import numpy as np -import DeepFried2 as df -from lbtoolbox.augmentation import AugmentationPipeline, Cropper -from df_extras import Flatten, Biternion - - -def mknet(): - return df.Sequential( # 184x76 - df.SpatialConvolution( 3, 24, (3, 3)), # 182x74 - df.BatchNormalization(24), - df.ReLU(), - df.SpatialConvolution(24, 24, (3, 3)), # 180x72 - df.SpatialMaxPooling((3, 3)), # 60x24 - df.BatchNormalization(24), - df.ReLU(), - df.SpatialConvolution(24, 48, (3, 3)), # 58x22 - df.BatchNormalization(48), - df.ReLU(), - df.SpatialConvolution(48, 48, (3, 3)), # 56x20 - df.SpatialMaxPooling((2, 2)), # 28x10 - df.BatchNormalization(48), - df.ReLU(), - df.SpatialConvolution(48, 64, (3, 3)), # 26x8 - df.BatchNormalization(64), - df.ReLU(), - df.SpatialConvolution(64, 64, (3, 3)), # 24x6 - df.SpatialMaxPooling((2, 2)), # 12x3 - df.BatchNormalization(64), - df.ReLU(), - df.SpatialConvolution(64, 64, (3, 2)), # 10x2 - df.BatchNormalization(64), - df.ReLU(), - df.Dropout(0.2), - Flatten(), - df.Linear(64*10*2, 512), - df.ReLU(), - df.Dropout(0.5), - df.Linear(512, 2, init=df.init.normal(0.01)), - Biternion() - ) - - -def mkaug(Xtr, ytr): - return AugmentationPipeline(Xtr, ytr, Cropper((184,76))) - -def preproc(im): - im = cv2.resize(im, (80, 200)) - im = np.rollaxis(im, 2, 0) - return im.astype(df.floatX)/255 - -def getrect(x,y,w,h): - # Here we use the full box. - # We know from the detector that full-height = 3x width. - # If that's more than is seen on camera, it will be clipped. - return x,y+int(w*0.8),w,2*w diff --git a/models/full_80_240.py b/models/full_80_240.py deleted file mode 100644 index 70470b4..0000000 --- a/models/full_80_240.py +++ /dev/null @@ -1,56 +0,0 @@ -import cv2 -import numpy as np -import DeepFried2 as df -from lbtoolbox.augmentation import AugmentationPipeline, Cropper -from df_extras import Flatten, Biternion - - -def mknet(): - return df.Sequential( # 220x76 - df.SpatialConvolution( 3, 24, (3, 3)), # 218x74 - df.BatchNormalization(24), - df.ReLU(), - df.SpatialConvolution(24, 24, (3, 3)), # 216x72 - df.SpatialMaxPooling((3, 3)), # 72x24 - df.BatchNormalization(24), - df.ReLU(), - df.SpatialConvolution(24, 48, (3, 3)), # 70x22 - df.BatchNormalization(48), - df.ReLU(), - df.SpatialConvolution(48, 48, (3, 3)), # 68x20 - df.SpatialMaxPooling((2, 2)), # 34x10 - df.BatchNormalization(48), - df.ReLU(), - df.SpatialConvolution(48, 64, (3, 3)), # 32x8 - df.BatchNormalization(64), - df.ReLU(), - df.SpatialConvolution(64, 64, (3, 3)), # 30x6 - df.SpatialMaxPooling((2, 2)), # 15x3 - df.BatchNormalization(64), - df.ReLU(), - df.SpatialConvolution(64, 64, (3, 2)), # 13x2 - df.BatchNormalization(64), - df.ReLU(), - df.Dropout(0.2), - Flatten(), - df.Linear(64*13*2, 512), - df.ReLU(), - df.Dropout(0.5), - df.Linear(512, 2, init=df.init.normal(0.01)), - Biternion() - ) - - -def mkaug(Xtr, ytr): - return AugmentationPipeline(Xtr, ytr, Cropper((220,76))) - -def preproc(im): - im = cv2.resize(im, (80, 240)) - im = np.rollaxis(im, 2, 0) - return im.astype(df.floatX)/255 - -def getrect(x,y,w,h): - # Here we use the full box. - # We know from the detector that full-height = 3x width. - # If that's more than is seen on camera, it will be clipped. - return x,y,w,3*w diff --git a/models/head_50_50.py b/models/head_50_50.py deleted file mode 100644 index 14ea76f..0000000 --- a/models/head_50_50.py +++ /dev/null @@ -1,50 +0,0 @@ -import cv2 -import numpy as np -import DeepFried2 as df -from lbtoolbox.augmentation import AugmentationPipeline, Cropper -from df_extras import Flatten, Biternion - - -def mknet(): - return df.Sequential( # 3@46 - df.SpatialConvolution( 3, 24, (3, 3)), # -> 24@44 - df.BatchNormalization(24), - df.ReLU(), - df.SpatialConvolution(24, 24, (3, 3)), # -> 24@42 - df.BatchNormalization(24), - df.SpatialMaxPooling((2, 2), ignore_border=False), # -> 24@21 - df.ReLU(), - df.SpatialConvolution(24, 48, (3, 3)), # -> 48@19 - df.BatchNormalization(48), - df.ReLU(), - df.SpatialConvolution(48, 48, (3, 3)), # -> 48@17 - df.BatchNormalization(48), - df.SpatialMaxPooling((2, 2), ignore_border=False), # -> 48@9 - df.ReLU(), - df.SpatialConvolution(48, 64, (3, 3)), # -> 64@7 - df.BatchNormalization(64), - df.ReLU(), - df.SpatialConvolution(64, 64, (3, 3)), # -> 64@5 - df.BatchNormalization(64), - df.ReLU(), - df.Dropout(0.2), - Flatten(), - df.Linear(64*5*5, 512), - df.ReLU(), - df.Dropout(0.5), - df.Linear(512, 2, init=df.init.normal(0.01)), - Biternion() - ) - - -def mkaug(Xtr, ytr): - return AugmentationPipeline(Xtr, ytr, Cropper((46,46))) - -def preproc(im): - im = cv2.resize(im, (50, 50)) - im = np.rollaxis(im, 2, 0) - return im.astype(df.floatX)/255 - -def getrect(x,y,w,h): - # Take only the square upper-body section. - return x,y,w,w diff --git a/models/head_80_50.py b/models/head_80_50.py deleted file mode 100644 index 6c27fd8..0000000 --- a/models/head_80_50.py +++ /dev/null @@ -1,53 +0,0 @@ -import cv2 -import numpy as np -import DeepFried2 as df -from lbtoolbox.augmentation import AugmentationPipeline, Cropper, Gamma -from df_extras import Flatten, Biternion - - -def mknet(): - return df.Sequential( # 48x70 (HxW) - df.SpatialConvolution( 3, 24, (3, 3)), # 46x68 - df.BatchNormalization(24), - df.ReLU(), - df.SpatialConvolution(24, 24, (3, 3)), # 44x66 - df.BatchNormalization(24), - df.SpatialMaxPooling((2, 3)), # 22x22 - df.ReLU(), - df.SpatialConvolution(24, 48, (3, 3)), # 20x20 - df.BatchNormalization(48), - df.ReLU(), - df.SpatialConvolution(48, 48, (3, 3)), # 18x18 - df.BatchNormalization(48), - df.SpatialMaxPooling((2, 2)), # 9x9 - df.ReLU(), - df.SpatialConvolution(48, 64, (3, 3)), # 7x7 - df.BatchNormalization(64), - df.ReLU(), - df.SpatialConvolution(64, 64, (3, 3)), # 5x5 - df.BatchNormalization(64), - df.ReLU(), - df.Dropout(0.2), - Flatten(), - df.Linear(64*5*5, 512), - df.ReLU(), - df.Dropout(0.5), - df.Linear(512, 2, init=df.init.normal(0.01)), - Biternion() - ) - - -def mkaug(Xtr, ytr): - return AugmentationPipeline(Xtr, ytr, - Cropper((48,70)), - # Gamma(), - ) - -def preproc(im): - im = cv2.resize(im, (80, 54)) - im = np.rollaxis(im, 2, 0) - return im.astype(df.floatX)/255 - -def getrect(x,y,w,h): - # Take only the square upper-body section. - return x,y,int(w*0.8),int(w*0.5) diff --git a/models/v1.py b/models/v1.py new file mode 100644 index 0000000..05baa27 --- /dev/null +++ b/models/v1.py @@ -0,0 +1,18 @@ +import numpy as np +import common as C + +class ModelV1: + def __call__(self, rgb, d, detrects): + if len(detrects) == 0: + return np.array([]), np.array([]) + + images = [] + for detrect in detrects: + detrect = self.getrect(*detrect) + images.append(self._preproc(C.cutout(rgb, *detrect), C.cutout(d, *detrect))) + images = np.array(images) + + bits = [self._net.forward(batch) for batch in self._aug.augbatch_pred(images, fast=True)] + preds = C.bit2deg(C.ensemble_biternions(bits)) - 90 # Subtract 90 to correct for "my weird" origin. + + return preds, np.full(len(detrects), 0.83) diff --git a/models/v1_body_80_200.py b/models/v1_body_80_200.py new file mode 100644 index 0000000..6109a18 --- /dev/null +++ b/models/v1_body_80_200.py @@ -0,0 +1,61 @@ +import cv2 +import numpy as np +import DeepFried2 as df +import common as C +from lbtoolbox.augmentation import AugmentationPipeline, Cropper +from df_extras import Flatten, Biternion + +from models.v1 import ModelV1 + +class Model(ModelV1): + def __init__(self, weightsname, *unused, **unused_kw): + self._net = df.Sequential( # 184x76 + df.SpatialConvolution( 3, 24, (3, 3)), # 182x74 + df.BatchNormalization(24), + df.ReLU(), + df.SpatialConvolution(24, 24, (3, 3)), # 180x72 + df.SpatialMaxPooling((3, 3)), # 60x24 + df.BatchNormalization(24), + df.ReLU(), + df.SpatialConvolution(24, 48, (3, 3)), # 58x22 + df.BatchNormalization(48), + df.ReLU(), + df.SpatialConvolution(48, 48, (3, 3)), # 56x20 + df.SpatialMaxPooling((2, 2)), # 28x10 + df.BatchNormalization(48), + df.ReLU(), + df.SpatialConvolution(48, 64, (3, 3)), # 26x8 + df.BatchNormalization(64), + df.ReLU(), + df.SpatialConvolution(64, 64, (3, 3)), # 24x6 + df.SpatialMaxPooling((2, 2)), # 12x3 + df.BatchNormalization(64), + df.ReLU(), + df.SpatialConvolution(64, 64, (3, 2)), # 10x2 + df.BatchNormalization(64), + df.ReLU(), + df.Dropout(0.2), + Flatten(), + df.Linear(64*10*2, 512), + df.ReLU(), + df.Dropout(0.5), + df.Linear(512, 2, init=df.init.normal(0.01)), + Biternion() + ) + + self._net.__setstate__(np.load(weightsname)) + self._net.evaluate() + + self._aug = AugmentationPipeline(None, None, Cropper((184,76))) + + def getrect(self, x, y, w, h): + # Here we use the full box. + # We know from the detector that full-height = 3x width. + # If that's more than is seen on camera, it will be clipped. + return x, y+int(w*0.8), w, 2*w + + def _preproc(self, det_rgb, det_d): + im = C.subtractbg(det_rgb, det_d, 1.0, 0.5) + im = cv2.resize(im, (80, 200)) + im = np.rollaxis(im, 2, 0) + return im.astype(df.floatX)/255 diff --git a/models/v1_full_80_240.py b/models/v1_full_80_240.py new file mode 100644 index 0000000..a3baef2 --- /dev/null +++ b/models/v1_full_80_240.py @@ -0,0 +1,61 @@ +import cv2 +import numpy as np +import DeepFried2 as df +import common as C +from lbtoolbox.augmentation import AugmentationPipeline, Cropper +from df_extras import Flatten, Biternion + +from models.v1 import ModelV1 + +class Model(ModelV1): + def __init__(self, weightsname, *unused, **unused_kw): + self._net = df.Sequential( # 220x76 + df.SpatialConvolution( 3, 24, (3, 3)), # 218x74 + df.BatchNormalization(24), + df.ReLU(), + df.SpatialConvolution(24, 24, (3, 3)), # 216x72 + df.SpatialMaxPooling((3, 3)), # 72x24 + df.BatchNormalization(24), + df.ReLU(), + df.SpatialConvolution(24, 48, (3, 3)), # 70x22 + df.BatchNormalization(48), + df.ReLU(), + df.SpatialConvolution(48, 48, (3, 3)), # 68x20 + df.SpatialMaxPooling((2, 2)), # 34x10 + df.BatchNormalization(48), + df.ReLU(), + df.SpatialConvolution(48, 64, (3, 3)), # 32x8 + df.BatchNormalization(64), + df.ReLU(), + df.SpatialConvolution(64, 64, (3, 3)), # 30x6 + df.SpatialMaxPooling((2, 2)), # 15x3 + df.BatchNormalization(64), + df.ReLU(), + df.SpatialConvolution(64, 64, (3, 2)), # 13x2 + df.BatchNormalization(64), + df.ReLU(), + df.Dropout(0.2), + Flatten(), + df.Linear(64*13*2, 512), + df.ReLU(), + df.Dropout(0.5), + df.Linear(512, 2, init=df.init.normal(0.01)), + Biternion() + ) + + self._net.__setstate__(np.load(weightsname)) + self._net.evaluate() + + self._aug = AugmentationPipeline(None, None, Cropper((220,76))) + + def getrect(self, x, y, w, h): + # Here we use the full box. + # We know from the detector that full-height = 3x width. + # If that's more than is seen on camera, it will be clipped. + return x, y, w, 3*w + + def _preproc(self, det_rgb, det_d): + im = C.subtractbg(det_rgb, det_d, 1.0, 0.5) + im = cv2.resize(im, (80, 240)) + im = np.rollaxis(im, 2, 0) + return im.astype(df.floatX)/255 diff --git a/models/v1_head_50_50.py b/models/v1_head_50_50.py new file mode 100644 index 0000000..22a7790 --- /dev/null +++ b/models/v1_head_50_50.py @@ -0,0 +1,55 @@ +import cv2 +import numpy as np +import DeepFried2 as df +import common as C +from lbtoolbox.augmentation import AugmentationPipeline, Cropper +from df_extras import Flatten, Biternion + +from models.v1 import ModelV1 + +class Model(ModelV1): + def __init__(self, weightsname, *unused, **unused_kw): + self._net = df.Sequential( # 3@46 + df.SpatialConvolution( 3, 24, (3, 3)), # -> 24@44 + df.BatchNormalization(24), + df.ReLU(), + df.SpatialConvolution(24, 24, (3, 3)), # -> 24@42 + df.BatchNormalization(24), + df.SpatialMaxPooling((2, 2), ignore_border=False), # -> 24@21 + df.ReLU(), + df.SpatialConvolution(24, 48, (3, 3)), # -> 48@19 + df.BatchNormalization(48), + df.ReLU(), + df.SpatialConvolution(48, 48, (3, 3)), # -> 48@17 + df.BatchNormalization(48), + df.SpatialMaxPooling((2, 2), ignore_border=False), # -> 48@9 + df.ReLU(), + df.SpatialConvolution(48, 64, (3, 3)), # -> 64@7 + df.BatchNormalization(64), + df.ReLU(), + df.SpatialConvolution(64, 64, (3, 3)), # -> 64@5 + df.BatchNormalization(64), + df.ReLU(), + df.Dropout(0.2), + Flatten(), + df.Linear(64*5*5, 512), + df.ReLU(), + df.Dropout(0.5), + df.Linear(512, 2, init=df.init.normal(0.01)), + Biternion() + ) + + self._net.__setstate__(np.load(weightsname)) + self._net.evaluate() + + self._aug = AugmentationPipeline(None, None, Cropper((46, 46))) + + def getrect(self, x, y, w, h): + # Take only the square upper-body section. + return x, y, w, w + + def _preproc(self, det_rgb, det_d): + im = C.subtractbg(det_rgb, det_d, 1.0, 0.5) + im = cv2.resize(im, (50, 50)) + im = np.rollaxis(im, 2, 0) + return im.astype(df.floatX)/255 diff --git a/models/v1_head_80_50.py b/models/v1_head_80_50.py new file mode 100644 index 0000000..7e93c03 --- /dev/null +++ b/models/v1_head_80_50.py @@ -0,0 +1,56 @@ +import cv2 +import numpy as np +import DeepFried2 as df +import common as C +from lbtoolbox.augmentation import AugmentationPipeline, Cropper +from df_extras import Flatten, Biternion + +from models.v1 import ModelV1 + +class Model(ModelV1): + def __init__(self, weightsname, *unused, **unused_kw): + self._net = df.Sequential( # 48x70 (HxW) + df.SpatialConvolution( 3, 24, (3, 3)), # 46x68 + df.BatchNormalization(24), + df.ReLU(), + df.SpatialConvolution(24, 24, (3, 3)), # 44x66 + df.BatchNormalization(24), + df.SpatialMaxPooling((2, 3)), # 22x22 + df.ReLU(), + df.SpatialConvolution(24, 48, (3, 3)), # 20x20 + df.BatchNormalization(48), + df.ReLU(), + df.SpatialConvolution(48, 48, (3, 3)), # 18x18 + df.BatchNormalization(48), + df.SpatialMaxPooling((2, 2)), # 9x9 + df.ReLU(), + df.SpatialConvolution(48, 64, (3, 3)), # 7x7 + df.BatchNormalization(64), + df.ReLU(), + df.SpatialConvolution(64, 64, (3, 3)), # 5x5 + df.BatchNormalization(64), + df.ReLU(), + df.Dropout(0.2), + Flatten(), + df.Linear(64*5*5, 512), + df.ReLU(), + df.Dropout(0.5), + df.Linear(512, 2, init=df.init.normal(0.01)), + Biternion() + ) + + self._net.__setstate__(np.load(weightsname)) + self._net.evaluate() + + self._aug = AugmentationPipeline(None, None, Cropper((48,70))) + + def getrect(self, x, y, w, h): + # Take only the square upper-body section. + # return x, y, int(w*0.8), int(w*0.5) + return int(x+.1*w), y, int(w*0.8), int(w*0.5) + + def _preproc(self, det_rgb, det_d): + im = C.subtractbg(det_rgb, det_d, 1.0, 0.5) + im = cv2.resize(im, (80, 54)) + im = np.rollaxis(im, 2, 0) + return im.astype(df.floatX)/255 diff --git a/models/v2_head_80_50.py b/models/v2_head_80_50.py new file mode 100644 index 0000000..d97db43 --- /dev/null +++ b/models/v2_head_80_50.py @@ -0,0 +1,155 @@ +import cv2 +import numpy as np + +import torch +import torch.nn as nn +from torch.autograd import Variable + +import common as C + + +# Monkey-patch because I trained with a newer version. +# This can be removed once PyTorch 0.4.x is out. +import torch._utils +try: + torch._utils._rebuild_tensor_v2 + OLD_TORCH = False +except AttributeError: + def _rebuild_tensor_v2(storage, storage_offset, size, stride, requires_grad, backward_hooks): + tensor = torch._utils._rebuild_tensor(storage, storage_offset, size, stride) + tensor.requires_grad = requires_grad + tensor._backward_hooks = backward_hooks + return tensor + torch._utils._rebuild_tensor_v2 = _rebuild_tensor_v2 + OLD_TORCH = True + + +def conv3(cin, cout, **kw): + kw.setdefault('bias', False) + kw.setdefault('padding', 1) + return nn.Conv2d(cin, cout, 3, **kw) + + +class Model: + def __init__(self, weightsname, GPU=False, hfactor=0.6, wfactor=1.0): + self._GPU = GPU + self._hfactor = hfactor + self._wfactor = wfactor + + torch.backends.cudnn.benchmark = True # Run benchmark to select fastest implementation of ops. + + # TODO: Load this from saved model settings. + pool = 'max' + dropout = (0.2, 0.2) + + if pool == 'avg': + head = nn.AvgPool2d((1,5)) + elif pool == 'max': + head = nn.MaxPool2d((1,5)) + elif pool == 'fc': + head = nn.Sequential(nn.Conv2d(128, 128, (1,5), bias=False), nn.BatchNorm2d(128), nn.ReLU()) + + self._net = nn.Sequential( # 50x80 -> 40x72 + conv3( 3, 16), nn.BatchNorm2d(16), nn.ReLU(), + conv3(16, 32), nn.BatchNorm2d(32), nn.ReLU(), + nn.MaxPool2d(2), # 20x36 + (nn.Dropout2d(dropout[0]) if dropout is not None else nn.Sequential()), + conv3(32, 32), nn.BatchNorm2d(32), nn.ReLU(), + conv3(32, 64), nn.BatchNorm2d(64), nn.ReLU(), + nn.MaxPool2d(2), # 10x18 + (nn.Dropout2d(dropout[0]) if dropout is not None else nn.Sequential()), + conv3(64, 64), nn.BatchNorm2d(64), nn.ReLU(), + conv3(64, 96), nn.BatchNorm2d(96), nn.ReLU(), + nn.MaxPool2d(2), # 5x9 + (nn.Dropout2d(dropout[0]) if dropout is not None else nn.Sequential()), + nn.Conv2d(96, 96, 3), nn.BatchNorm2d(96), nn.ReLU(), + nn.Conv2d(96,128, 3), nn.BatchNorm2d(128), nn.ReLU(), + head, + (nn.Dropout(dropout[1]) if dropout is not None else nn.Sequential()), + nn.Conv2d(128, 3, 1) + ) + + # Load the network weights onto the CPU first, no matter where they trained. + checkpoint = torch.load(weightsname, map_location='cpu') + self._net.load_state_dict(checkpoint) + + # And finally move the thing to the GPU, if needed. + self._net = self._maybe_gpu(self._net) + + self._net.eval() + + def __call__(self, rgb, d, detrects): + if len(detrects) == 0: + return np.array([]), np.array([]) + + images = np.array([self._preproc(*self._cutout(rgb, d, *rect)) for rect in detrects]) + + preds = self._forward(images) + + # Need to split out the raw predictions into what they really are, + # and apply their respecitve non-linearities. + biternions, confidences = preds[:,:2], preds[:,2] + biternions = C.normalized(biternions, axis=1) + print(biternions[0,:,0,0]) + confidences = self._conf_nonlin(confidences) + + # The predictions still contain the spatial dimension: BCHW. + # We just average these out, though more advanced things could be done, + # such as weighted average weighting by certainty, etc. + biternions = np.mean(biternions, axis=(-2,-1)) + confidences = np.mean(confidences, axis=(-2,-1)) + + # Need to normalize again in order to have real biternions + biternions = C.normalized(biternions) + + # Convert the biternions into an angle as a common API. + angles = C.bit2deg(biternions) + + return angles, confidences + + def getrect(self, x, y, w, h): + """ + Transform an original detection/tracking box into what we need for + making predictions, i.e. should be the same as during dumping. + """ + x, y, w, h = C.cutout_hwfact(x, y, w, h, self._hfactor, self._wfactor) + return x, y, w, h + + def _cutout(self, rgb, d, x, y, w, h): + rect = self.getrect(x, y, w, h) + return C.cutout(rgb, *rect), C.cutout(d, *rect) + + def _preproc(self, rgb, d): + """ Transforms an rgb+d cut-out into a network-input. + + We could also do background-subtraction here, but I trained the model without. + """ + # First, resize the image to a fixed size. + # Note the argument is (w,h) but the shape is HWC + rgb = cv2.resize(rgb, (80, 50)) + # From HWC to CHW, and then [0,255] to [0,1] + return rgb.transpose(2,0,1).astype(np.float32)/255 + + def _conf_nonlin(self, k): + MAXK = 8 + return MAXK*C.sigmoid(float(np.log(MAXK-1))/(MAXK/2)*(k-(MAXK/2))) + + def _forward(self, x): + """ Send numpy-array x through the network and return numpy-array of results. """ + if OLD_TORCH: + x_var = Variable(self._maybe_gpu(torch.from_numpy(x)), volatile=True) + return self._net(x_var).data.cpu().numpy() + else: + with torch.no_grad(): + x_var = Variable(self._maybe_gpu(torch.from_numpy(x))) + return self._net(x_var).data.cpu().numpy() + + def _maybe_gpu(self, whatever): + """ Moves `what` to CUDA and returns it, if `use_cuda` and it's available. + + Actually, `use_cuda` is the GPU-index to be used, which means `0` uses the + first GPU. To not use GPUs, set `use_cuda` to `False` instead. + """ + if self._GPU is not False and torch.cuda.is_available(): + whatever = whatever.cuda(device=self._GPU, **kw) + return whatever diff --git a/scripts/common.py b/scripts/common.py index 56f4a13..c575d06 100644 --- a/scripts/common.py +++ b/scripts/common.py @@ -1,18 +1,26 @@ import numpy as np +def normalized(x, axis=-1): + return x / (np.linalg.norm(x, axis=axis, keepdims=True) + 1e-12) + + +def sigmoid(x): + return np.exp(-np.logaddexp(0, -x)) + + def deg2bit(deg): rad = np.deg2rad(deg) return np.array([np.cos(rad), np.sin(rad)]).T -def bit2deg(angles_bit): - return (np.rad2deg(np.arctan2(angles_bit[:,1], angles_bit[:,0])) + 360) % 360 +def bit2deg(bit): + return (np.rad2deg(np.arctan2(bit[...,1], bit[...,0])) + 360) % 360 def flipbiternions(bits): bits = bits.copy() - bits[:,1] *= -1 + bits[...,1] *= -1 return bits diff --git a/scripts/predict.py b/scripts/predict.py index 6392392..85b1856 100755 --- a/scripts/predict.py +++ b/scripts/predict.py @@ -19,10 +19,6 @@ from biternion.msg import HeadOrientations from visualization_msgs.msg import Marker -import DeepFried2 as df - -from common import bit2deg, ensemble_biternions, subtractbg, cutout - # Distinguish between STRANDS and SPENCER. try: from rwth_perception_people_msgs.msg import UpperBodyDetector @@ -62,18 +58,11 @@ def __init__(self): # Create and load the network. netlib = import_module(modelname) - self.net = netlib.mknet() - self.net.__setstate__(np.load(weightsname)) - self.net.evaluate() - - self.aug = netlib.mkaug(None, None) - self.preproc = netlib.preproc - self.getrect = netlib.getrect - - # Do a fake forward-pass for precompilation. - im = cutout(np.zeros((480,640,3), np.uint8), 0, 0, 150, 450) - im = next(self.aug.augimg_pred(self.preproc(im), fast=True)) - self.net.forward(np.array([im])) + self.model = netlib.Model(weightsname, GPU=False) + + # Do a fake forward-pass for precompilation/GPU init/... + self.model(np.zeros((480,640,3), np.uint8), + np.zeros((480,640), np.float32), [(0,0,150,450)]) rospy.loginfo("BiternionNet initialized") src = rospy.get_param("~src", "tra") @@ -103,9 +92,9 @@ def __init__(self): def cb(self, src, rgb, d, caminfo, *more): # Ugly workaround because approximate sync sometimes jumps back in time. - if rgb.header.stamp <= self.last_stamp: - rospy.logwarn("Jump back in time detected and dropped like it's hot") - return + #if rgb.header.stamp <= self.last_stamp: + # rospy.logwarn("Jump back in time detected and dropped like it's hot") + # return self.last_stamp = rgb.header.stamp @@ -124,45 +113,31 @@ def cb(self, src, rgb, d, caminfo, *more): header = rgb.header bridge = CvBridge() - rgb = bridge.imgmsg_to_cv2(rgb)[:,:,::-1] # Need to do BGR-RGB conversion manually. + rgb = bridge.imgmsg_to_cv2(rgb, desired_encoding='rgb8') d = bridge.imgmsg_to_cv2(d) - imgs = [] - for detrect in detrects: - detrect = self.getrect(*detrect) - det_rgb = cutout(rgb, *detrect) - det_d = cutout(d, *detrect) - - # Preprocess and stick into the minibatch. - im = subtractbg(det_rgb, det_d, 1.0, 0.5) - im = self.preproc(im) - imgs.append(im) - sys.stderr.write("\r{}".format(self.counter)) ; sys.stderr.flush() - self.counter += 1 - - # TODO: We could further optimize by putting all augmentations in a - # single batch and doing only one forward pass. Should be easy. - if len(detrects): - bits = [self.net.forward(batch) for batch in self.aug.augbatch_pred(np.array(imgs), fast=True)] - preds = bit2deg(ensemble_biternions(bits)) - 90 # Subtract 90 to correct for "my weird" origin. - # print(preds) - else: - preds = [] + # Do the extraction and prediction + preds, confs = self.model(rgb, d, detrects) + self.counter += len(preds) + sys.stderr.write("\r{}".format(self.counter)) ; sys.stderr.flush() + + # Publish angle predictions if 0 < self.pub.get_num_connections(): self.pub.publish(HeadOrientations( header=header, angles=list(preds), - confidences=[0.83] * len(imgs) + confidences=list(confs), )) # Visualization + # TODO: Visualize confidence, too. if 0 < self.pub_vis.get_num_connections(): - rgb_vis = rgb[:,:,::-1].copy() + rgb_vis = rgb.copy() for detrect, alpha in zip(detrects, preds): - l, t, w, h = self.getrect(*detrect) + l, t, w, h = self.model.getrect(*detrect) px = int(round(np.cos(np.deg2rad(alpha))*w/2)) py = -int(round(np.sin(np.deg2rad(alpha))*h/2)) - cv2.rectangle(rgb_vis, (detrect[0], detrect[1]), (detrect[0]+detrect[2],detrect[1]+detrect[3]), (0,255,255), 1) + cv2.rectangle(rgb_vis, (detrect[0], detrect[1]), (detrect[0]+detrect[2],detrect[1]+detrect[3]), (0,0,255), 1) cv2.rectangle(rgb_vis, (l,t), (l+w,t+h), (0,255,0), 2) cv2.line(rgb_vis, (l+w//2, t+h//2), (l+w//2+px,t+h//2+py), (0,255,0), 2) # cv2.putText(rgb_vis, "{:.1f}".format(alpha), (l, t+25), cv2.FONT_HERSHEY_SIMPLEX, 1, (255,0,255), 2) @@ -177,7 +152,7 @@ def cb(self, src, rgb, d, caminfo, *more): poseArray = PoseArray(header=header) for (dx, dy, dw, dh, dd), alpha in zip(get_rects(src, with_depth=True), preds): - dx, dy, dw, dh = self.getrect(dx, dy, dw, dh) + dx, dy, dw, dh = self.model.getrect(dx, dy, dw, dh) # PoseArray message for boundingbox centres poseArray.poses.append(Pose(