<gh_stars>1-10
"""
3D Faster R-CNN
Train on the CT scan of legs dataset
"""

import os
import sys
import json
import datetime
import numpy as np
import PIL
import imgaug

ROOT_DIR = os.getcwd()
sys.path.append(ROOT_DIR)

from faster_rcnn.config import Config
from faster_rcnn import model as modellib, utils

# Path to trained weights file
CT_MODEL_PATH = os.path.join(ROOT_DIR, "faster_rcnn_ct.h5")

DEFAULT_LOGS_DIR = os.path.join(ROOT_DIR, 'logs')

from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval


############################################################
#  Configurations
############################################################

class CTLandmarkConfig(Config):

    """
    Configuration for training of the landmark dataset/
    Derive from the base config an overrides some methods
    """
    # Give the configuration a recognizable name
    NAME = 'ct_landmark'

    IMAGES_PER_GPU = 2

    # Number of classes (including background)
    NUM_CLASSES = 1 + 34  # Background + balloon

    # Number of training steps per epoch
    STEPS_PER_EPOCH = 100

    # Skip detections with < 90% confidence
    DETECTION_MIN_CONFIDENCE = 0.9


############################################################
#  Dataset
############################################################

class CTLandmarkDataset(utils.Dataset):



    def load_ct(self, dataset_dir, subset, class_ids = None):
        """Load a subset of the CT scan dataset.
        dataset_dir: Root directory of the dataset.
        subset: Subset to load: train or val
        """
        coco = COCO("{}/annotations/{}".format(dataset_dir, subset))
        image_dir = "{}/{}".format(dataset_dir, subset)
        

        # Load all classes or a subset?
        if not class_ids:
            # All classes
            class_ids = sorted(coco.getCatIds())

        # All images or a subset?
        if class_ids:
            image_ids = []
            for id in class_ids:
                image_ids.extend(list(coco.getImgIds(catIds=[id])))
            # Remove duplicates
            image_ids = list(set(image_ids))
        else:
            # All images
            image_ids = list(coco.imgs.keys())

        # Add classes
        for i in class_ids:
            self.add_class("coco", i, coco.loadCats(i)[0]["name"])

        # Add images
        for i in image_ids:
            self.add_image(
                "coco", image_id=i,
                path=os.path.join(image_dir, coco.imgs[i]['file_name']),
                width=coco.imgs[i]["width"],
                height=coco.imgs[i]["height"],
                depth=coco.imgs[i]["depth"],
                annotations=coco.loadAnns(coco.getAnnIds(
                    imgIds=[i], catIds=class_ids, iscrowd=None)))
    
    # TODO: Check the below function once again
    def load_annotations(self, image_id):
        """Load bounding boxes for the give image.
        The annotations are in the form [z, y, x, depth,  height, width]
        """
        bboxes = []
        class_ids = []
        annots = self.image_info[image_id]['annotations']
        for annot in annots:
            class_id = annot['category_id']
            bbox = annot['bbox']
            class_ids.append(class_id)
            bboxes.append(bbox)
        bboxes = np.stack(bboxes, axis = 0)
        class_ids = np.array(class_ids, dtype=np.int32)
        return bboxes, class_ids 

############################################################
#  COCO Evaluation
############################################################

def build_coco_results(dataset, image_ids, rois, class_ids, scores, masks):
    """Arrange resutls to match COCO specs in http://cocodataset.org/#format
    """
    # If no results, return an empty list
    if rois is None:
        return []

    results = []
    for image_id in image_ids:
        # Loop through detections
        for i in range(rois.shape[0]):
            class_id = class_ids[i]
            score = scores[i]
            bbox = np.around(rois[i], 1)
            mask = masks[:, :, i]

            result = {
                "image_id": image_id,
                "category_id": dataset.get_source_class_id(class_id, "coco"),
                "bbox": [bbox[1], bbox[0], bbox[3] - bbox[1], bbox[2] - bbox[0]],
                "score": score,
                "segmentation": maskUtils.encode(np.asfortranarray(mask))
            }
            results.append(result)
    return results


