def get_imagenet_tfdata(datadir, name, batch_size, mapper=None, parallel=None):
"""
Args:
mapper: a symbolic function that takes a tf.string (the raw bytes read from file) and produces a BGR image.
Defaults to `fbresnet_mapper(isTrain)`.
Returns:
A `tf.data.Dataset`. If training, the dataset is infinite.
The dataset contains BGR images and labels.
"""
def get_imglist(dir, name):
"""
Returns:
[(full filename, label)]
"""
dir = os.path.join(dir, name)
meta = dataset.ILSVRCMeta()
imglist = meta.get_image_list(
name,
dataset.ILSVRCMeta.guess_dir_structure(dir))
def _filter(fname):
# png
return 'n02105855_2933.JPEG' in fname
ret = []
for fname, label in imglist:
if _filter(fname):
logger.info("Image {} was filtered out.".format(fname))
continue
fname = os.path.join(dir, fname)
ret.append((fname, label))
return ret
assert name in ['train', 'val', 'test']
assert datadir is not None
isTrain = name == 'train'
if mapper is None:
mapper = fbresnet_mapper(isTrain)
if parallel is None:
parallel = min(40, multiprocessing.cpu_count() // 2) # assuming hyperthreading
imglist = get_imglist(datadir, name)
N = len(imglist)
filenames = tf.constant([k[0] for k in imglist], name='filenames')
labels = tf.constant([k[1] for k in imglist], dtype=tf.int32, name='labels')
ds = tf.data.Dataset.from_tensor_slices((filenames, labels))
if isTrain:
ds = ds.shuffle(N, reshuffle_each_iteration=True).repeat()
ds = ds.apply(
tf.data.experimental.map_and_batch(
lambda fname, label: (mapper(tf.read_file(fname)), label),
batch_size=batch_size,
num_parallel_batches=parallel))
ds = ds.prefetch(100)
return ds
def image_preprocess(self, image):
with tf.name_scope('image_preprocess'):
if image.dtype.base_dtype != tf.float32:
image = tf.cast(image, tf.float32)
mean = [0.485, 0.456, 0.406] # rgb
std = [0.229, 0.224, 0.225]
if self.image_bgr:
mean = mean[::-1]
std = std[::-1]
image_mean = tf.constant(mean, dtype=tf.float32) * 255.
image_std = tf.constant(std, dtype=tf.float32) * 255.
image = (image - image_mean) / image_std
return image
def build_graph(self, image, label):
xys = np.array([(y, x, 1) for y in range(WARP_TARGET_SIZE)
for x in range(WARP_TARGET_SIZE)], dtype='float32')
xys = tf.constant(xys, dtype=tf.float32, name='xys') # p x 3
image = image / 255.0 - 0.5 # bhw2
def get_stn(image):
stn = (LinearWrap(image)
.AvgPooling('downsample', 2)
.Conv2D('conv0', 20, 5, padding='VALID')
.MaxPooling('pool0', 2)
.Conv2D('conv1', 20, 5, padding='VALID')
.FullyConnected('fc1', 32)
.FullyConnected('fct', 6, activation=tf.identity,
kernel_initializer=tf.constant_initializer(),
bias_initializer=tf.constant_initializer([1, 0, HALF_DIFF, 0, 1, HALF_DIFF]))())
# output 6 parameters for affine transformation
stn = tf.reshape(stn, [-1, 2, 3], name='affine') # bx2x3
stn = tf.reshape(tf.transpose(stn, [2, 0, 1]), [3, -1]) # 3 x (bx2)
coor = tf.reshape(tf.matmul(xys, stn),
[WARP_TARGET_SIZE, WARP_TARGET_SIZE, -1, 2])
coor = tf.transpose(coor, [2, 0, 1, 3], 'sampled_coords') # b h w 2
sampled = GridSample('warp', [image, coor], borderMode='constant')
return sampled
with argscope([Conv2D, FullyConnected], activation=tf.nn.relu):
with tf.variable_scope('STN1'):
sampled1 = get_stn(image)
with tf.variable_scope('STN2'):
sampled2 = get_stn(image)
# For visualization in tensorboard
with tf.name_scope('visualization'):
padded1 = tf.pad(sampled1, [[0, 0], [HALF_DIFF, HALF_DIFF], [HALF_DIFF, HALF_DIFF], [0, 0]])
padded2 = tf.pad(sampled2, [[0, 0], [HALF_DIFF, HALF_DIFF], [HALF_DIFF, HALF_DIFF], [0, 0]])
img_orig = tf.concat([image[:, :, :, 0], image[:, :, :, 1]], 1) # b x 2h x w
transform1 = tf.concat([padded1[:, :, :, 0], padded1[:, :, :, 1]], 1)
transform2 = tf.concat([padded2[:, :, :, 0], padded2[:, :, :, 1]], 1)
stacked = tf.concat([img_orig, transform1, transform2], 2, 'viz')
tf.summary.image('visualize',
tf.expand_dims(stacked, -1), max_outputs=30)
sampled = tf.concat([sampled1, sampled2], 3, 'sampled_concat')
logits = (LinearWrap(sampled)
.FullyConnected('fc1', 256, activation=tf.nn.relu)
.FullyConnected('fc2', 128, activation=tf.nn.relu)
.FullyConnected('fct', 19, activation=tf.identity)())
tf.nn.softmax(logits, name='prob')
cost = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=label)
cost = tf.reduce_mean(cost, name='cross_entropy_loss')
wrong = tf.cast(tf.logical_not(tf.nn.in_top_k(logits, label, 1)), tf.float32, name='incorrect_vector')
summary.add_moving_summary(tf.reduce_mean(wrong, name='train_error'))
wd_cost = tf.multiply(1e-5, regularize_cost('fc.*/W', tf.nn.l2_loss),
name='regularize_loss')
summary.add_moving_summary(cost, wd_cost)
return tf.add_n([wd_cost, cost], name='cost')
def GridSample(inputs, borderMode='repeat'):
"""
Sample the images using the given coordinates, by bilinear interpolation.
This was described in the paper:
`Spatial Transformer Networks <http://arxiv.org/abs/1506.02025>`_.
This is equivalent to `torch.nn.functional.grid_sample`,
up to some non-trivial coordinate transformation.
This implementation returns pixel value at pixel (1, 1) for a floating point coordinate (1.0, 1.0).
Note that this may not be what you need.
Args:
inputs (list): [images, coords]. images has shape NHWC.
coords has shape (N, H', W', 2), where each pair of the last dimension is a (y, x) real-value
coordinate.
borderMode: either "repeat" or "constant" (zero-filled)
Returns:
tf.Tensor: a tensor named ``output`` of shape (N, H', W', C).
"""
image, mapping = inputs
assert image.get_shape().ndims == 4 and mapping.get_shape().ndims == 4
input_shape = image.get_shape().as_list()[1:]
assert None not in input_shape, \
"Images in GridSample layer must have fully-defined shape"
assert borderMode in ['repeat', 'constant']
orig_mapping = mapping
mapping = tf.maximum(mapping, 0.0)
lcoor = tf.floor(mapping)
ucoor = lcoor + 1
diff = mapping - lcoor
neg_diff = 1.0 - diff # bxh2xw2x2
lcoory, lcoorx = tf.split(lcoor, 2, 3)
ucoory, ucoorx = tf.split(ucoor, 2, 3)
lyux = tf.concat([lcoory, ucoorx], 3)
uylx = tf.concat([ucoory, lcoorx], 3)
diffy, diffx = tf.split(diff, 2, 3)
neg_diffy, neg_diffx = tf.split(neg_diff, 2, 3)
ret = tf.add_n([sample(image, lcoor) * neg_diffx * neg_diffy,
sample(image, ucoor) * diffx * diffy,
sample(image, lyux) * neg_diffy * diffx,
sample(image, uylx) * diffy * neg_diffx], name='sampled')
if borderMode == 'constant':
max_coor = tf.constant([input_shape[0] - 1, input_shape[1] - 1], dtype=tf.float32)
mask = tf.greater_equal(orig_mapping, 0.0)
mask2 = tf.less_equal(orig_mapping, max_coor)
mask = tf.logical_and(mask, mask2) # bxh2xw2x2
mask = tf.reduce_all(mask, [3]) # bxh2xw2 boolean
mask = tf.expand_dims(mask, 3)
ret = ret * tf.cast(mask, tf.float32)
return tf.identity(ret, name='output')
def build_graph(self, image, label):
drop_rate = tf.constant(0.5 if self.training else 0.0)
if self.training:
tf.summary.image("train_image", image, 10)
if tf.test.is_gpu_available():
image = tf.transpose(image, [0, 3, 1, 2])
data_format = 'channels_first'
else:
data_format = 'channels_last'
image = image / 4.0 # just to make range smaller
with argscope(Conv2D, activation=BNReLU, use_bias=False, kernel_size=3), \
argscope([Conv2D, MaxPooling, BatchNorm], data_format=data_format):
logits = LinearWrap(image) \
.Conv2D('conv1.1', filters=64) \
.Conv2D('conv1.2', filters=64) \
.MaxPooling('pool1', 3, stride=2, padding='SAME') \
.Conv2D('conv2.1', filters=128) \
.Conv2D('conv2.2', filters=128) \
.MaxPooling('pool2', 3, stride=2, padding='SAME') \
.Conv2D('conv3.1', filters=128, padding='VALID') \
.Conv2D('conv3.2', filters=128, padding='VALID') \
.FullyConnected('fc0', 1024 + 512, activation=tf.nn.relu) \
.Dropout(rate=drop_rate) \
.FullyConnected('fc1', 512, activation=tf.nn.relu) \
.FullyConnected('linear', out_dim=self.cifar_classnum)()
cost = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits,
labels=label)
cost = tf.reduce_mean(cost, name='cross_entropy_loss')
correct = tf.cast(tf.nn.in_top_k(predictions=logits,
targets=label,
k=1),
tf.float32,
name='correct')
# monitor training error
add_moving_summary(tf.reduce_mean(correct, name='accuracy'))
# weight decay on all W of fc layers
wd_cost = regularize_cost('fc.*/W',
l2_regularizer(4e-4),
name='regularize_loss')
add_moving_summary(cost, wd_cost)
add_param_summary(('.*/W', ['histogram'])) # monitor W
return tf.add_n([cost, wd_cost], name='cost')
def sample(img, coords):
"""
Args:
img: bxhxwxc
coords: bxh2xw2x2. each coordinate is (y, x) integer.
Out of boundary coordinates will be clipped.
Return:
bxh2xw2xc image
"""
shape = img.get_shape().as_list()[1:] # h, w, c
batch = tf.shape(img)[0]
shape2 = coords.get_shape().as_list()[1:3] # h2, w2
assert None not in shape2, coords.get_shape()
max_coor = tf.constant([shape[0] - 1, shape[1] - 1], dtype=tf.float32)
coords = tf.clip_by_value(coords, 0., max_coor) # borderMode==repeat
coords = tf.cast(coords, tf.int32)
batch_index = tf.range(batch, dtype=tf.int32)
batch_index = tf.reshape(batch_index, [-1, 1, 1, 1])
batch_index = tf.tile(batch_index, [1, shape2[0], shape2[1], 1]) # bxh2xw2x1
indices = tf.concat([batch_index, coords], axis=3) # bxh2xw2x3
sampled = tf.gather_nd(img, indices)
return sampled
def roi_heads(self, image, features, proposals, targets):
image_shape2d = tf.shape(image)[2:] # h,w
assert len(features) == 5, "Features have to be P23456!"
gt_boxes, gt_labels, *_ = targets
if self.training:
proposals = sample_fast_rcnn_targets(proposals.boxes, gt_boxes,
gt_labels)
fastrcnn_head_func = getattr(model_frcnn, cfg.FPN.FRCNN_HEAD_FUNC)
if not cfg.FPN.CASCADE:
roi_feature_fastrcnn = multilevel_roi_align(
features[:4], proposals.boxes, 7)
head_feature = fastrcnn_head_func('fastrcnn', roi_feature_fastrcnn)
fastrcnn_label_logits, fastrcnn_box_logits = fastrcnn_outputs(
'fastrcnn/outputs', head_feature, cfg.DATA.NUM_CATEGORY)
fastrcnn_head = FastRCNNHead(
proposals, fastrcnn_box_logits, fastrcnn_label_logits,
gt_boxes,
tf.constant(cfg.FRCNN.BBOX_REG_WEIGHTS, dtype=tf.float32))
else:
def roi_func(boxes):
return multilevel_roi_align(features[:4], boxes, 7)
fastrcnn_head = CascadeRCNNHead(proposals, roi_func,
fastrcnn_head_func,
(gt_boxes, gt_labels),
image_shape2d,
cfg.DATA.NUM_CATEGORY)
if self.training:
all_losses = fastrcnn_head.losses()
if cfg.MODE_MASK:
gt_masks = targets[2]
# maskrcnn loss
roi_feature_maskrcnn = multilevel_roi_align(
features[:4],
proposals.fg_boxes(),
14,
name_scope='multilevel_roi_align_mask')
maskrcnn_head_func = getattr(model_mrcnn,
cfg.FPN.MRCNN_HEAD_FUNC)
mask_logits = maskrcnn_head_func(
'maskrcnn', roi_feature_maskrcnn,
cfg.DATA.NUM_CATEGORY) # #fg x #cat x 28 x 28
target_masks_for_fg = crop_and_resize(
tf.expand_dims(gt_masks, 1),
proposals.fg_boxes(),
proposals.fg_inds_wrt_gt,
28,
pad_border=False) # fg x 1x28x28
target_masks_for_fg = tf.squeeze(target_masks_for_fg, 1,
'sampled_fg_mask_targets')
all_losses.append(
maskrcnn_loss(mask_logits, proposals.fg_labels(),
target_masks_for_fg))
return all_losses
else:
decoded_boxes = fastrcnn_head.decoded_output_boxes()
decoded_boxes = clip_boxes(decoded_boxes,
image_shape2d,
name='fastrcnn_all_boxes')
label_scores = fastrcnn_head.output_scores(
name='fastrcnn_all_scores')
final_boxes, final_scores, final_labels = fastrcnn_predictions(
decoded_boxes, label_scores, name_scope='output')
if cfg.MODE_MASK:
# Cascade inference needs roi transform with refined boxes.
roi_feature_maskrcnn = multilevel_roi_align(
features[:4], final_boxes, 14)
maskrcnn_head_func = getattr(model_mrcnn,
cfg.FPN.MRCNN_HEAD_FUNC)
mask_logits = maskrcnn_head_func(
'maskrcnn', roi_feature_maskrcnn,
cfg.DATA.NUM_CATEGORY) # #fg x #cat x 28 x 28
indices = tf.stack([
tf.range(tf.size(final_labels)),
tf.cast(final_labels, tf.int32) - 1
],
axis=1)
final_mask_logits = tf.gather_nd(mask_logits,
indices) # #resultx28x28
tf.sigmoid(final_mask_logits, name='output/masks')
return []
def roi_heads(self, image, features, proposals, targets):
image_shape2d = tf.shape(image)[2:] # h,w
featuremap = features[0]
gt_boxes, gt_labels, *_ = targets
if self.training:
# sample proposal boxes in training
proposals = sample_fast_rcnn_targets(proposals.boxes, gt_boxes,
gt_labels)
# The boxes to be used to crop RoIs.
# Use all proposal boxes in inference
boxes_on_featuremap = proposals.boxes * (1.0 / cfg.RPN.ANCHOR_STRIDE)
roi_resized = roi_align(featuremap, boxes_on_featuremap, 14)
feature_fastrcnn = resnet_conv5(
roi_resized, cfg.BACKBONE.RESNET_NUM_BLOCKS[-1]) # nxcx7x7
# Keep C5 feature to be shared with mask branch
feature_gap = GlobalAvgPooling('gap',
feature_fastrcnn,
data_format='channels_first')
fastrcnn_label_logits, fastrcnn_box_logits = fastrcnn_outputs(
'fastrcnn', feature_gap, cfg.DATA.NUM_CATEGORY)
fastrcnn_head = FastRCNNHead(
proposals, fastrcnn_box_logits, fastrcnn_label_logits, gt_boxes,
tf.constant(cfg.FRCNN.BBOX_REG_WEIGHTS, dtype=tf.float32))
if self.training:
all_losses = fastrcnn_head.losses()
if cfg.MODE_MASK:
gt_masks = targets[2]
# maskrcnn loss
# In training, mask branch shares the same C5 feature.
fg_feature = tf.gather(feature_fastrcnn, proposals.fg_inds())
mask_logits = maskrcnn_upXconv_head(
'maskrcnn', fg_feature, cfg.DATA.NUM_CATEGORY,
num_convs=0) # #fg x #cat x 14x14
target_masks_for_fg = crop_and_resize(
tf.expand_dims(gt_masks, 1),
proposals.fg_boxes(),
proposals.fg_inds_wrt_gt,
14,
pad_border=False) # nfg x 1x14x14
target_masks_for_fg = tf.squeeze(target_masks_for_fg, 1,
'sampled_fg_mask_targets')
all_losses.append(
maskrcnn_loss(mask_logits, proposals.fg_labels(),
target_masks_for_fg))
return all_losses
else:
decoded_boxes = fastrcnn_head.decoded_output_boxes()
decoded_boxes = clip_boxes(decoded_boxes,
image_shape2d,
name='fastrcnn_all_boxes')
label_scores = fastrcnn_head.output_scores(
name='fastrcnn_all_scores')
final_boxes, final_scores, final_labels = fastrcnn_predictions(
decoded_boxes, label_scores, name_scope='output')
if cfg.MODE_MASK:
roi_resized = roi_align(
featuremap, final_boxes * (1.0 / cfg.RPN.ANCHOR_STRIDE),
14)
feature_maskrcnn = resnet_conv5(
roi_resized, cfg.BACKBONE.RESNET_NUM_BLOCKS[-1])
mask_logits = maskrcnn_upXconv_head(
'maskrcnn', feature_maskrcnn, cfg.DATA.NUM_CATEGORY,
0) # #result x #cat x 14x14
indices = tf.stack([
tf.range(tf.size(final_labels)),
tf.cast(final_labels, tf.int32) - 1
],
axis=1)
final_mask_logits = tf.gather_nd(mask_logits,
indices) # #resultx14x14
tf.sigmoid(final_mask_logits, name='output/masks')
return []