<jupyter_start><jupyter_text>Setup  Basic imports and constants<jupyter_code>import os
import sys
import random
import math
import re
import time
import numpy as np
import cv2
import matplotlib
import matplotlib.pyplot as plt
import pandas as pd
from keras.callbacks import EarlyStopping

# 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.config import Config
from mrcnn import utils
import mrcnn.model as modellib
from mrcnn import visualize
from mrcnn.model import log

from openimages2019 import setup as st
from openimages2019 import utils as u

from skimage.draw import rectangle

%matplotlib inline 

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

DATA_DIR = os.path.join(ROOT_DIR, "../data")

# Local path to trained weights file
COCO_MODEL_PATH = os.path.join(DATA_DIR, "mask_rcnn_coco.h5")

MASK_DIR = os.path.join(DATA_DIR, "segmentation")


#################################
USE_MASKS = True
#################################

#Make GPUs visible
!export HIP_VISIBLE_DEVICES=1,2,3

#Set which GPU devices' memory should be accessible to running GPUs
os.environ["CUDA_VISIBLE_DEVICES"]="1,2,3"


import mlflow
#add mlflow stuff

MLFLOW_SERVER = os.environ['mlflow_tracker_host']

os.environ['NO_PROXY'] = MLFLOW_SERVER
mlflow.tracking.set_tracking_uri('http://' + MLFLOW_SERVER + ':5005')
EXPERIMENT_NAME = 'kaggle_openimage_mask_rcnn_v1.0'
mlflow.set_experiment(EXPERIMENT_NAME)

# os.environ['AZURE_STORAGE_ACCESS_KEY'] = ''


def get_ax(rows=1, cols=1, size=8):
    """Return a Matplotlib Axes array to be used in
    all visualizations in the notebook. Provide a
    central point to control graph sizes.
    
    Change the default size attribute to control the size
    of rendered images
    """
    _, ax = plt.subplots(rows, cols, figsize=(size*cols, size*rows))
    return ax<jupyter_output><empty_output><jupyter_text>Define augmentation and train functions  Image Augmentation<jupyter_code># Example taken from https://imgaug.readthedocs.io/en/latest/source/examples_basics.html
# not actually used ... though

import numpy as np
import imgaug as ia
import imgaug.augmenters as iaa


ia.seed(1)

# Example batch of images.
# The array has shape (32, 64, 64, 3) and dtype uint8.
images = np.array(
    [ia.quokka(size=(64, 64)) for _ in range(32)],
    dtype=np.uint8
)

seq = iaa.Sequential([
    iaa.Fliplr(0.5), # horizontal flips
    iaa.Crop(percent=(0, 0.1)), # random crops
    # Small gaussian blur with random sigma between 0 and 0.5.
    # But we only blur about 50% of all images.
    iaa.Sometimes(0.5,
        iaa.GaussianBlur(sigma=(0, 0.5))
    ),
    # Strengthen or weaken the contrast in each image.
    iaa.ContrastNormalization((0.75, 1.5)),
    # Add gaussian noise.
    # For 50% of all images, we sample the noise once per pixel.
    # For the other 50% of all images, we sample the noise per pixel AND
    # channel. This can change the color (not only brightness) of the
    # pixels.
    iaa.AdditiveGaussianNoise(loc=0, scale=(0.0, 0.05*255), per_channel=0.5),
    # Make some images brighter and some darker.
    # In 20% of all cases, we sample the multiplier once per channel,
    # which can end up changing the color of the images.
    iaa.Multiply((0.8, 1.2), per_channel=0.2),
    # Apply affine transformations to each image.
    # Scale/zoom them, translate/move them, rotate them and shear them.
    iaa.Affine(
        scale={"x": (0.8, 1.2), "y": (0.8, 1.2)},
        translate_percent={"x": (-0.2, 0.2), "y": (-0.2, 0.2)},
        rotate=(-25, 25),
        shear=(-8, 8)
    )
], random_order=True) # apply augmenters in random order<jupyter_output><empty_output><jupyter_text>General training functions<jupyter_code>import imgaug

