def append_visulize(self, imgpath, boxesPre, labelsPre, scoresPre, boxGT, labelGT, savepath=None): imPre = np.array(Image.open(imgpath).convert('RGB')) imGT = imPre.copy() # scoreGT = np.ones(shape=(labelGT.shape[0],)) visualize_boxes(image=imPre, boxes=boxesPre, labels=labelsPre, probs=scoresPre, class_labels=self.cateNames) # visualize_boxes(image=imGT, boxes=boxGT, labels=labelGT, probs=scoreGT, # class_labels=self.cateNames) whitepad = np.zeros(shape=(imPre.shape[0], 10, 3), dtype=np.uint8) imshow = np.concatenate((imGT, whitepad, imPre), axis=1) self.visual_imgs.append(imshow) # plt.imshow(imPre) # plt.show() if savepath: import os plt.imsave( os.path.join(savepath, '{}.png'.format(len(self.visual_imgs))), imshow)
def evaluate(cfg, args): device = select_device(args.device) # Initialize/load model if cfg.MODEL.META_ARCHITECTURE: # Initialize model model = YOLOv3(cfg).to(device) # Load weights if cfg.TEST.WEIGHTS.endswith(".pth"): state = torch.load(cfg.TEST.WEIGHTS, map_location=device) model.load_state_dict(state["state_dict"]) else: load_darknet_weights(model, cfg.TEST.WEIGHTS) if device.type != "cpu" and torch.cuda.device_count() > 1: model = nn.DataParallel(model) else: warnings.warn( "WARNNING: Backbone network cannot be empty! " f"Default load Darknet53 meta architecture for `{cfg.CONFIG_FILE}`!" ) model = YOLOv3(cfg).to(device) if cfg.TEST.TASK == "visual": images = os.listdir(os.path.join(os.getcwd(), "data", "test")) for filename in images: path = os.path.join( os.path.join(os.getcwd(), "data", "test", filename)) images = cv2.imread(path) assert images is not None bboxes_prd = Evaluator(model, cfg=cfg).get_bbox(images) if bboxes_prd.shape[0] != 0: boxes = bboxes_prd[..., :4] class_inds = bboxes_prd[..., 5].astype(np.int32) scores = bboxes_prd[..., 4] visualize_boxes(image=images, boxes=boxes, labels=class_inds, probs=scores, class_labels=cfg.CLASSES) path = os.path.join(f"./outputs/{filename}") cv2.imwrite(path, images) elif cfg.TEST.TASK == "eval": maps = 0. with torch.no_grad(): aps = Evaluator(model, visiual=True, cfg=cfg).calculate_aps(cfg.TEST.MULTI_SCALE, cfg.TEST.FLIP) for i in aps: print(f"{i:25s} --> mAP : {aps[i]:.4f}") maps += aps[i] maps = maps / len(cfg.CLASSES) print(f'mAP:{maps:.6f}') return maps
def _inference_epoch(self, imgdir, outdir=None, multiscale=True, flip=True): from dataset import get_imgdir from utils.visualize import visualize_boxes self.model.eval() dataloader = get_imgdir(imgdir, batch_size=8, net_size=self.net_size) for i, (imgpath, imgs, ori_shapes) in enumerate(dataloader): if i == 50: break ori_shapes = ori_shapes.cuda() if not multiscale: INPUT_SIZES = [self.net_size] else: INPUT_SIZES = [ self.net_size - 32, self.net_size, self.net_size + 32 ] pyramids = makeImgPyramids(imgs.numpy().transpose(0, 2, 3, 1), scales=INPUT_SIZES, flip=flip) # produce outputFeatures for each scale img2multi = defaultdict(list) for idx, pyramid in enumerate(pyramids): pyramid = torch.from_numpy(pyramid.transpose(0, 3, 1, 2)).cuda() with torch.no_grad(): grids = self.model(pyramid) for imgidx in range(imgs.shape[0]): img2multi[imgidx].append([grid[imgidx] for grid in grids]) # append prediction for each image per scale/flip for imgidx, scalegrids in img2multi.items(): allboxes = [] allscores = [] for _grids, _scale in zip(scalegrids[:len(INPUT_SIZES)], INPUT_SIZES): _boxes, _scores = predict_yolo( _grids, self.anchors, _scale, ori_shapes[imgidx], num_classes=self.num_classes) allboxes.append(_boxes) allscores.append(_scores) if flip: for _grids, _scale in zip(scalegrids[len(INPUT_SIZES):], INPUT_SIZES): _boxes, _scores = predict_yolo( _grids, self.anchors, _scale, ori_shapes[imgidx], num_classes=self.num_classes) _boxes = bbox_flip(_boxes.squeeze(0), flip_x=True, size=ori_shapes[imgidx]) _boxes = _boxes[np.newaxis, :] allboxes.append(_boxes) allscores.append(_scores) nms_boxes, nms_scores, nms_labels = torch_nms( torch.cat(allboxes, dim=1), torch.cat(allscores, dim=1), num_classes=self.num_classes) if nms_boxes is not None: detected_img = visualize_boxes( np.array(Image.open(imgpath[imgidx]).convert('RGB')), boxes=nms_boxes.cpu().numpy(), labels=nms_labels.cpu().numpy(), probs=nms_scores.cpu().numpy(), class_labels=self.labels) if outdir is not None: plt.imsave( os.path.join(outdir, imgpath[imgidx].split('/')[-1]), detected_img)