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