def step_epoch(self): ''' all models should be in cuda() write to coco_bbox_results.json ''' with torch.no_grad(): results = [] for i in range(len(self.dataset)): img, bbox, label, loc, scale = self.dataset[i] img = img.cuda().view(1, img.shape[0], img.shape[1], img.shape[2]) loc = loc.cuda().view(1, -1) temp = self.net(img, loc) pred_cls_i, pred_cls_p, pred_reg = get_pred(temp, self.net.nms_th, self.net.nms_iou) pred_cls_i = pred_cls_i[0].cpu() pred_cls_p = pred_cls_p[0].cpu() pred_reg = pred_reg[0].cpu() if pred_reg.shape[0] > 0: ymin, xmin, ymax, xmax = pred_reg.split([1, 1, 1, 1], dim=1) h, w = ymax - ymin, xmax - xmin pred_reg = torch.cat([xmin - loc[0, 1].cpu(), ymin - loc[0, 0].cpu(), w, h], dim=1) pred_reg = pred_reg / float(scale) for box_id in range(pred_reg.shape[0]): score = float(pred_cls_p[box_id]) label = int(pred_cls_i[box_id]) box = pred_reg[box_id, :] image_result = { 'image_id' : self.val_image_ids[i], 'category_id' : self.coco_labels[str(label)], 'score' : float(score), 'bbox' : box.tolist(), } results.append(image_result) print('step:%d/%d' % (i, len(self.dataset)), end='\r') json.dump(results, open('coco_bbox_results.json', 'w'), indent=4)
def step_epoch(self): ''' return map_mean, map_50, map_75 note: this function will set self.net.train() at last ''' with torch.no_grad(): self.net.eval() pred_bboxes, pred_labels, pred_scores, gt_bboxes, gt_labels = [], [], [], [], [] for i, (imgs, boxes, labels, locs, scales) in enumerate(self.loader): if i == 0: batch_size = int(imgs.shape[0]) temp = self.net(imgs, locs) pred_cls_i, pred_cls_p, pred_reg = get_pred( temp, self.net.module.nms_th, self.net.module.nms_iou) # DataParallel for idx in range(len(pred_cls_i)): pred_cls_i[idx] = pred_cls_i[idx].cpu().detach().numpy() pred_cls_p[idx] = pred_cls_p[idx].cpu().detach().numpy() pred_reg[idx] = pred_reg[idx].cpu().detach().numpy() _boxes, _label = [], [] for idx in range(boxes.shape[0]): mask = labels[idx] > 0 _boxes.append(boxes[idx][mask].detach().numpy()) _label.append(labels[idx][mask].detach().numpy()) pred_bboxes += pred_reg pred_labels += pred_cls_i pred_scores += pred_cls_p gt_bboxes += _boxes gt_labels += _label print(' Eval: {}/{}'.format(i * batch_size, len(self.dataset)), end='\r') ap_iou = [0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95] ap_res = [] for iou_th in ap_iou: res = eval_detection(pred_bboxes, pred_labels, pred_scores, gt_bboxes, gt_labels, iou_th=iou_th) ap_res.append(res) ap_sum = 0.0 for i in range(len(ap_res)): ap_sum += float(ap_res[i]['map']) map_mean = ap_sum / float(len(ap_res)) map_50 = float(ap_res[0]['map']) map_75 = float(ap_res[5]['map']) print('map_mean:', map_mean, 'map_50:', map_50, 'map_75:', map_75) self.net.train() return map_mean, map_50, map_75
def pred(self, img_pil): _boxes = torch.zeros(0, 4) img_pil, boxes, loc, scale = center_fix(img_pil, _boxes, self.net.view_size) img = transforms.ToTensor()(img_pil) img = self.normalizer(img).view(1, img.shape[0], img.shape[1], img.shape[2]) # img = img.cuda() loc = loc.view(1, -1) # .cuda() with torch.no_grad(): temp = self.net(img, loc) cls_i_preds, cls_p_preds, reg_preds = get_pred(temp, self.net.nms_th, self.net.nms_iou) reg_preds[0][:, 0::2] -= loc[0, 0] reg_preds[0][:, 1::2] -= loc[0, 1] reg_preds[0] /= scale return cls_i_preds[0], cls_p_preds[0], reg_preds[0]
def pred(self, img_pil): ''' all models should be in cuda() return cls_i_preds, cls_p_preds, reg_preds ''' _boxes = torch.zeros(0, 4) img_pil, boxes, loc, scale = center_fix(img_pil, _boxes, self.net.view_size) img = transforms.ToTensor()(img_pil) img = self.normalizer(img).view(1, img.shape[0], img.shape[1], img.shape[2]) img = img.cuda() loc = loc.view(1, -1).cuda() with torch.no_grad(): temp = self.net(img, loc) pred_cls_i, pred_cls_p, pred_reg = get_pred(temp, self.net.nms_th, self.net.nms_iou) pred_reg[0][:, 0::2] -= loc[0, 0] pred_reg[0][:, 1::2] -= loc[0, 1] pred_reg[0] /= scale return pred_cls_i[0], pred_cls_p[0], pred_reg[0]
shuffle=False, num_workers=0, collate_fn=dataset_eval.collate_fn) # Eval with torch.no_grad(): net.eval() pred_bboxes = [] pred_labels = [] pred_scores = [] gt_bboxes = [] gt_labels = [] for i, (img, bbox, label, scale, oh, ow) in enumerate(loader_eval): temp = net(img) cls_i_preds, cls_p_preds, reg_preds = get_pred(temp, net.module.nms_th, net.module.nms_iou, oh, ow) for idx in range(len(cls_i_preds)): cls_i_preds[idx] = cls_i_preds[idx].cpu().detach().numpy() cls_p_preds[idx] = cls_p_preds[idx].cpu().detach().numpy() reg_preds[idx] = reg_preds[idx].cpu().detach().numpy() _boxes = [] _label = [] for idx in range(bbox.shape[0]): mask = label[idx] > 0 _boxes.append(bbox[idx][mask].detach().numpy()) _label.append(label[idx][mask].detach().numpy()) pred_bboxes += reg_preds pred_labels += cls_i_preds pred_scores += cls_p_preds gt_bboxes += _boxes