def evaluate_coco(model, dataset, coco, eval_type="bbox", limit=0, image_ids=None):
    """Runs official COCO evaluation.
    dataset: A Dataset object with valiadtion data
    eval_type: "bbox" or "segm" for bounding box or segmentation evaluation
    limit: if not 0, it's the number of images to use for evaluation
    """
    # Pick COCO images from the dataset
    image_ids = image_ids or dataset.image_ids

    # Limit to a subset
    if limit:
        image_ids = image_ids[:limit]

    # Get corresponding COCO image IDs.
    coco_image_ids = [dataset.image_info[id]["id"] for id in image_ids]

    t_prediction = 0
    t_start = time.time()

    results = []
    for i, image_id in enumerate(image_ids):
        # Load image
        image = dataset.load_image(image_id)

        # Run detection
        t = time.time()
        r = model.detect([image], verbose=0)[0]
        t_prediction += (time.time() - t)

        # Convert results to COCO format
        # Cast masks to uint8 because COCO tools errors out on bool
        image_results = build_coco_results(dataset, coco_image_ids[i:i + 1],
                                           r["rois"], r["class_ids"],
                                           r["scores"],
                                           r["masks"].astype(np.uint8))
        results.extend(image_results)

    # Load results. This modifies results with additional attributes.
    coco_results = coco.loadRes(results)

    # Evaluate
    cocoEval = COCOeval(coco, coco_results, eval_type)
    cocoEval.params.imgIds = coco_image_ids
    cocoEval.evaluate()
    cocoEval.accumulate()
    cocoEval.summarize()

    print("Prediction time: {}. Average {}/image".format(
        t_prediction, t_prediction / len(image_ids)))
    print("Total time: ", time.time() - t_start)



############################################################
#  Training
############################################################


if __name__ == '__main__':
    import argparse

    # Parse command line arguments
    parser = argparse.ArgumentParser(
        description='Train 3D Faster R-CNN on 3D CT scans.')
    parser.add_argument("command",
                        metavar="<command>",
                        help="'train' or 'evaluate' on 3D CT scans")
    parser.add_argument('--dataset', required=True,
                        metavar="/path/to/ct_scans/",
                        help='Directory of the CT Scans dataset')
    parser.add_argument('--model', required=True,
                        metavar="/path/to/weights.h5",
                        help="Path to weights .h5 file")
    parser.add_argument('--logs', required=False,
                        default=DEFAULT_LOGS_DIR,
                        metavar="/path/to/logs/",
                        help='Logs and checkpoints directory (default=logs/)')
    parser.add_argument('--limit', required=False,
                        default=5,
                        metavar="<image count>",
                        help='Images to use for evaluation (default=500)')

    args = parser.parse_args()
    print("Command: ", args.command)
    print("Model: ", args.model)
    print("Dataset: ", args.dataset)
    print("Logs: ", args.logs)


    # Configurations
    if args.command == "train":
        config = CTLandmarkConfig()
    else:
        class InferenceConfig(CTLandmarkConfig):
            # Set batch size to 1 since we'll be running inference on
            # one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU
            GPU_COUNT = 1
            IMAGES_PER_GPU = 1
            DETECTION_MIN_CONFIDENCE = 0
        config = InferenceConfig()
    config.display()

    # Create model
    if args.command == "train":
        model = modellib.FasterRCNN(mode="training", config=config,
                                  model_dir=args.logs)
    else:
        model = modellib.FasterRCNN(mode="inference", config=config,
                                  model_dir=args.logs)

    # Select weights file to load        
    model_path = args.model

    # Load weights
    print("Loading weights ", model_path)
    model.load_weights(model_path, by_name=True)

    # Train or evaluate
    if args.command == "train":
        # Training dataset. Use the training set and 35K from the
        # validation set, as as in the Mask RCNN paper.
        dataset_train = CTLandmarkDataset()
        dataset_train.load_ct(args.dataset, subset = "train")
        dataset_train.prepare()

        # Validation dataset
        dataset_val = CTLandmarkDataset()
        val_type = "val" 
        dataset_val.load_ct(args.dataset, val_type)
        dataset_val.prepare()

        # Image Augmentation
        # Right/Left flip 50% of the time
        augmentation = imgaug.augmenters.Fliplr(0.5)

        # *** This training schedule is an example. Update to your needs ***

        # Training - Stage 1
        print("Training network heads")
        model.train(dataset_train, dataset_val,
                    learning_rate=config.LEARNING_RATE,
                    epochs=40,
                    layers='heads',
                    augmentation=augmentation)

        # Training - Stage 2
        # Finetune layers from ResNet stage 4 and up
        print("Fine tune Resnet stage 4 and up")
        model.train(dataset_train, dataset_val,
                    learning_rate=config.LEARNING_RATE,
                    epochs=120,
                    layers='4+',
                    augmentation=augmentation)

        # Training - Stage 3
        # Fine tune all layers
        print("Fine tune all layers")
        model.train(dataset_train, dataset_val,
                    learning_rate=config.LEARNING_RATE / 10,
                    epochs=160,
                    layers='all',
                    augmentation=augmentation)

    elif args.command == "evaluate":
        # Validation dataset
        dataset_val = CTLandmarkDataset()
        val_type = "val" if args.year in '2017' else "minival"
        coco = dataset_val.load_coco(args.dataset, val_type, year=args.year, return_coco=True, auto_download=args.download)
        dataset_val.prepare()
        print("Running COCO evaluation on {} images.".format(args.limit))
        evaluate_coco(model, dataset_val, coco, "bbox", limit=int(args.limit))
    else:
        print("'{}' is not recognized. "
              "Use 'train' or 'evaluate'".format(args.command))

    

