class Detector(object): def __init__(self, weights_path): self.model = SSD('test') self.model.cuda().eval() state = torch.load(weights_path, map_location=lambda storage, loc: storage) state = {key: value.float() for key, value in state.items()} self.model.load_state_dict(state) self.transform = GeneralizedRCNNTransform(DETECTOR_MIN_SIZE, DETECTOR_MAX_SIZE, DETECTOR_MEAN, DETECTOR_STD) self.transform.eval() def detect(self, images): images = torch.stack( [torch.from_numpy(image).cuda() for image in images]) images = images.transpose(1, 3).transpose(2, 3).float() original_image_sizes = [img.shape[-2:] for img in images] images, _ = self.transform(images, None) with torch.no_grad(): detections_batch = self.model(images.tensors).cpu().numpy() result = [] for detections, image_size in zip(detections_batch, images.image_sizes): scores = detections[1, :, 0] keep_idxs = scores > DETECTOR_THRESHOLD detections = detections[1, keep_idxs, :] detections = detections[:, [1, 2, 3, 4, 0]] detections[:, 0] *= image_size[1] detections[:, 1] *= image_size[0] detections[:, 2] *= image_size[1] detections[:, 3] *= image_size[0] result.append({ 'scores': torch.from_numpy(detections[:, 4]), 'boxes': torch.from_numpy(detections[:, :4]) }) result = self.transform.postprocess(result, images.image_sizes, original_image_sizes) return result
def evaluate_yolo_2017(model, data_loader, device): n_threads = torch.get_num_threads() # FIXME remove this and make paste_masks_in_image run on the GPU torch.set_num_threads(1) cpu_device = torch.device("cpu") model.eval() metric_logger = utils.MetricLogger(delimiter=" ") header = 'Test:' coco = get_coco_api_from_dataset(data_loader.dataset) iou_types = _get_iou_types(model) coco_evaluator = CocoEvaluator(coco, iou_types) transform = GeneralizedRCNNTransform(416, 416, [0, 0, 0], [1, 1, 1]) transform.eval() for image, targets in metric_logger.log_every(data_loader, 100, header): image = list(img.to(device) for img in image) original_image_sizes = [img.shape[-2:] for img in image] targets = [{k: v.to(device) for k, v in t.items()} for t in targets] torch.cuda.synchronize() model_time = time.time() transformed_img = transform(image) transformed_shape = transformed_img[0].tensors.shape[-2:] inf_out, _ = model(transformed_img[0].tensors) # Run NMS output = non_max_suppression(inf_out, conf_thres=0.001, iou_thres=0.6) # Statistics per image predictions = [] for si, pred in enumerate(output): prediction = {'boxes': [], 'labels': [], 'scores': []} if pred is None: continue # Append to text file # with open('test.txt', 'a') as file: # [file.write('%11.5g' * 7 % tuple(x) + '\n') for x in pred] # Clip boxes to image bounds clip_coords(pred, transformed_shape) # Append to pycocotools JSON dictionary # [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ... image_id = int(targets[si]['image_id']) box = pred[:, :4].clone() # xyxy # scale_coords(transformed_shape, box, shapes[si][0], shapes[si][1]) # to original shape # box = xyxy2xywh(box) # xywh # box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner for di, d in enumerate(pred): box_T = [floatn(x, 3) for x in box[di]] label = coco91class[int(d[5])] score = floatn(d[4], 5) prediction['boxes'].append(box_T) prediction['labels'].append(label) prediction['scores'].append(score) prediction['boxes'] = torch.tensor(prediction['boxes']) prediction['labels'] = torch.tensor(prediction['labels']) prediction['scores'] = torch.tensor(prediction['scores']) predictions.append(prediction) outputs = transform.postprocess(predictions, transformed_img[0].image_sizes, original_image_sizes) outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in predictions] model_time = time.time() - model_time res = { target["image_id"].item(): output for target, output in zip(targets, outputs) } evaluator_time = time.time() coco_evaluator.update(res) evaluator_time = time.time() - evaluator_time metric_logger.update(model_time=model_time, evaluator_time=evaluator_time) # gather the stats from all processes metric_logger.synchronize_between_processes() print("Averaged stats:", metric_logger) coco_evaluator.synchronize_between_processes() # accumulate predictions from all images coco_evaluator.accumulate() coco_evaluator.summarize() torch.set_num_threads(n_threads) return coco_evaluator