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