"""
3D Faster R-CNN
Train on the CT scan of legs dataset
"""

import os
import sys
import json
import datetime
import numpy as np
import PIL
import imgaug

ROOT_DIR = os.getcwd()
sys.path.append(ROOT_DIR)

from faster_rcnn.config import Config
from faster_rcnn import model as modellib, utils

# Path to trained weights file
CT_MODEL_PATH = os.path.join(ROOT_DIR, "faster_rcnn_ct.h5")

DEFAULT_LOGS_DIR = os.path.join(ROOT_DIR, 'logs')

from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval


############################################################
#  Configurations
############################################################

class CTLandmarkConfig(Config):

    """
    Configuration for training of the landmark dataset/
    Derive from the base config an overrides some methods
    """
    # Give the configuration a recognizable name
    NAME = 'ct_landmark'

    IMAGES_PER_GPU = 2

    # Number of classes (including background)
    NUM_CLASSES = 1 + 34  # Background + balloon

    # Number of training steps per epoch
    STEPS_PER_EPOCH = 100

    # Skip detections with < 90% confidence
    DETECTION_MIN_CONFIDENCE = 0.9


############################################################
#  Dataset
############################################################

class CTLandmarkDataset(utils.Dataset):



    def load_ct(self, dataset_dir, subset, class_ids = None):
        """Load a subset of the CT scan dataset.
        dataset_dir: Root directory of the dataset.
        subset: Subset to load: train or val
        """
        coco = COCO("{}/annotations/{}".format(dataset_dir, subset))
        image_dir = "{}/{}".format(dataset_dir, subset)
        

        # Load all classes or a subset?
        if not class_ids:
            # All classes
            class_ids = sorted(coco.getCatIds())

        # All images or a subset?
        if class_ids:
            image_ids = []
            for id in class_ids:
                image_ids.extend(list(coco.getImgIds(catIds=[id])))
            # Remove duplicates
            image_ids = list(set(image_ids))
        else:
            # All images
            image_ids = list(coco.imgs.keys())

        # Add classes
        for i in class_ids:
            self.add_class("coco", i, coco.loadCats(i)[0]["name"])

        # Add images
        for i in image_ids:
            self.add_image(
                "coco", image_id=i,
                path=os.path.join(image_dir, coco.imgs[i]['file_name']),
                width=coco.imgs[i]["width"],
                height=coco.imgs[i]["height"],
                depth=coco.imgs[i]["depth"],
                annotations=coco.loadAnns(coco.getAnnIds(
                    imgIds=[i], catIds=class_ids, iscrowd=None)))
    
    # TODO: Check the below function once again
    def load_annotations(self, image_id):
        """Load bounding boxes for the give image.
        The annotations are in the form [z, y, x, depth,  height, width]
        """
        bboxes = []
        class_ids = []
        annots = self.image_info[image_id]['annotations']
        for annot in annots:
            class_id = annot['category_id']
            bbox = annot['bbox']
            class_ids.append(class_id)
            bboxes.append(bbox)
        bboxes = np.stack(bboxes, axis = 0)
        class_ids = np.array(class_ids, dtype=np.int32)
        return bboxes, class_ids 

