<gh_stars>0
import os
import json
import datetime
import numpy as np
import skimage.draw
from mrcnn.config import Config
from mrcnn import model as modellib, utils
from tensorflow.python.client import device_lib
import warnings

import tensorflow as tf

os.environ['CUDA_VISIBLE_DEVICES'] = "0"
os.environ["XLA_FLAGS"] = "--xla_gpu_cuda_data_dir=/usr/local/cuda"

warnings.filterwarnings('ignore', category=DeprecationWarning)
warnings.filterwarnings('ignore', category=FutureWarning)
warnings.filterwarnings('ignore', category=Warning)

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

# Path to trained weights file
WEIGHTS_PATH = "mask_rcnn_object_0049.h5"

# Directory to save logs and model checkpoints, if not provided through the command line argument --logs
DEFAULT_LOGS_DIR = "logs"


############################################################
#  Configurations
############################################################
class CustomConfig(Config):
    """Configuration for training on the chart  dataset.
    Derives from the base Config class and overrides some values.
    """
    # Give the configuration a recognizable name
    NAME = "object"

    GPU_COUNT = 1

    IMAGES_PER_GPU = 1

    # Number of classes (including background)
    NUM_CLASSES = 1 + 37

    # Number of training steps per epoch
    STEPS_PER_EPOCH = 100

    # Skip detections with < 90% confidence
    DETECTION_MIN_CONFIDENCE = 0.9


