bbox_detection.py
Go to the documentation of this file.
1 from __future__ import print_function
2 
3 import os
4 import os.path as osp
5 
6 from chainercv.chainer_experimental.datasets.sliceable import GetterDataset
7 import cv2
8 import numpy as np
9 
10 import xml.etree.ElementTree as ET
11 
12 class BboxDetectionDataset(GetterDataset):
13 
14  def __init__(self, root_dir):
15  super(BboxDetectionDataset, self).__init__()
16  self.root_dir = root_dir
17 
18  class_names_path = osp.join(root_dir, 'class_names.txt')
19  with open(class_names_path, 'r') as f:
20  class_names = f.readlines()
21  self.class_names = [name.rstrip() for name in class_names]
22  self.fg_class_names = np.array(self.class_names[1:])
23 
24  self._imgs = []
25  self._annotations = []
26 
27  imgs_dir = osp.join(root_dir, 'JPEGImages')
28  annotation_dir = osp.join(root_dir, 'Annotations')
29 
30  for img_ in sorted(os.listdir(imgs_dir)):
31  img_path = osp.join(imgs_dir, img_)
32  basename = img_.rstrip('.jpg')
33  annotation_path = osp.join(annotation_dir, basename + '.xml')
34  self._imgs.append(img_path)
35  self._annotations.append(annotation_path)
36 
37  self.add_getter(('img', 'bbox', 'label'), self._get_example)
38 
39  def __len__(self):
40  return len(self._imgs)
41 
42  def _get_example(self, i):
43  img_path = self._imgs[i]
44  annotation_path = self._annotations[i]
45 
46  img = cv2.imread(img_path)
47  assert img.dtype == np.uint8
48  assert img.ndim == 3
49 
50  label = []
51  bbox = []
52 
53  tree = ET.parse(annotation_path)
54  root = tree.getroot()
55 
56  for obj in root.findall("object"):
57  label.append(np.where(self.fg_class_names == obj.find("name").text)[0][0])
58  bbox.append([obj.find("bndbox").find("ymin").text,
59  obj.find("bndbox").find("xmin").text,
60  obj.find("bndbox").find("ymax").text,
61  obj.find("bndbox").find("xmax").text])
62 
63  label = np.array(label, dtype=np.int32)
64  bbox = np.array(bbox, dtype=np.float32)
65 
66  img = img[:, :, ::-1]
67  img = img.transpose((2, 0, 1))
68  return img, bbox, label ## this is input


jsk_recognition_utils
Author(s):
autogenerated on Wed Mar 3 2021 04:00:53