def _generate_examples()

in projects/vision-ai-edge-platform/pipelines/segmentation/deeplabv3plus/trainer/vertexai_image_segmentation_dataset.py [0:0]

• 46 lines of code
• 9 McCabe index (conditional complexity)

    def _generate_examples(self, split):
        """Yields examples."""
        dataset_items = [{'uri': el['imageGcsUri'],
                          'classes': [pa['displayName']
                                      for pa in el['polygonAnnotations']],
                          'polygons': _vertexes_to_xy(el['polygonAnnotations'])}
                         for el in self.dataset_metadata[split]]

        #TODO: Turn into tf.function
        def create_mask(image, polygons, classes):
            size = tuple(image.shape[:-1])
            class_masks = [Image.new('L', size) for c in self.dataset_classes]
            class_draws = [ImageDraw.Draw(mask) for mask in class_masks]

            #TODO: consider cv2.drawContours()
            for i, xy in enumerate(polygons):
                class_id = self.dataset_classes.index(classes[i])
                scaled_xy = np.round(np.multiply(xy, size)).astype(int)
                tupled_xy = [(t[0], t[1]) for t in scaled_xy]
                class_draws[class_id].polygon(tupled_xy, fill = 255)

            mask_img = tf.concat([
                tf.cast(tf.keras.utils.img_to_array(mask), np.uint8)
                for mask in class_masks[1:]], axis=-1)

            #calculate background class channel
            mask_img = tf.concat([
                tf.expand_dims(
                    tf.math.scalar_mul(
                        255,
                        tf.cast(tf.reduce_sum(mask_img, axis=-1) == 0,
                                dtype=np.uint8)),
                    axis=-1),
                mask_img
            ], axis=-1)

            mask_img = tf.cast(
                tf.reshape(mask_img, (*size, len(self.dataset_classes))),
                np.uint8)

            return mask_img

        @tf.function
        def resize_image(img):
            return tf.cast(
                tf.image.resize_with_pad(img,
                                         self.image_height,
                                         self.image_width),
                np.uint8)


        for el in dataset_items:
            image = _load_image(el['uri'])
            mask = create_mask(image, el['polygons'], el['classes']).numpy()
            yield ('split_' + el['uri'],
                   {
                       'image': resize_image(image).numpy(),
                       'segmentation_mask': resize_image(mask).numpy()
                    })