segmentation.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 import chainer
7 try:
8  import chainer_mask_rcnn as cmr
9 except ImportError:
10  print('chainer_mask_rcnn cannot be imported.')
11 import cv2
12 import numpy as np
13 
14 
15 class SemanticSegmentationDataset(chainer.dataset.DatasetMixin):
16 
17  def __init__(self, root_dir):
18  self.root_dir = root_dir
19 
20  class_names_path = osp.join(root_dir, 'class_names.txt')
21  with open(class_names_path, 'r') as f:
22  class_names = f.readlines()
23  self.class_names = [name.rstrip() for name in class_names]
24 
25  self._images = []
26  self._labels = []
27  images_dir = osp.join(root_dir, 'JPEGImages')
28  labels_dir = osp.join(root_dir, 'SegmentationClass')
29  for image_ in sorted(os.listdir(images_dir)):
30  image_path = osp.join(images_dir, image_)
31  basename = image_.rstrip('.jpg')
32  label_path = osp.join(labels_dir, basename + '.npy')
33  self._images.append(image_path)
34  self._labels.append(label_path)
35 
36  def __len__(self):
37  return len(self._images)
38 
39  def get_example(self, i):
40  image_path = self._images[i]
41  label_path = self._labels[i]
42 
43  image = cv2.imread(image_path)
44  assert image.dtype == np.uint8
45  assert image.ndim == 3
46 
47  label = np.load(label_path)
48  assert label.dtype == np.int32
49  assert label.ndim == 2
50 
51  return image, label
52 
53 
54 class InstanceSegmentationDataset(chainer.dataset.DatasetMixin):
55 
56  def __init__(self, root_dir):
57  self.root_dir = root_dir
58 
59  class_names_path = osp.join(root_dir, 'class_names.txt')
60  with open(class_names_path, 'r') as f:
61  class_names = f.readlines()
62  # instance id 0 is '_background_' and should be ignored.
63  self.fg_class_names = [name.rstrip() for name in class_names][1:]
64 
65  self._images = []
66  self._class_labels = []
67  self._instance_labels = []
68  images_dir = osp.join(root_dir, 'JPEGImages')
69  class_labels_dir = osp.join(root_dir, 'SegmentationClass')
70  instance_labels_dir = osp.join(root_dir, 'SegmentationObject')
71  for image_ in sorted(os.listdir(images_dir)):
72  image_path = osp.join(images_dir, image_)
73  basename = image_.rstrip('.jpg')
74  class_label_path = osp.join(
75  class_labels_dir, basename + '.npy')
76  instance_label_path = osp.join(
77  instance_labels_dir, basename + '.npy')
78  self._images.append(image_path)
79  self._class_labels.append(class_label_path)
80  self._instance_labels.append(instance_label_path)
81 
82  def __len__(self):
83  return len(self._images)
84 
85  def get_example(self, i):
86  image_path = self._images[i]
87  class_label_path = self._class_labels[i]
88  instance_label_path = self._instance_labels[i]
89 
90  image = cv2.imread(image_path)
91  assert image.dtype == np.uint8
92  assert image.ndim == 3
93 
94  class_label = np.load(class_label_path)
95  assert class_label.dtype == np.int32
96  assert class_label.ndim == 2
97 
98  instance_label = np.load(instance_label_path)
99  # instance id 0 is '_background_' and should be ignored.
100  instance_label[instance_label == 0] = -1
101  assert instance_label.dtype == np.int32
102  assert instance_label.ndim == 2
103 
104  assert image.shape[:2] == class_label.shape == instance_label.shape
105 
106  labels, bboxes, masks = cmr.utils.label2instance_boxes(
107  label_instance=instance_label, label_class=class_label,
108  return_masks=True,
109  )
110  masks = masks.astype(np.int32, copy=False)
111  labels = labels.astype(np.int32, copy=False)
112  labels -= 1 # background: 0 -> -1
113  bboxes = bboxes.astype(np.float32, copy=False)
114 
115  return image, bboxes, labels, masks


jsk_recognition_utils
Author(s):
autogenerated on Fri Dec 6 2019 04:02:51