def draw_kp(hm, pt1, pt2, boxes, img_path, pose_classes):
scores = tensor([[0.999]] * (boxes.shape[0]))
boxes = boxes.float()
preds_hm, preds_img, preds_scores = getPrediction(hm, pt1, pt2,
opt.inputResH,
opt.inputResW,
opt.outputResH,
opt.outputResW)
kps, score = pose_nms(boxes, scores, preds_img, preds_scores, pose_classes)
orig_img = cv2.imread(img_path)
if kps:
cond = True
kpv = KeyPointVisualizer()
img = kpv.vis_ske(orig_img, kps, score)
else:
img = orig_img
cond = False
return img, cond
def prediction(model):
model.eval()
dataset = Mscoco_minival()
minival_loader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=False,
num_workers=20,
pin_memory=True)
minival_loader_desc = tqdm(minival_loader)
final_result = []
tmp_inp = {}
for i, (inp, box, im_name, metaData) in enumerate(minival_loader_desc):
#inp = torch.autograd.Variable(inp.cuda(), volatile=True)
pt1, pt2, ori_inp = metaData
#with torch.autograd.profiler.profile(use_cuda=True) as prof:
if im_name[0] in tmp_inp.keys():
inps = tmp_inp[im_name[0]]['inps']
ori_inps = tmp_inp[im_name[0]]['ori_inps']
boxes = tmp_inp[im_name[0]]['boxes']
pt1s = tmp_inp[im_name[0]]['pt1s']
pt2s = tmp_inp[im_name[0]]['pt2s']
tmp_inp[im_name[0]]['inps'] = torch.cat((inps, inp), dim=0)
tmp_inp[im_name[0]]['pt1s'] = torch.cat((pt1s, pt1), dim=0)
tmp_inp[im_name[0]]['pt2s'] = torch.cat((pt2s, pt2), dim=0)
tmp_inp[im_name[0]]['ori_inps'] = torch.cat((ori_inps, ori_inp),
dim=0)
tmp_inp[im_name[0]]['boxes'] = torch.cat((boxes, box), dim=0)
else:
tmp_inp[im_name[0]] = {
'inps': inp,
'ori_inps': ori_inp,
'boxes': box,
'pt1s': pt1,
'pt2s': pt2
}
for im_name, item in tqdm(tmp_inp.items()):
inp = item['inps']
pt1 = item['pt1s']
pt2 = item['pt2s']
box = item['boxes']
ori_inp = item['ori_inps']
with torch.no_grad():
try:
kp_preds = model(inp)
kp_preds = kp_preds.data[:, :17, :]
except RuntimeError as e:
'''
Divide inputs into two batches
'''
assert str(e) == 'CUDA error: out of memory'
bn = inp.shape[0]
inp1 = inp[:bn // 2]
inp2 = inp[bn // 2:]
kp_preds1 = model(inp1)
kp_preds2 = model(inp2)
kp_preds = torch.cat((kp_preds1, kp_preds2), dim=0)
kp_preds = kp_preds.data[:, :17, :]
# kp_preds = gaussian_kernel(F.relu(kp_preds))
# Get predictions
# location prediction (n, kp, 2) | score prediction (n, kp, 1)
preds, preds_img, preds_scores = getPrediction(
kp_preds.cpu().data, pt1, pt2, opt.inputResH, opt.inputResW,
opt.outputResH, opt.outputResW)
result = pose_nms(box, preds_img, preds_scores)
result = {'imgname': im_name, 'result': result}
#img = display_frame(orig_img, result, opt.outputpath)
#ori_inp = np.transpose(
# ori_inp[0][:3].clone().numpy(), (1, 2, 0)) * 255
#img = vis_frame(ori_inp, result)
#cv2.imwrite(os.path.join(
# './val', 'vis', im_name), img)
final_result.append(result)
write_json(final_result, './val', for_eval=True)
return getmap()
def generateSampleBox(img_path,
bndbox,
part,
nJoints,
imgset,
scale_factor,
dataset,
train=True):
nJoints_coco = 17
nJoints_mpii = 16
img = load_image(img_path)
if train:
img[0].mul_(random.uniform(0.7, 1.3)).clamp_(0, 1)
img[1].mul_(random.uniform(0.7, 1.3)).clamp_(0, 1)
img[2].mul_(random.uniform(0.7, 1.3)).clamp_(0, 1)
ori_img = img.clone()
img[0].add_(-0.406)
img[1].add_(-0.457)
img[2].add_(-0.480)
upLeft = torch.Tensor((int(bndbox[0][0]), int(bndbox[0][1])))
bottomRight = torch.Tensor((int(bndbox[0][2]), int(bndbox[0][3])))
ht = bottomRight[1] - upLeft[1]
width = bottomRight[0] - upLeft[0]
imght = img.shape[1]
imgwidth = img.shape[2]
scaleRate = random.uniform(*scale_factor)
upLeft[0] = max(0, upLeft[0] - width * scaleRate / 2)
upLeft[1] = max(0, upLeft[1] - ht * scaleRate / 2)
bottomRight[0] = min(imgwidth - 1, bottomRight[0] + width * scaleRate / 2)
bottomRight[1] = min(imght - 1, bottomRight[1] + ht * scaleRate / 2)
# Doing Random Sample
if opt.addDPG:
PatchScale = random.uniform(0, 1)
if PatchScale > 0.85:
ratio = ht / width
if (width < ht):
patchWidth = PatchScale * width
patchHt = patchWidth * ratio
else:
patchHt = PatchScale * ht
patchWidth = patchHt / ratio
xmin = upLeft[0] + random.uniform(0, 1) * (width - patchWidth)
ymin = upLeft[1] + random.uniform(0, 1) * (ht - patchHt)
xmax = xmin + patchWidth + 1
ymax = ymin + patchHt + 1
else:
xmin = max(
1,
min(upLeft[0] + np.random.normal(-0.0142, 0.1158) * width,
imgwidth - 3))
ymin = max(
1,
min(upLeft[1] + np.random.normal(0.0043, 0.068) * ht,
imght - 3))
xmax = min(
max(xmin + 2,
bottomRight[0] + np.random.normal(0.0154, 0.1337) * width),
imgwidth - 3)
ymax = min(
max(ymin + 2,
bottomRight[1] + np.random.normal(-0.0013, 0.0711) * ht),
imght - 3)
upLeft[0] = xmin
upLeft[1] = ymin
bottomRight[0] = xmax
bottomRight[1] = ymax
# Counting Joints number
jointNum = 0
if imgset == 'coco':
for i in range(17):
if part[i][0] > 0 and part[i][0] > upLeft[0] and part[i][1] > upLeft[1] \
and part[i][0] < bottomRight[0] and part[i][1] < bottomRight[1]:
jointNum += 1
else:
for i in range(16):
if part[i][0] > 0 and part[i][0] > upLeft[0] and part[i][1] > upLeft[1] \
and part[i][0] < bottomRight[0] and part[i][1] < bottomRight[1]:
jointNum += 1
# Doing Random Crop
if opt.addDPG:
if jointNum > 13 and train:
switch = random.uniform(0, 1)
if switch > 0.96:
bottomRight[0] = (upLeft[0] + bottomRight[0]) / 2
bottomRight[1] = (upLeft[1] + bottomRight[1]) / 2
elif switch > 0.92:
upLeft[0] = (upLeft[0] + bottomRight[0]) / 2
bottomRight[1] = (upLeft[1] + bottomRight[1]) / 2
elif switch > 0.88:
upLeft[1] = (upLeft[1] + bottomRight[1]) / 2
bottomRight[0] = (upLeft[0] + bottomRight[0]) / 2
elif switch > 0.84:
upLeft[0] = (upLeft[0] + bottomRight[0]) / 2
upLeft[1] = (upLeft[1] + bottomRight[1]) / 2
elif switch > 0.80:
bottomRight[0] = (upLeft[0] + bottomRight[0]) / 2
elif switch > 0.76:
upLeft[0] = (upLeft[0] + bottomRight[0]) / 2
elif switch > 0.72:
bottomRight[1] = (upLeft[1] + bottomRight[1]) / 2
elif switch > 0.68:
upLeft[1] = (upLeft[1] + bottomRight[1]) / 2
ori_inp = cropBox(ori_img, upLeft, bottomRight, opt.inputResH,
opt.inputResW)
inp = cropBox(img, upLeft, bottomRight, opt.inputResH, opt.inputResW)
if jointNum == 0:
inp = torch.zeros(3, opt.inputResH, opt.inputResW)
out_bigcircle = torch.zeros(nJoints, opt.outputResH, opt.outputResW)
out_smallcircle = torch.zeros(nJoints, opt.outputResH, opt.outputResW)
out = torch.zeros(nJoints, opt.outputResH, opt.outputResW)
setMask = torch.zeros(nJoints, opt.outputResH, opt.outputResW)
# Draw Label
if imgset == 'coco':
for i in range(nJoints_coco):
if part[i][0] > 0 and part[i][0] > upLeft[0] and part[i][1] > upLeft[1] \
and part[i][0] < bottomRight[0] and part[i][1] < bottomRight[1]:
out_bigcircle[i] = drawBigCircle(
out_bigcircle[i],
transformBox(part[i], upLeft, bottomRight, opt.inputResH,
opt.inputResW, opt.outputResH,
opt.outputResW), opt.hmGauss * 2)
out_smallcircle[i] = drawSmallCircle(
out_smallcircle[i],
transformBox(part[i], upLeft, bottomRight, opt.inputResH,
opt.inputResW, opt.outputResH,
opt.outputResW), opt.hmGauss)
out[i] = drawGaussian(
out[i],
transformBox(part[i], upLeft, bottomRight, opt.inputResH,
opt.inputResW, opt.outputResH,
opt.outputResW), opt.hmGauss)
setMask[i].add_(1)
elif imgset == 'mpii':
for i in range(nJoints_coco, nJoints_coco + nJoints_mpii):
if part[i - nJoints_coco][0] > 0 and part[i - nJoints_coco][0] > upLeft[0] and part[i - nJoints_coco][1] > upLeft[1] \
and part[i - nJoints_coco][0] < bottomRight[0] and part[i - nJoints_coco][1] < bottomRight[1]:
out_bigcircle[i] = drawBigCircle(
out_bigcircle[i],
transformBox(part[i - nJoints_coco], upLeft, bottomRight,
opt.inputResH, opt.inputResW, opt.outputResH,
opt.outputResW), opt.hmGauss * 2)
out_smallcircle[i] = drawSmallCircle(
out_smallcircle[i],
transformBox(part[i - nJoints_coco], upLeft, bottomRight,
opt.inputResH, opt.inputResW, opt.outputResH,
opt.outputResW), opt.hmGauss)
out[i] = drawGaussian(
out[i],
transformBox(part[i - nJoints_coco], upLeft, bottomRight,
opt.inputResH, opt.inputResW, opt.outputResH,
opt.outputResW), opt.hmGauss)
setMask[i].add_(1)
else:
for i in range(nJoints_coco, nJoints_coco + nJoints_mpii):
if part[i - nJoints_coco][0] > 0 and part[i - nJoints_coco][0] > upLeft[0] and part[i - nJoints_coco][1] > upLeft[1] \
and part[i - nJoints_coco][0] < bottomRight[0] and part[i - nJoints_coco][1] < bottomRight[1]:
out_bigcircle[i] = drawBigCircle(
out_bigcircle[i],
transformBox(part[i - nJoints_coco], upLeft, bottomRight,
opt.inputResH, opt.inputResW, opt.outputResH,
opt.outputResW), opt.hmGauss * 2)
out_smallcircle[i] = drawSmallCircle(
out_smallcircle[i],
transformBox(part[i - nJoints_coco], upLeft, bottomRight,
opt.inputResH, opt.inputResW, opt.outputResH,
opt.outputResW), opt.hmGauss)
out[i] = drawGaussian(
out[i],
transformBox(part[i - nJoints_coco], upLeft, bottomRight,
opt.inputResH, opt.inputResW, opt.outputResH,
opt.outputResW), opt.hmGauss)
if i != 6 + nJoints_coco and i != 7 + nJoints_coco:
setMask[i].add_(1)
if opt.debug:
preds_hm, preds_img, preds_scores = getPrediction(
out.unsqueeze(0), upLeft.unsqueeze(0), bottomRight.unsqueeze(0),
opt.inputResH, opt.inputResW, opt.outputResH, opt.outputResW)
tmp_preds = preds_hm.mul(opt.inputResH / opt.outputResH)
drawCOCO(ori_inp.unsqueeze(0), tmp_preds, preds_scores)
if train:
# Flip
if random.uniform(0, 1) < 0.5:
inp = flip(inp)
ori_inp = flip(ori_inp)
out_bigcircle = shuffleLR(flip(out_bigcircle), dataset)
out_smallcircle = shuffleLR(flip(out_smallcircle), dataset)
out = shuffleLR(flip(out), dataset)
# Rotate
r = rnd(opt.rotate)
if random.uniform(0, 1) < 0.6:
r = 0
if r != 0:
inp = cv_rotate(inp, r, opt.inputResW, opt.inputResH)
out_bigcircle = cv_rotate(out_bigcircle, r, opt.outputResW,
opt.outputResH)
out_smallcircle = cv_rotate(out_smallcircle, r, opt.outputResW,
opt.outputResH)
out = cv_rotate(out, r, opt.outputResW, opt.outputResH)
return inp, out_bigcircle, out_smallcircle, out, setMask