class TrainConfig(Config):
    
    NAME = "kaggle"
    GPU_COUNT = 3
    IMAGES_PER_GPU = 4
    IMAGE_MIN_DIM = 512
    IMAGE_MAX_DIM = 512
    

def log_params(pz):
    for k,v in pz.items():
        mlflow.log_param(k,v)

# Image Augmentation ... pulled from coco example
# Right/Left flip 50% of the time
augmentation = imgaug.augmenters.Fliplr(0.5)
        
        
def train(model, inf_config, train_data, val_data, params):
    with mlflow.start_run():
        
#         es = EarlyStopping(monitor='val_loss', mode='min', verbose=1, patience=20)
        
        model.train(train_data, val_data, 
                    learning_rate=params['learning_rate'], 
                    epochs=params['epochs'],
                    augmentation=augmentation,
                    layers=params['layers']
#                     layers=params['layers'],
#                     custom_callbacks=[es]
                   )

        log_params(params)
        
        #Inference to get mAP
        inf_model = modellib.MaskRCNN(mode="inference", config=inf_config, model_dir=MODEL_DIR)
        model_path = inf_model.find_last()
        inf_model.load_weights(model_path, by_name=True)
        mAP = u.eval_mAP(inf_model, val_data, inf_config, params['mAP_sample_size'])
        
        mlflow.log_metric('mAP', mAP)
        mlflow.log_param('Model Path', model_path)
        
        return model_path, mAP

    
    
def train_on_class_subset(class_set, params, init_with="coco"):
    
    class KaggleConfig(TrainConfig):
        NUM_CLASSES = len(class_set) + 1 # + 1 for background class

    class InferenceConfig(KaggleConfig):
        GPU_COUNT = 1

    train_config = KaggleConfig()
    inf_config = InferenceConfig()    

    model = modellib.MaskRCNN(mode="training", config=train_config, model_dir=MODEL_DIR)

    if init_with == "imagenet":
        model.load_weights(model.get_imagenet_weights(), by_name=True)
    elif init_with == "coco":
        model.load_weights(COCO_MODEL_PATH, by_name=True,
                           exclude=["mrcnn_class_logits", "mrcnn_bbox_fc", 
                                    "mrcnn_bbox", "mrcnn_mask"])
    else: #use a base model other than coco
        model.load_weights(init_with, by_name=True,
                          exclude=["mrcnn_class_logits", "mrcnn_bbox_fc", 
                                    "mrcnn_bbox", "mrcnn_mask"])
    
    
    if USE_MASKS:
        anns = st.load_annotations_by_image(class_set, use_masks=True)
        train_data = st.load_dataset(anns, DATA_DIR, class_set, is_train=True,mask_path=MASK_DIR)
        val_data = st.load_dataset(anns, DATA_DIR, class_set, is_train=False, mask_path=MASK_DIR)
    else:
        anns = st.load_annotations_by_image(class_set)
        train_data = st.load_dataset(anns, DATA_DIR, class_set, is_train=True)
        val_data = st.load_dataset(anns, DATA_DIR, class_set, is_train=False)
    
    
    return train(model, inf_config, train_data, val_data, params)<jupyter_output><empty_output><jupyter_text>Partition the classes according to frequency<jupyter_code># class_sets = st.partition_classes()

if USE_MASKS:
    all_classes = st.load_classes(path_to_csv=os.path.join(DATA_DIR,'seg_class_descriptions.csv'))
    anns = st.load_annotations_by_image(classes=all_classes, use_masks=True)
else:
    all_classes = st.load_classes()
    anns = st.load_annotations_by_image(classes=all_classes, use_masks=False)

    
cnts = anns['LabelName'].value_counts()

