def detect(self, frame):
# We get the detections from the frame
dets = self.model(frame)
results = []
batch_size = dets['class_ids'].shape[0]
for b in range(batch_size):
image_meta = {k: v[b].cpu().numpy() if not isinstance(v, list) else v[b]
for k, v in frame['image_meta'].items()}
det = {k: v[b] for k, v in dets.items()}
det = self.filter(det)
if det is None:
results.append({'image_meta': image_meta})
continue
det = {k: v.cpu().numpy() for k, v in det.items()}
det['boxes'] = boxes_postprocess(det['boxes'], image_meta)
det['image_meta'] = image_meta
results.append(det)
if self.args.debug==2:
image = image_postprocess(frame['image'][b].cpu().numpy().transpose(1, 2, 0), image_meta)
save_path = os.path.join(self.args.debug_dir, image_meta['image_id'] + '.png')
visualize_boxes(image, det['class_ids'], det['boxes'], det['scores'],
class_names = self.args.class_names,
save_path = save_path,
show = self.args.mode=='demo')
return results
def __getitem__(self, index):
"""Generates one sample of data"""
image, image_id = self.load_image(index)
gt_class_ids, gt_boxes = self.load_annotations(index)
image_meta = {
'index': index,
'image_id': image_id,
'orig_size': np.array(image.shape, dtype=np.int32)
}
image, image_meta, gt_boxes = self.preprocess(image, image_meta,
gt_boxes)
gt = self.prepare_annotations(gt_class_ids, gt_boxes)
inp = {
'image': image.transpose(2, 0, 1),
'image_meta': image_meta,
'gt': gt
}
if self.args.debug == 1:
image = image * image_meta['rgb_std'] + image_meta['rgb_mean']
save_path = os.path.join(self.args.debug_dir,
image_meta['image_id'] + '.png')
visualize_boxes(image,
gt_class_ids,
gt_boxes,
class_names=self.class_names,
save_path=save_path)
return inp