############################################################
#  Dataset
############################################################
class CustomDataset(utils.Dataset):

    def load_custom(self, dataset_dir, subset):
        CHARACTERS = [
            "1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11",
            "A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L",
            "M", "N", "O", "P", "Q", "R", "S", "T", "U", "V", "W", "X",
            "Y", "Z"
        ]
        self.add_class("object", 1, "1")
        self.add_class("object", 2, "2")
        self.add_class("object", 3, "3")
        self.add_class("object", 4, "4")
        self.add_class("object", 5, "5")
        self.add_class("object", 6, "6")
        self.add_class("object", 7, "7")
        self.add_class("object", 8, "8")
        self.add_class("object", 9, "9")
        self.add_class("object", 10, "0")
        self.add_class("object", 11, "11")
        self.add_class("object", 12, "A")
        self.add_class("object", 13, "B")
        self.add_class("object", 14, "C")
        self.add_class("object", 15, "D")
        self.add_class("object", 16, "E")
        self.add_class("object", 17, "F")
        self.add_class("object", 18, "G")
        self.add_class("object", 19, "H")
        self.add_class("object", 20, "I")
        self.add_class("object", 21, "J")
        self.add_class("object", 22, "K")
        self.add_class("object", 23, "L")
        self.add_class("object", 24, "M")
        self.add_class("object", 25, "N")
        self.add_class("object", 26, "O")
        self.add_class("object", 27, "P")
        self.add_class("object", 28, "Q")
        self.add_class("object", 29, "R")
        self.add_class("object", 30, "S")
        self.add_class("object", 31, "T")
        self.add_class("object", 32, "U")
        self.add_class("object", 33, "V")
        self.add_class("object", 34, "W")
        self.add_class("object", 35, "X")
        self.add_class("object", 36, "Y")
        self.add_class("object", 37, "Z")

        # Train or validation dataset?
        assert subset in ["train", "val"]
        dataset_dir = os.path.join(dataset_dir, subset)

        # Load annotations
        # VGG Image Annotator saves each image in the form:
        # { 'filename': '28503151_5b5b7ec140_b.jpg',
        #   'regions': {
        #       '0': {
        #           'region_attributes': {},
        #           'shape_attributes': {
        #               'all_points_x': [...],
        #               'all_points_y': [...],
        #               'name': 'polygon'}},
        #       ... more regions ...
        #   },
        #   'size': 100202
        # }
        # We mostly care about the x and y coordinates of each region
        annotations = json.load(open("annotattion.json"))
        # print(annotations1)
        annotations = list(annotations.values())  # don't need the dict keys

        # The VIA tool saves images in the JSON even if they don't have any
        # annotations. Skip unannotated images.
        annotations = [a for a in annotations if a['regions']]

        # Add images
        labels = []
        pid = []
        for a in annotations:
            print(a, '\n\n', a.keys(), '\n\n')
            # Get the x, y coordinates of points of the polygons that make up
            # the outline of each object instance. There are stores in the
            # shape_attributes (see json format above
            polygons = [r['shape_attributes'] for r in a['regions']]
            # print(polygons)
            objects = []
            for s in a['regions']:
                try:
                    objects.append(s['region_attributes']['CHARACTER'])
                except KeyError:
                    # print('=====>', s, )
                    print(9 / 0)
            objects = [s['region_attributes']['CHARACTER'] for s in a['regions']]
            # print("objects:", objects)
            labels.extend(objects)
            name_dict = {CHARACTERS[index - 1]: index for index in range(1, len(CHARACTERS) + 1)}
            pattern_ids = [name_dict[a] for a in objects]

            # print("numids", pattern_ids)
            image_path = os.path.join(dataset_dir, a['filename'])
            image = skimage.io.imread(image_path)
            height, width = image.shape[:2]
            # print('IMAGE DIMENSIONS ', width, height)
            self.add_image(
                "object",  ## for a single class just add the name here
                image_id=a['filename'],  # use file name as a unique image id
                path=image_path,
                width=width, height=height,
                polygons=polygons,
                pattern_ids=np.array(pattern_ids, dtype=np.int32))

    def load_mask(self, image_id):
        """Generate instance masks for an image.
       Returns:
        masks: A bool array of shape [height, width, instance count] with
            one mask per instance.
        class_ids: a 1D array of class IDs of the instance masks.
        """
        # If not a bottle dataset image, delegate to parent class.
        image_info = self.image_info[image_id]
        if image_info["source"] != "object":
            return super(self.__class__, self).load_mask(image_id)

        # Convert polygons to a bitmap mask of shape
        # [height, width, instance_count]
        info = self.image_info[image_id]
        if info["source"] != "object":
            return super(self.__class__, self).load_mask(image_id)
        num_ids = info['pattern_ids']
        mask = np.zeros([info["height"], info["width"], len(info["polygons"])],
                        dtype=np.uint8)
        for i, p in enumerate(info["polygons"]):
            # Get indexes of pixels inside the polygon and set them to 1
            rr, cc = skimage.draw.polygon(p['all_points_y'], p['all_points_x'])

            mask[rr, cc, i] = 1

        # Return mask, and array of class IDs of each instance. Since we have
        # one class ID only, we return an array of 1s
        # Map class names to class IDs.
        num_ids = np.array(num_ids, dtype=np.int32)
        return mask, num_ids

    def image_reference(self, image_id):
        """Return the path of the image."""
        info = self.image_info[image_id]
        if info["source"] == "object":
            return info["path"]
        else:
            super(self.__class__, self).image_reference(image_id)


def train(model):
    """Train the model."""
    # Training dataset.
    dataset_train = CustomDataset()
    data_path = 'Dataset'
    dataset_train.load_custom(data_path, "train")
    dataset_train.prepare()

    # Validation dataset
    dataset_val = CustomDataset()
    dataset_val.load_custom(data_path, "val")
    dataset_val.prepare()

    # *** This training schedule is an example. Update to your needs ***
    # Since we're using a very small dataset, and starting from
    # COCO trained weights, we don't need to train too long. Also,
    # no need to train all layers, just the heads should do it.
    print("Training network heads")

    model.train(dataset_train, dataset_val, learning_rate=config.LEARNING_RATE, epochs=50, layers='heads')


