def gendefault(annot, idx, img_in, chg=None):
# Initialize sample parameters
c = annot['center'][idx]
s = annot['scale'][idx]
flip, r = False, 0
flip_idxs = ref.flipped_parts[annot.attrs['name']]
# Handle data augmentation
if chg is not None:
# Flipping
if 'flip' in chg:
if np.random.rand() < .5:
flip = True
# Scaling
if 'scale' in chg:
s *= min(
1 + chg['scale'],
max(1 - chg['scale'], (np.random.randn() * chg['scale']) + 1))
# Rotation
if 'rotate' in chg:
if chg['rotate'] == -1:
# Force vertical orientation
r = annot['torsoangle'][idx]
else:
r = np.random.randint(-chg['rotate'], chg['rotate'] + 1)
# Translation
if 'translate' in chg:
for i in xrange(2):
offset = np.random.randint(-chg['translate'],
chg['translate'] + 1)
c[i] += offset
# Generate input image
cropped = img.crop(img_in, c, s, ref.in_res, rot=r)
inp = np.zeros((3, ref.in_res[0], ref.in_res[1]))
for i in xrange(3):
inp[i, :, :] = cropped[:, :, i]
# Generate part heatmap output
num_parts = annot['part'].shape[1]
out = np.zeros((num_parts, ref.out_res[0], ref.out_res[1]))
for i in xrange(num_parts):
pt = annot['part'][idx, i]
if pt[0] > 0:
draw.gaussian(out[i], img.transform(pt, c, s, ref.out_res, rot=r),
2)
# Flip sample
if flip:
inp = np.array([np.fliplr(inp[i]) for i in xrange(len(inp))])
out = np.array(
[np.fliplr(out[flip_idxs[i]]) for i in xrange(len(out))])
return inp, out
def gendetect(annot, idx, img_in, chg=None):
img_c = [img_in.shape[1] / 2, img_in.shape[0] / 2]
img_s = max(img_in.shape) / 200
flip, r = False, 0
idxs = np.where(annot['index'][:] == annot['index'][idx])[0]
# Handle data augmentation
if chg is not None:
# Flipping
if 'flip' in chg:
if np.random.rand() < .5:
flip = True
# Scaling
if 'scale' in chg:
img_s *= min(
1 + chg['scale'],
max(1 - chg['scale'], (np.random.randn() * chg['scale']) + 1))
# Rotation
# if 'rotate' in chg:
# r = np.random.randint(-chg['rotate'], chg['rotate'] + 1)
# Translation
if 'translate' in chg:
for i in xrange(2):
offset = np.random.randint(-chg['translate'],
chg['translate'] + 1)
c[i] += offset
img_c[0] += img_s * np.random.randint(-10, 10)
img_c[1] += img_s * np.random.randint(-10, 10)
cropped = img.crop(img_in, img_c, img_s, ref.in_res)
inp = np.zeros((3, ref.in_res[0], ref.in_res[1]))
for i in xrange(3):
inp[i, :, :] = cropped[:, :, i]
out = np.zeros((2, ref.out_res[0], ref.out_res[1]))
for i in idxs:
pt = img.transform(annot['center'][i], img_c, img_s, ref.out_res)
draw.gaussian(out[0], pt, 1)
out[1, pt[1] - 1:pt[1] + 1,
pt[0] - 1:pt[0] + 1] = annot['scale'][i] / img_s
if flip:
inp = np.array([np.fliplr(inp[i]) for i in xrange(len(inp))])
out = np.array([np.fliplr(out[i]) for i in xrange(len(out))])
return inp, out
