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