|
import caffe |
|
import numpy as np |
|
import math |
|
|
|
#Layer config is json |
|
try: |
|
import simplejson as json |
|
except ImportError: |
|
import json |
|
|
|
face_scale_thresh = np.array([0.1, 0.2,]) |
|
person_scale_thresh = np.array([0.2, 0.4,]) |
|
offset = 0.01 |
|
|
|
class MultiScaleSpatialMapLayer(caffe.Layer): |
|
|
|
def setup(self, bottom, top): |
|
if len(bottom) >= 1: |
|
raise Exception("This is a data layer, it takes no bottom blob") |
|
|
|
self.json_param = json.loads(self.param_str) |
|
|
|
self.data_file = self.json_param['source'] |
|
self.map_size = self.json_param['map_size'] |
|
self.batch_size = self.json_param['batch_size'] |
|
|
|
self.map_ch = 6 |
|
|
|
if len(top) == 1: |
|
self.give_gt = False |
|
elif len(top) == 2: |
|
self.give_gt = True |
|
else: |
|
raise Exception("Too many top blobs, expect 1 or 2") |
|
|
|
tmp = json.load(open(self.data_file)) |
|
self.det_data = tmp['data'] |
|
self.data_cursor = iter(self.det_data) |
|
|
|
def get_next_img(self): |
|
try: |
|
next_img = self.data_cursor.next() |
|
except StopIteration, e: |
|
#rewind if reached the end |
|
self.data_cursor = iter(self.det_data) |
|
next_img = self.data_cursor.next() |
|
return next_img |
|
|
|
def reshape(self, bottom, top): |
|
|
|
top[0].reshape(self.batch_size, self.map_ch, self.map_size, self.map_size) |
|
|
|
if self.give_gt: |
|
top[1].reshape(self.batch_size, 1) |
|
|
|
self.data_buffer = np.empty((self.batch_size, self.map_ch, self.map_size, self.map_size), |
|
dtype=np.float32) |
|
|
|
def build_det_map(self, idx, img_info): |
|
|
|
img_label = img_info['gt'] |
|
|
|
def norm_coord(c, ratio, map_size, up=False): |
|
return min(math.ceil(c / ratio), map_size) if up else min(math.floor(c / ratio), map_size) |
|
|
|
def norm_bbox(bbox, w_ratio, h_ratio, map_size): |
|
return [ norm_coord(bbox[0], w_ratio, map_size), norm_coord(bbox[1], h_ratio, map_size), |
|
norm_coord(bbox[2], w_ratio, map_size, True), norm_coord(bbox[3], h_ratio, map_size, True) ] |
|
|
|
def get_scale(bbox, img_size, scale_ticks): |
|
area = float(bbox[2] - bbox[0]) * float(bbox[3] - bbox[1]) / img_size[0] / img_size[1] |
|
scale = (scale_ticks <= area).sum() |
|
return math.sqrt(scale) |
|
|
|
|
|
#face |
|
if img_info['face'] is not None: |
|
face_img_size = img_info['face']['size'] |
|
w_ratio = face_img_size[0] / float(self.map_size) |
|
h_ratio = face_img_size[1] / float(self.map_size) |
|
for fbox in img_info['face']['bbox']: |
|
nfbox = norm_bbox(fbox, w_ratio, h_ratio, self.map_size) |
|
b_scale = get_scale(fbox, face_img_size, face_scale_thresh) |
|
self.data_buffer[idx, b_scale, nfbox[1]:nfbox[3], nfbox[0]:nfbox[2]] += 1 |
|
|
|
#person |
|
if img_info['person'] is not None: |
|
person_img_size = img_info['person']['size'] |
|
w_ratio = person_img_size[0] / float(self.map_size) |
|
h_ratio = person_img_size[1] / float(self.map_size) |
|
for pbox in img_info['person']['bbox']: |
|
npbox = norm_bbox(pbox, w_ratio, h_ratio, self.map_size) |
|
b_scale = get_scale(pbox, person_img_size, person_scale_thresh) |
|
self.data_buffer[idx, b_scale + 3, npbox[1]:npbox[3], npbox[0]:npbox[2]] += 1 |
|
return img_label |
|
|
|
def forward(self, bottom, top): |
|
self.data_buffer[:] = offset |
|
gt_label = np.empty((self.batch_size, 1), dtype=np.float32) |
|
|
|
for n in xrange(self.batch_size): |
|
img_info = self.get_next_img() |
|
gt_label[n] = self.build_det_map(n, img_info) |
|
|
|
# copy buffer to output data |
|
top[0].data[:] = self.data_buffer[:] |
|
if self.give_gt: |
|
top[1].data[:] = gt_label[:] |