############################################################
#  COCO Evaluation
############################################################

def build_coco_results(dataset, image_ids, rois, class_ids, scores, masks):
    """Arrange resutls to match COCO specs in http://cocodataset.org/#format
    """
    # If no results, return an empty list
    if rois is None:
        return []

    results = []
    for image_id in image_ids:
        # Loop through detections
        for i in range(rois.shape[0]):
            class_id = class_ids[i]
            score = scores[i]
            bbox = np.around(rois[i], 1)
            mask = masks[:, :, i]

            result = {
                "image_id": image_id,
                "category_id": dataset.get_source_class_id(class_id, "coco"),
                "bbox": [bbox[1], bbox[0], bbox[3] - bbox[1], bbox[2] - bbox[0]],
                "score": score,
                "segmentation": maskUtils.encode(np.asfortranarray(mask))
            }
            results.append(result)
    return results


def evaluate_coco(model, dataset, coco, eval_type="bbox", limit=0, image_ids=None):
    """Runs official COCO evaluation.
    dataset: A Dataset object with valiadtion data
    eval_type: "bbox" or "segm" for bounding box or segmentation evaluation
    limit: if not 0, it's the number of images to use for evaluation
    """
    # Pick COCO images from the dataset
    image_ids = image_ids or dataset.image_ids

    # Limit to a subset
    if limit:
        image_ids = image_ids[:limit]

    # Get corresponding COCO image IDs.
    coco_image_ids = [dataset.image_info[id]["id"] for id in image_ids]

    t_prediction = 0
    t_start = time.time()

    results = []
    for i, image_id in enumerate(image_ids):
        # Load image
        image = dataset.load_image(image_id)

        # Run detection
        t = time.time()
        r = model.detect([image], verbose=0)[0]
        t_prediction += (time.time() - t)

        # Convert results to COCO format
        # Cast masks to uint8 because COCO tools errors out on bool
        image_results = build_coco_results(dataset, coco_image_ids[i:i + 1],
                                           r["rois"], r["class_ids"],
                                           r["scores"],
                                           r["masks"].astype(np.uint8))
        results.extend(image_results)

    # Load results. This modifies results with additional attributes.
    coco_results = coco.loadRes(results)

    # Evaluate
    cocoEval = COCOeval(coco, coco_results, eval_type)
    cocoEval.params.imgIds = coco_image_ids
    cocoEval.evaluate()
    cocoEval.accumulate()
    cocoEval.summarize()

    print("Prediction time: {}. Average {}/image".format(
        t_prediction, t_prediction / len(image_ids)))
    print("Total time: ", time.time() - t_start)



############################################################
#  Training
############################################################