def color_splash(image, mask):
    """Apply color splash effect.
    image: RGB image [height, width, 3]
    mask: instance segmentation mask [height, width, instance count]
    Returns result image.
    """
    # Make a grayscale copy of the image. The grayscale copy still
    # has 3 RGB channels, though.
    gray = skimage.color.gray2rgb(skimage.color.rgb2gray(image)) * 255
    # We're treating all instances as one, so collapse the mask into one layer
    mask = (np.sum(mask, -1, keepdims=True) >= 1)
    # Copy color pixels from the original color image where mask is set
    if mask.shape[0] > 0:
        splash = np.where(mask, image, gray).astype(np.uint8)
    else:
        splash = gray
    return splash


def detect_and_color_splash(model, image_path=None, video_path=None):
    assert image_path or video_path

    # Image or video?
    if image_path:
        # Run model detection and generate the color splash effect
        print("Running on {}".format(args.image))
        # Read image
        image = skimage.io.imread(args.image)
        # Detect objects
        r = model.detect([image], verbose=1)[0]
        # Color splash
        splash = color_splash(image, r['masks'])
        # Save output
        file_name = "splash_{:%Y%m%dT%H%M%S}.png".format(datetime.datetime.now())
        skimage.io.imsave(file_name, splash)
    elif video_path:
        import cv2
        # Video capture
        vcapture = cv2.VideoCapture(video_path)
        width = int(vcapture.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(vcapture.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = vcapture.get(cv2.CAP_PROP_FPS)

        # Define codec and create video writer
        file_name = "splash_{:%Y%m%dT%H%M%S}.avi".format(datetime.datetime.now())
        vwriter = cv2.VideoWriter(file_name,
                                  cv2.VideoWriter_fourcc(*'MJPG'),
                                  fps, (width, height))

        count = 0
        success = True
        while success:
            print("frame: ", count)
            # Read next image
            success, image = vcapture.read()
            if success:
                # OpenCV returns images as BGR, convert to RGB
                image = image[..., ::-1]
                # Detect objects
                r = model.detect([image], verbose=0)[0]
                # Color splash
                splash = color_splash(image, r['masks'])
                # RGB -> BGR to save image to video
                splash = splash[..., ::-1]
                # Add image to video writer
                vwriter.write(splash)
                count += 1
        vwriter.release()
    print("Saved to ", file_name)


############################################################
#  Training
############################################################
import tensorflow as tf

if __name__ == '__main__':
    # Configurations
    config = CustomConfig()
    config.display()

    # Create model
    model = modellib.MaskRCNN(mode="training", config=config, model_dir='./logs')
    weights_path = WEIGHTS_PATH
    # Load weights
    print("Loading weights ", weights_path)
    with tf.device('/device:XLA_GPU:0'):
        model.load_weights(weights_path, by_name=True,
                           exclude=["mrcnn_class_logits", "mrcnn_bbox_fc", "mrcnn_bbox", "mrcnn_mask"])
        train(model)



import os
import json
import datetime
import numpy as np
import skimage.draw
from mrcnn.config import Config
from mrcnn import model as modellib, utils
from tensorflow.python.client import device_lib
import warnings

import tensorflow as tf

os.environ['CUDA_VISIBLE_DEVICES'] = "0"
os.environ["XLA_FLAGS"] = "--xla_gpu_cuda_data_dir=/usr/local/cuda"

warnings.filterwarnings('ignore', category=DeprecationWarning)
warnings.filterwarnings('ignore', category=FutureWarning)
warnings.filterwarnings('ignore', category=Warning)

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

# Path to trained weights file
WEIGHTS_PATH = "mask_rcnn_object_0049.h5"

# Directory to save logs and model checkpoints, if not provided through the command line argument --logs
DEFAULT_LOGS_DIR = "logs"


############################################################
#  Configurations
############################################################
class CustomConfig(Config):
    """Configuration for training on the chart  dataset.
    Derives from the base Config class and overrides some values.
    """
    # Give the configuration a recognizable name
    NAME = "object"

    GPU_COUNT = 1

    IMAGES_PER_GPU = 1

    # Number of classes (including background)
    NUM_CLASSES = 1 + 37

    # Number of training steps per epoch
    STEPS_PER_EPOCH = 100

    # Skip detections with < 90% confidence
    DETECTION_MIN_CONFIDENCE = 0.9


############################################################
#  Dataset
############################################################
class CustomDataset(utils.Dataset):

    def load_custom(self, dataset_dir, subset):
        CHARACTERS = [
            "1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11",
            "A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L",
            "M", "N", "O", "P", "Q", "R", "S", "T", "U", "V", "W", "X",
            "Y", "Z"
        ]
        self.add_class("object", 1, "1")
        self.add_class("object", 2, "2")
        self.add_class("object", 3, "3")
        self.add_class("object", 4, "4")
        self.add_class("object", 5, "5")
        self.add_class("object", 6, "6")
        self.add_class("object", 7, "7")
        self.add_class("object", 8, "8")
        self.add_class("object", 9, "9")
        self.add_class("object", 10, "0")
        self.add_class("object", 11, "11")
        self.add_class("object", 12, "A")
        self.add_class("object", 13, "B")
        self.add_class("object", 14, "C")
        self.add_class("object", 15, "D")
        self.add_class("object", 16, "E")
        self.add_class("object", 17, "F")
        self.add_class("object", 18, "G")
        self.add_class("object", 19, "H")
        self.add_class("object", 20, "I")
        self.add_class("object", 21, "J")
        self.add_class("object", 22, "K")
        self.add_class("object", 23, "L")
        self.add_class("object", 24, "M")
        self.add_class("object", 25, "N")
        self.add_class("object", 26, "O")
        self.add_class("object", 27, "P")
        self.add_class("object", 28, "Q")
        self.add_class("object", 29, "R")
        self.add_class("object", 30, "S")
        self.add_class("object", 31, "T")
        self.add_class("object", 32, "U")
        self.add_class("object", 33, "V")
        self.add_class("object", 34, "W")
        self.add_class("object", 35, "X")
        self.add_class("object", 36, "Y")
        self.add_class("object", 37, "Z")

        # Train or validation dataset?
        assert subset in ["train", "val"]
        dataset_dir = os.path.join(dataset_dir, subset)

        # Load annotations
        # VGG Image Annotator saves each image in the form:
        # { 'filename': '28503151_5b5b7ec140_b.jpg',
        #   'regions': {
        #       '0': {
        #           'region_attributes': {},
        #           'shape_attributes': {
        #               'all_points_x': [...],
        #               'all_points_y': [...],
        #               'name': 'polygon'}},
        #       ... more regions ...
        #   },
        #   'size': 100202
        # }
        # We mostly care about the x and y coordinates of each region
        annotations = json.load(open("annotattion.json"))
        # print(annotations1)
        annotations = list(annotations.values())  # don't need the dict keys

        # The VIA tool saves images in the JSON even if they don't have any
        # annotations. Skip unannotated images.
        annotations = [a for a in annotations if a['regions']]

        # Add images
        labels = []
        pid = []
        for a in annotations:
            print(a, '\n\n', a.keys(), '\n\n')
            # Get the x, y coordinates of points of the polygons that make up
            # the outline of each object instance. There are stores in the
            # shape_attributes (see json format above
            polygons = [r['shape_attributes'] for r in a['regions']]
            # print(polygons)
            objects = []
            for s in a['regions']:
                try:
                    objects.append(s['region_attributes']['CHARACTER'])
                except KeyError:
                    # print('=====>', s, )
                    print(9 / 0)
            objects = [s['region_attributes']['CHARACTER'] for s in a['regions']]
            # print("objects:", objects)
            labels.extend(objects)
            name_dict = {CHARACTERS[index - 1]: index for index in range(1, len(CHARACTERS) + 1)}
            pattern_ids = [name_dict[a] for a in objects]

            # print("numids", pattern_ids)
            image_path = os.path.join(dataset_dir, a['filename'])
            image = skimage.io.imread(image_path)
            height, width = image.shape[:2]
            # print('IMAGE DIMENSIONS ', width, height)
            self.add_image(
                "object",  ## for a single class just add the name here
                image_id=a['filename'],  # use file name as a unique image id
                path=image_path,
                width=width, height=height,
                polygons=polygons,
                pattern_ids=np.array(pattern_ids, dtype=np.int32))

    def load_mask(self, image_id):
        """Generate instance masks for an image.
       Returns:
        masks: A bool array of shape [height, width, instance count] with
            one mask per instance.
        class_ids: a 1D array of class IDs of the instance masks.
        """
        # If not a bottle dataset image, delegate to parent class.
        image_info = self.image_info[image_id]
        if image_info["source"] != "object":
            return super(self.__class__, self).load_mask(image_id)

        # Convert polygons to a bitmap mask of shape
        # [height, width, instance_count]
        info = self.image_info[image_id]
        if info["source"] != "object":
            return super(self.__class__, self).load_mask(image_id)
        num_ids = info['pattern_ids']
        mask = np.zeros([info["height"], info["width"], len(info["polygons"])],
                        dtype=np.uint8)
        for i, p in enumerate(info["polygons"]):
            # Get indexes of pixels inside the polygon and set them to 1
            rr, cc = skimage.draw.polygon(p['all_points_y'], p['all_points_x'])

            mask[rr, cc, i] = 1

        # Return mask, and array of class IDs of each instance. Since we have
        # one class ID only, we return an array of 1s
        # Map class names to class IDs.
        num_ids = np.array(num_ids, dtype=np.int32)
        return mask, num_ids

    def image_reference(self, image_id):
        """Return the path of the image."""
        info = self.image_info[image_id]
        if info["source"] == "object":
            return info["path"]
        else:
            super(self.__class__, self).image_reference(image_id)


def train(model):
    """Train the model."""
    # Training dataset.
    dataset_train = CustomDataset()
    data_path = 'Dataset'
    dataset_train.load_custom(data_path, "train")
    dataset_train.prepare()

    # Validation dataset
    dataset_val = CustomDataset()
    dataset_val.load_custom(data_path, "val")
    dataset_val.prepare()

    # *** This training schedule is an example. Update to your needs ***
    # Since we're using a very small dataset, and starting from
    # COCO trained weights, we don't need to train too long. Also,
    # no need to train all layers, just the heads should do it.
    print("Training network heads")

    model.train(dataset_train, dataset_val, learning_rate=config.LEARNING_RATE, epochs=50, layers='heads')


def color_splash(image, mask):
    """Apply color splash effect.
    image: RGB image [height, width, 3]
    mask: instance segmentation mask [height, width, instance count]
    Returns result image.
    """
    # Make a grayscale copy of the image. The grayscale copy still
    # has 3 RGB channels, though.
    gray = skimage.color.gray2rgb(skimage.color.rgb2gray(image)) * 255
    # We're treating all instances as one, so collapse the mask into one layer
    mask = (np.sum(mask, -1, keepdims=True) >= 1)
    # Copy color pixels from the original color image where mask is set
    if mask.shape[0] > 0:
        splash = np.where(mask, image, gray).astype(np.uint8)
    else:
        splash = gray
    return splash


def detect_and_color_splash(model, image_path=None, video_path=None):
    assert image_path or video_path

    # Image or video?
    if image_path:
        # Run model detection and generate the color splash effect
        print("Running on {}".format(args.image))
        # Read image
        image = skimage.io.imread(args.image)
        # Detect objects
        r = model.detect([image], verbose=1)[0]
        # Color splash
        splash = color_splash(image, r['masks'])
        # Save output
        file_name = "splash_{:%Y%m%dT%H%M%S}.png".format(datetime.datetime.now())
        skimage.io.imsave(file_name, splash)
    elif video_path:
        import cv2
        # Video capture
        vcapture = cv2.VideoCapture(video_path)
        width = int(vcapture.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(vcapture.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = vcapture.get(cv2.CAP_PROP_FPS)

        # Define codec and create video writer
        file_name = "splash_{:%Y%m%dT%H%M%S}.avi".format(datetime.datetime.now())
        vwriter = cv2.VideoWriter(file_name,
                                  cv2.VideoWriter_fourcc(*'MJPG'),
                                  fps, (width, height))

        count = 0
        success = True
        while success:
            print("frame: ", count)
            # Read next image
            success, image = vcapture.read()
            if success:
                # OpenCV returns images as BGR, convert to RGB
                image = image[..., ::-1]
                # Detect objects
                r = model.detect([image], verbose=0)[0]
                # Color splash
                splash = color_splash(image, r['masks'])
                # RGB -> BGR to save image to video
                splash = splash[..., ::-1]
                # Add image to video writer
                vwriter.write(splash)
                count += 1
        vwriter.release()
    print("Saved to ", file_name)


############################################################
#  Training
############################################################
import tensorflow as tf

if __name__ == '__main__':
    # Configurations
    config = CustomConfig()
    config.display()

    # Create model
    model = modellib.MaskRCNN(mode="training", config=config, model_dir='./logs')
    weights_path = WEIGHTS_PATH
    # Load weights
    print("Loading weights ", weights_path)
    with tf.device('/device:XLA_GPU:0'):
        model.load_weights(weights_path, by_name=True,
                           exclude=["mrcnn_class_logits", "mrcnn_bbox_fc", "mrcnn_bbox", "mrcnn_mask"])
        train(model)
# Path to trained weights file # Directory to save logs and model checkpoints, if not provided through the command line argument --logs ############################################################ #  Configurations ############################################################ Configuration for training on the chart  dataset.     Derives from the base Config class and overrides some values. # Give the configuration a recognizable name # Number of classes (including background) # Number of training steps per epoch # Skip detections with < 90% confidence ############################################################ #  Dataset ############################################################ # Train or validation dataset? # Load annotations # VGG Image Annotator saves each image in the form: # { 'filename': '28503151_5b5b7ec140_b.jpg', #   'regions': { #       '0': { #           'region_attributes': {}, #           'shape_attributes': { #               'all_points_x': [...], #               'all_points_y': [...], #               'name': 'polygon'}}, #       ... more regions ... #   }, #   'size': 100202 # } # We mostly care about the x and y coordinates of each region # print(annotations1) # don't need the dict keys # The VIA tool saves images in the JSON even if they don't have any # annotations. Skip unannotated images. # Add images # Get the x, y coordinates of points of the polygons that make up # the outline of each object instance. There are stores in the # shape_attributes (see json format above # print(polygons) # print('=====>', s, ) # print("objects:", objects) # print("numids", pattern_ids) # print('IMAGE DIMENSIONS ', width, height) ## for a single class just add the name here # use file name as a unique image id Generate instance masks for an image.        Returns:         masks: A bool array of shape [height, width, instance count] with             one mask per instance.         class_ids: a 1D array of class IDs of the instance masks. # If not a bottle dataset image, delegate to parent class. # Convert polygons to a bitmap mask of shape # [height, width, instance_count] # Get indexes of pixels inside the polygon and set them to 1 # Return mask, and array of class IDs of each instance. Since we have # one class ID only, we return an array of 1s # Map class names to class IDs. Return the path of the image. Train the model. # Training dataset. # Validation dataset # *** This training schedule is an example. Update to your needs *** # Since we're using a very small dataset, and starting from # COCO trained weights, we don't need to train too long. Also, # no need to train all layers, just the heads should do it. Apply color splash effect.     image: RGB image [height, width, 3]     mask: instance segmentation mask [height, width, instance count]     Returns result image. # Make a grayscale copy of the image. The grayscale copy still # has 3 RGB channels, though. # We're treating all instances as one, so collapse the mask into one layer # Copy color pixels from the original color image where mask is set # Image or video? # Run model detection and generate the color splash effect # Read image # Detect objects # Color splash # Save output # Video capture # Define codec and create video writer # Read next image # OpenCV returns images as BGR, convert to RGB # Detect objects # Color splash # RGB -> BGR to save image to video # Add image to video writer ############################################################ #  Training ############################################################ # Configurations # Create model # Load weights