Load a subset of the Balloon dataset. dataset_dir: Root directory of the dataset. subset: Subset to load: train or val
(self, dataset_dir, subset)
| 86 | class BalloonDataset(utils.Dataset): |
| 87 | |
| 88 | def load_balloon(self, dataset_dir, subset): |
| 89 | """Load a subset of the Balloon dataset. |
| 90 | dataset_dir: Root directory of the dataset. |
| 91 | subset: Subset to load: train or val |
| 92 | """ |
| 93 | # Add classes. We have only one class to add. |
| 94 | self.add_class("balloon", 1, "balloon") |
| 95 | |
| 96 | # Train or validation dataset? |
| 97 | assert subset in ["train", "val"] |
| 98 | dataset_dir = os.path.join(dataset_dir, subset) |
| 99 | |
| 100 | # Load annotations |
| 101 | # VGG Image Annotator saves each image in the form: |
| 102 | # { 'filename': '28503151_5b5b7ec140_b.jpg', |
| 103 | # 'regions': { |
| 104 | # '0': { |
| 105 | # 'region_attributes': {}, |
| 106 | # 'shape_attributes': { |
| 107 | # 'all_points_x': [...], |
| 108 | # 'all_points_y': [...], |
| 109 | # 'name': 'polygon'}}, |
| 110 | # ... more regions ... |
| 111 | # }, |
| 112 | # 'size': 100202 |
| 113 | # } |
| 114 | # We mostly care about the x and y coordinates of each region |
| 115 | annotations = json.load(open(os.path.join(dataset_dir, "via_region_data.json"))) |
| 116 | annotations = list(annotations.values()) # don't need the dict keys |
| 117 | |
| 118 | # The VIA tool saves images in the JSON even if they don't have any |
| 119 | # annotations. Skip unannotated images. |
| 120 | annotations = [a for a in annotations if a['regions']] |
| 121 | |
| 122 | # Add images |
| 123 | for a in annotations: |
| 124 | # Get the x, y coordinaets of points of the polygons that make up |
| 125 | # the outline of each object instance. There are stores in the |
| 126 | # shape_attributes (see json format above) |
| 127 | polygons = [r['shape_attributes'] for r in a['regions'].values()] |
| 128 | |
| 129 | # load_mask() needs the image size to convert polygons to masks. |
| 130 | # Unfortunately, VIA doesn't include it in JSON, so we must read |
| 131 | # the image. This is only managable since the dataset is tiny. |
| 132 | image_path = os.path.join(dataset_dir, a['filename']) |
| 133 | image = skimage.io.imread(image_path) |
| 134 | height, width = image.shape[:2] |
| 135 | |
| 136 | self.add_image( |
| 137 | "balloon", |
| 138 | image_id=a['filename'], # use file name as a unique image id |
| 139 | path=image_path, |
| 140 | width=width, height=height, |
| 141 | polygons=polygons) |
| 142 | |
| 143 | def load_mask(self, image_id): |
| 144 | """Generate instance masks for an image. |