if __name__ == '__main__':
    import argparse

    # Parse command line arguments
    parser = argparse.ArgumentParser(
        description='Train 3D Faster R-CNN on 3D CT scans.')
    parser.add_argument("command",
                        metavar="<command>",
                        help="'train' or 'evaluate' on 3D CT scans")
    parser.add_argument('--dataset', required=True,
                        metavar="/path/to/ct_scans/",
                        help='Directory of the CT Scans dataset')
    parser.add_argument('--model', required=True,
                        metavar="/path/to/weights.h5",
                        help="Path to weights .h5 file")
    parser.add_argument('--logs', required=False,
                        default=DEFAULT_LOGS_DIR,
                        metavar="/path/to/logs/",
                        help='Logs and checkpoints directory (default=logs/)')
    parser.add_argument('--limit', required=False,
                        default=5,
                        metavar="<image count>",
                        help='Images to use for evaluation (default=500)')

    args = parser.parse_args()
    print("Command: ", args.command)
    print("Model: ", args.model)
    print("Dataset: ", args.dataset)
    print("Logs: ", args.logs)


    # Configurations
    if args.command == "train":
        config = CTLandmarkConfig()
    else:
        class InferenceConfig(CTLandmarkConfig):
            # Set batch size to 1 since we'll be running inference on
            # one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU
            GPU_COUNT = 1
            IMAGES_PER_GPU = 1
            DETECTION_MIN_CONFIDENCE = 0
        config = InferenceConfig()
    config.display()

    # Create model
    if args.command == "train":
        model = modellib.FasterRCNN(mode="training", config=config,
                                  model_dir=args.logs)
    else:
        model = modellib.FasterRCNN(mode="inference", config=config,
                                  model_dir=args.logs)

    # Select weights file to load        
    model_path = args.model

    # Load weights
    print("Loading weights ", model_path)
    model.load_weights(model_path, by_name=True)

    # Train or evaluate
    if args.command == "train":
        # Training dataset. Use the training set and 35K from the
        # validation set, as as in the Mask RCNN paper.
        dataset_train = CTLandmarkDataset()
        dataset_train.load_ct(args.dataset, subset = "train")
        dataset_train.prepare()

        # Validation dataset
        dataset_val = CTLandmarkDataset()
        val_type = "val" 
        dataset_val.load_ct(args.dataset, val_type)
        dataset_val.prepare()

        # Image Augmentation
        # Right/Left flip 50% of the time
        augmentation = imgaug.augmenters.Fliplr(0.5)

        # *** This training schedule is an example. Update to your needs ***

        # Training - Stage 1
        print("Training network heads")
        model.train(dataset_train, dataset_val,
                    learning_rate=config.LEARNING_RATE,
                    epochs=40,
                    layers='heads',
                    augmentation=augmentation)

        # Training - Stage 2
        # Finetune layers from ResNet stage 4 and up
        print("Fine tune Resnet stage 4 and up")
        model.train(dataset_train, dataset_val,
                    learning_rate=config.LEARNING_RATE,
                    epochs=120,
                    layers='4+',
                    augmentation=augmentation)

        # Training - Stage 3
        # Fine tune all layers
        print("Fine tune all layers")
        model.train(dataset_train, dataset_val,
                    learning_rate=config.LEARNING_RATE / 10,
                    epochs=160,
                    layers='all',
                    augmentation=augmentation)

    elif args.command == "evaluate":
        # Validation dataset
        dataset_val = CTLandmarkDataset()
        val_type = "val" if args.year in '2017' else "minival"
        coco = dataset_val.load_coco(args.dataset, val_type, year=args.year, return_coco=True, auto_download=args.download)
        dataset_val.prepare()
        print("Running COCO evaluation on {} images.".format(args.limit))
        evaluate_coco(model, dataset_val, coco, "bbox", limit=int(args.limit))
    else:
        print("'{}' is not recognized. "
              "Use 'train' or 'evaluate'".format(args.command))
3D Faster R-CNN Train on the CT scan of legs dataset # Path to trained weights file ############################################################ #  Configurations ############################################################ Configuration for training of the landmark dataset/     Derive from the base config an overrides some methods # Give the configuration a recognizable name # Number of classes (including background) # Background + balloon # Number of training steps per epoch # Skip detections with < 90% confidence ############################################################ #  Dataset ############################################################ Load a subset of the CT scan dataset.         dataset_dir: Root directory of the dataset.         subset: Subset to load: train or val # Load all classes or a subset? # All classes # All images or a subset? # Remove duplicates # All images # Add classes # Add images # TODO: Check the below function once again Load bounding boxes for the give image.         The annotations are in the form [z, y, x, depth,  height, width] ############################################################ #  COCO Evaluation ############################################################ Arrange resutls to match COCO specs in http://cocodataset.org/#format # If no results, return an empty list # Loop through detections Runs official COCO evaluation.     dataset: A Dataset object with valiadtion data     eval_type: "bbox" or "segm" for bounding box or segmentation evaluation     limit: if not 0, it's the number of images to use for evaluation # Pick COCO images from the dataset # Limit to a subset # Get corresponding COCO image IDs. # Load image # Run detection # Convert results to COCO format # Cast masks to uint8 because COCO tools errors out on bool # Load results. This modifies results with additional attributes. # Evaluate ############################################################ #  Training ############################################################ # Parse command line arguments # Configurations # Set batch size to 1 since we'll be running inference on # one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU # Create model # Select weights file to load # Load weights # Train or evaluate # Training dataset. Use the training set and 35K from the # validation set, as as in the Mask RCNN paper. # Validation dataset # Image Augmentation # Right/Left flip 50% of the time # *** This training schedule is an example. Update to your needs *** # Training - Stage 1 # Training - Stage 2 # Finetune layers from ResNet stage 4 and up # Training - Stage 3 # Fine tune all layers # Validation dataset