File size: 3,328 Bytes
17cd746 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
#!/usr/bin/env python
# Copyright (c) Xuangeng Chu ([email protected])
# Modified based on code from Orest Kupyn (University of Oxford).
import os
import torch
import numpy as np
import torchvision
from .utils_vgghead import nms
from .utils_lmks_detector import LmksDetector
class VGGHeadDetector(torch.nn.Module):
def __init__(self, device,
vggheadmodel_path=None):
super().__init__()
self.image_size = 640
self._device = device
self.vggheadmodel_path = vggheadmodel_path
self._init_models()
def _init_models(self,):
# vgg_heads_l
self.model = torch.load(self.vggheadmodel_path, map_location='cpu')
self.model.to(self._device).eval()
@torch.no_grad()
def forward(self, image_tensor, image_key, conf_threshold=0.5):
if not hasattr(self, 'model'):
self._init_models()
image_tensor = image_tensor.to(self._device).float()
image, padding, scale = self._preprocess(image_tensor)
bbox, scores, flame_params = self.model(image)
bbox, vgg_results = self._postprocess(bbox, scores, flame_params, conf_threshold)
if bbox is None:
print('VGGHeadDetector: No face detected: {}!'.format(image_key))
return None, None, None
vgg_results['normalize'] = {'padding': padding, 'scale': scale}
# bbox
bbox = bbox.clip(0, self.image_size)
bbox[[0, 2]] -= padding[0]; bbox[[1, 3]] -= padding[1]; bbox /= scale
bbox = bbox.clip(0, self.image_size / scale)
return vgg_results, bbox, None
def _preprocess(self, image):
_, h, w = image.shape
if h > w:
new_h, new_w = self.image_size, int(w * self.image_size / h)
else:
new_h, new_w = int(h * self.image_size / w), self.image_size
scale = self.image_size / max(h, w)
image = torchvision.transforms.functional.resize(image, (new_h, new_w), antialias=True)
pad_w = self.image_size - image.shape[2]
pad_h = self.image_size - image.shape[1]
image = torchvision.transforms.functional.pad(image, (pad_w // 2, pad_h // 2, pad_w - pad_w // 2, pad_h - pad_h // 2), fill=127)
image = image.unsqueeze(0).float() / 255.0
return image, np.array([pad_w // 2, pad_h // 2]), scale
def _postprocess(self, bbox, scores, flame_params, conf_threshold):
# flame_params = {"shape": 300, "exp": 100, "rotation": 6, "jaw": 3, "translation": 3, "scale": 1}
bbox, scores, flame_params = nms(bbox, scores, flame_params, confidence_threshold=conf_threshold)
if bbox.shape[0] == 0:
return None, None
max_idx = ((bbox[:, 3] - bbox[:, 1]) * (bbox[:, 2] - bbox[:, 0])).argmax().long()
bbox, flame_params = bbox[max_idx], flame_params[max_idx]
if bbox[0] < 5 and bbox[1] < 5 and bbox[2] > 635 and bbox[3] > 635:
return None, None
# flame
posecode = torch.cat([flame_params.new_zeros(3), flame_params[400:403]])
vgg_results = {
'rotation_6d': flame_params[403:409], 'translation': flame_params[409:412], 'scale': flame_params[412:],
'shapecode': flame_params[:300], 'expcode': flame_params[300:400], 'posecode': posecode,
}
return bbox, vgg_results
|