# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Mask R-CNN - Inspect Ballon Trained Model
#
# Code and visualizations to test, debug, and evaluate the Mask R-CNN model.

# +
import os
import sys
import random
import math
import re
import time
import numpy as np
import tensorflow as tf
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.patches as patches

# Root directory of the project
ROOT_DIR = os.path.abspath("../../")

# Import Mask RCNN
sys.path.append(ROOT_DIR)  # To find local version of the library
from mrcnn import utils
from mrcnn import visualize
from mrcnn.visualize import display_images
import mrcnn.model as modellib
from mrcnn.model import log

from samples.balloon import balloon

# %matplotlib inline 

# Directory to save logs and trained model
MODEL_DIR = os.path.join(ROOT_DIR, "logs")

# Path to Ballon trained weights
# You can download this file from the Releases page
# https://github.com/matterport/Mask_RCNN/releases
#BALLON_WEIGHTS_PATH = "/path/to/mask_rcnn_balloon.h5"  # TODO: update this path
# -

# ## Configurations

config = balloon.BalloonConfig()
BALLOON_DIR = os.path.join(ROOT_DIR, "datasets/balloon")


# +
# Override the training configurations with a few
# changes for inferencing.
class InferenceConfig(config.__class__):
    # Run detection on one image at a time
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1

config = InferenceConfig()
config.display()
# -

# ## Notebook Preferences

# +
# Device to load the neural network on.
# Useful if you're training a model on the same 
# machine, in which case use CPU and leave the
# GPU for training.
DEVICE = "/cpu:0"  # /cpu:0 or /gpu:0

# Inspect the model in training or inference modes
# values: 'inference' or 'training'
# TODO: code for 'training' test mode not ready yet
TEST_MODE = "training"


# -

def get_ax(rows=1, cols=1, size=16):
    """Return a Matplotlib Axes array to be used in
    all visualizations in the notebook. Provide a
    central point to control graph sizes.
    
    Adjust the size attribute to control how big to render images
    """
    _, ax = plt.subplots(rows, cols, figsize=(size*cols, size*rows))
    return ax


# ## Load Training Dataset

# +
# Load validation dataset
dataset = balloon.BalloonDataset()
dataset.load_balloon(BALLOON_DIR, "train")

# Must call before using the dataset
dataset.prepare()

print("Images: {}\nClasses: {}".format(len(dataset.image_ids), dataset.class_names))
# -

# ## Load Model

# Create model in inference mode
with tf.device(DEVICE):
    model = modellib.MaskRCNN(mode="training", model_dir=MODEL_DIR,
                              config=config)

# ## Step by Step Prediction

# ## Stage 3: Generating Masks
#
# This stage takes the detections (refined bounding boxes and class IDs) from the previous layer and runs the mask head to generate segmentation masks for every instance.

# Masks
det_boxes = utils.denorm_boxes(mrcnn["detections"][0, :, :4], image.shape[:2])
det_mask_specific = np.array([mrcnn["masks"][0, i, :, :, c] 
                              for i, c in enumerate(det_class_ids)])
det_masks = np.array([utils.unmold_mask(m, det_boxes[i], image.shape)
                      for i, m in enumerate(det_mask_specific)])
log("det_mask_specific", det_mask_specific)
log("det_masks", det_masks)

image_id = random.choice(dataset.image_ids)
image, image_meta, gt_class_id, gt_bbox, gt_mask, orig_mask =\
    modellib.load_image_gt(dataset, config, image_id, use_mini_mask=False)
info = dataset.image_info[image_id]
print("image ID: {}.{} ({}) {}".format(info["source"], info["id"], image_id, 
                                       dataset.image_reference(image_id)))

# ### 3.a Mask Targets
#
# These are the training targets for the mask branch

display_images(np.transpose(gt_mask, [2, 0, 1]), cmap="Blues")


# ### 3.b Predicted Masks