class_sets = []

n_partitions = 10

p_size = int(len(all_classes) / n_partitions)

for i in range(n_partitions):
    s = i*p_size
    idxs = cnts.iloc[s:(s+p_size)].index.values
    tmp_set = all_classes[all_classes['LabelName'].isin(idxs)]
    tmp_set = tmp_set.reset_index()
    tmp_set['LabelID'] = tmp_set.index + 1
    class_sets.append(tmp_set)<jupyter_output><empty_output><jupyter_text>Training  Train all models<jupyter_code>params = {
        'learning_rate' : 0.001,
        'epochs' : 50,
        'layers' : 'heads',
        'mAP_sample_size' : 250,
        'GPU_COUNT' : 4,
        'IMAGES_PER_GPU' : 4,
        'IMAGE_MIN_DIM' : 512,
        'IMAGE_MAX_DIM' : 512
         }

base_model = os.path.join(DATA_DIR, "models/omni_seg_base_0586.h5")

#TODO log the entire Config object with mlflow, not just params
for i,class_set in enumerate(class_sets):
    params['class_set_index'] = i
#     train_on_class_subset(class_set, params)
    train_on_class_subset(class_set, params, init_with=base_model)<jupyter_output><empty_output><jupyter_text>Train a single model<jupyter_code>#Note that changes to ALL_CAPS properties in params do not affect the model.  You will need to change the corresponding
# values above (e.g in TrainConfig) to have any effect.

params = {
        'learning_rate' : 0.001,
        'epochs' : 100,
        'layers' : 'heads',
        'mAP_sample_size' : 250,
        'GPU_COUNT' : 3,
        'IMAGES_PER_GPU' : 4,
        'IMAGE_MIN_DIM' : 512,
        'IMAGE_MAX_DIM' : 512
         }

# base_model = os.path.join(DATA_DIR, "models/omni_seg_base_0586.h5")

set_num = 9

params['class_set_index'] = set_num
train_on_class_subset(class_sets[set_num], params)<jupyter_output><empty_output><jupyter_text>Train a subset<jupyter_code>params = {
        'learning_rate' : 0.001,
        'epochs' : 50,
        'layers' : 'heads',
        'mAP_sample_size' : 250,
        'GPU_COUNT' : 3,
        'IMAGES_PER_GPU' : 4,
        'IMAGE_MIN_DIM' : 512,
        'IMAGE_MAX_DIM' : 512
         }

base_model = os.path.join(DATA_DIR, "models/omni_seg_base_0586.h5")

for i in range(2,10):
    params['class_set_index'] = i
#     train_on_class_subset(class_set, params)
    train_on_class_subset(class_sets[i], params, init_with=base_model)<jupyter_output><empty_output><jupyter_text>MISC  Determine number of images containing any object of a class set<jupyter_code>anns = st.load_annotations_by_image()

val_anns = anns[anns['RelativePath'].str.contains('validation',regex=False)]
train_anns = anns[anns['RelativePath'].str.contains('train',regex=False)]

def count_em(anns):
    rval = []
    
    for cs in class_sets:
        z = anns[anns['LabelName'].isin(cs['LabelName'].values)]['ImageID'].nunique()
        rval.append(z)
        
    return rval

num_class_set_images_train = count_em(train_anns)
num_class_set_images_val = count_em(val_anns)

display(num_class_set_images_train)
display(num_class_set_images_val)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1137974560636955, 'fraction_numerical': 0.019225167491989513, 'mean_word_length': 3.449244060475162, 'pattern_counts': {'":': 6, '<': 33, '<?xml version=': 0, '>': 33, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '315661', 'license': 'MIT', 'n_tokens_mistral': 3757, 'n_tokens_neox': 3568, 'n_words': 785, 'path': 'samples/openimages2019/ensemble_train.ipynb', 'repo_name': 'Tshort76/Mask_RCNN', 'size': '10299'}