def build_batch_data(image, image_meta, gt_class_ids, gt_boxes, gt_masks, gt_orig_masks):
    
    backbone_shapes = modellib.compute_backbone_shapes(config, config.IMAGE_SHAPE)
    anchors = utils.generate_pyramid_anchors(config.RPN_ANCHOR_SCALES,
                                             config.RPN_ANCHOR_RATIOS,
                                             backbone_shapes,
                                             config.BACKBONE_STRIDES,
                                             config.RPN_ANCHOR_STRIDE)
    
    
    
    rpn_match, rpn_bbox = modellib.build_rpn_targets(image.shape, anchors,
                                                    gt_class_ids, gt_boxes, config)
    inst_ids = modellib.build_kira_targets(image.shape, gt_class_ids, gt_orig_masks, config)
    
    batch_image_meta = np.zeros(
        (1,) + image_meta.shape, dtype=image_meta.dtype)
    batch_rpn_match = np.zeros(
        [1, anchors.shape[0], 1], dtype=rpn_match.dtype)
    batch_rpn_bbox = np.zeros(
        [1, config.RPN_TRAIN_ANCHORS_PER_IMAGE, 4], dtype=rpn_bbox.dtype)
    batch_images = np.zeros(
        (1,) + image.shape, dtype=np.float32)
    batch_gt_class_ids = np.zeros(
        (1, config.MAX_GT_INSTANCES), dtype=np.int32)
    batch_gt_boxes = np.zeros(
        (1, config.MAX_GT_INSTANCES, 4), dtype=np.int32)
    batch_gt_masks = np.zeros(
        (1, gt_masks.shape[0], gt_masks.shape[1],
         config.MAX_GT_INSTANCES), dtype=gt_masks.dtype)
    
    batch_inst_ids = np.zeros(
        (1, image.shape[0], image.shape[1]), dtype=np.int32)
    
    b= 0
    batch_image_meta[b] = image_meta
    batch_rpn_match[b] = rpn_match[:, np.newaxis]
    batch_rpn_bbox[b] = rpn_bbox
    batch_images[b] = modellib.mold_image(image.astype(np.float32), config)
    batch_gt_class_ids[b, :gt_class_ids.shape[0]] = gt_class_ids
    batch_gt_boxes[b, :gt_boxes.shape[0]] = gt_boxes
    batch_gt_masks[b, :, :, :gt_masks.shape[-1]] = gt_masks
    
    batch_inst_ids[b] = inst_ids
    
    inputs = [batch_images, batch_image_meta, batch_rpn_match, batch_rpn_bbox,
                          batch_gt_class_ids, batch_gt_boxes, batch_gt_masks, batch_inst_ids]
    return inputs

from collections import OrderedDict
import keras.backend as K
def run_train_graph(input_batch, outputs, image_metas=None):
    """Runs a sub-set of the computation graph that computes the given
    outputs.
    image_metas: If provided, the images are assumed to be already
        molded (i.e. resized, padded, and normalized)
    outputs: List of tuples (name, tensor) to compute. The tensors are
        symbolic TensorFlow tensors and the names are for easy tracking.
    Returns an ordered dict of results. Keys are the names received in the
    input and values are Numpy arrays.
    """
    keras_model = model.keras_model

    # Organize desired outputs into an ordered dict
    outputs = OrderedDict(outputs)
    for o in outputs.values():
        assert o is not None

    # Build a Keras function to run parts of the computation graph
    inputs = keras_model.inputs
    if keras_model.uses_learning_phase and not isinstance(K.learning_phase(), int):
        inputs += [K.learning_phase()]
    kf = K.function(keras_model.inputs, list(outputs.values()))

    # Prepare inputs
    if image_metas is None:
        print(input_batch[0].shape)
        molded_images, image_metas, _ = model.mold_inputs([input_batch[0][0]])
    else:
        molded_images = images
    image_shape = molded_images[0].shape
    # Anchors
    anchors = model.get_anchors(image_shape)
    # Duplicate across the batch dimension because Keras requires it
    # TODO: can this be optimized to avoid duplicating the anchors?
    anchors = np.broadcast_to(anchors, (model.config.BATCH_SIZE,) + anchors.shape)
    model_in = [molded_images, image_metas, anchors]

    # Run inference
    if keras_model.uses_learning_phase and not isinstance(K.learning_phase(), int):
        model_in.append(0.)
    outputs_np = kf(model_in)

    # Pack the generated Numpy arrays into a a dict and log the results.
    outputs_np = OrderedDict([(k, v)
                              for k, v in zip(outputs.keys(), outputs_np)])
    for k, v in outputs_np.items():
        log(k, v)
    return outputs_np


# +
inputs = build_batch_data(image, image_meta, gt_class_id, gt_bbox, gt_mask, orig_mask)


# Get predictions of mask head
mrcnn = run_train_graph(inputs, [
    ("input_image",        model.keras_model.get_layer("input_image").output)
])

# Get detection class IDs. Trim zero padding.
det_class_ids = mrcnn['detections'][0, :, 4].astype(np.int32)
det_count = np.where(det_class_ids == 0)[0][0]
det_class_ids = det_class_ids[:det_count]

print("{} detections: {}".format(
    det_count, np.array(dataset.class_names)[det_class_ids]))
# -

display_images(det_mask_specific[:4] * 255, cmap="Blues", interpolation="none")

display_images(det_masks[:4] * 255, cmap="Blues", interpolation="none")

# ## Visualize Activations
#
# In some cases it helps to look at the output from different layers and visualize them to catch issues and odd patterns.

# Get activations of a few sample layers
activations = model.run_graph([image], [
    ("input_image",        model.keras_model.get_layer("input_image").output),
    ("res2c_out",          model.keras_model.get_layer("res2c_out").output),
    ("res3c_out",          model.keras_model.get_layer("res3c_out").output),
    ("res4w_out",          model.keras_model.get_layer("res4w_out").output),  # for resnet100
    ("rpn_bbox",           model.keras_model.get_layer("rpn_bbox").output),
    ("roi",                model.keras_model.get_layer("ROI").output),
])

# Input image (normalized)
_ = plt.imshow(modellib.unmold_image(activations["input_image"][0],config))

# Backbone feature map
display_images(np.transpose(activations["res2c_out"][0,:,:,:4], [2, 0, 1]), cols=4)



StarcoderdataJupyterScriptsDedupFiltered