# ---
# 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
# ---

# + [markdown] id="YUuPSqnCmrwh"
# # Metric Learning 3D-CNN for clustering

# + [markdown] id="Z0MEYWGlmxx1"
# ## for Google Colab

# + id="9a117f4c"
from google.colab import drive
drive.mount('/content/drive')

# + id="c9c2e582"
# !unzip -q /content/drive/MyDrive/data/mn10_64.zip

# + id="ShzHa4E6unZu"
# !pip install volumentations-3D tensorflow-addons tensorflow-determinism kaleido

# + [markdown] id="IW_eeZRxm1Fp"
# ## Import modules and set parameters

# + id="b2afbb09"
import os
from glob import glob
import re
import numpy as np
import pandas as pd
from volumentations import *

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import tensorflow_addons as tfa


# + id="A47FOhxTnkLx"
def set_seed(seed=200):
    tf.random.set_seed(seed)

    # optional
    # for numpy.random
    np.random.seed(seed)
    # for built-in random
    random.seed(seed)
    # for hash seed
    os.environ["PYTHONHASHSEED"] = str(seed)
    os.environ['TF_DETERMINISTIC_OPS'] = '1'
    
set_seed(123)

# + id="7eef2634"
EXPERIMENT_DIR = 'MetricLearning_64'

NUM_SAMPLES = 10
K = 4
BATCH_SIZE = NUM_SAMPLES * K
EPOCHS = 10

DATA_DIR = '/content/mn10/64'
IMAGE_SIZE = 64
NUM_CHANNELS = 1

# + id="d16dc394"
os.makedirs(EXPERIMENT_DIR, exist_ok=True)


# + [markdown] id="bf270b35"
# ## Define models

# + id="Lv8WgZDDZ3NP"
class L2Constraint(keras.layers.Layer):
    def __init__(self, alpha=16):
        super(L2Constraint, self).__init__()
        self.alpha = alpha

    def call(self, x):
        x = self.alpha * tf.nn.l2_normalize(x)
        return x


# + id="a8d48419"
input = layers.Input(shape=(IMAGE_SIZE, IMAGE_SIZE, IMAGE_SIZE, NUM_CHANNELS))
x = layers.Conv3D(filters=64, kernel_size=3, activation='relu')(input)
x = layers.MaxPool3D(pool_size=2)(x)
x = layers.BatchNormalization()(x)
x = layers.Conv3D(filters=64, kernel_size=3, activation='relu')(x)
x = layers.MaxPool3D(pool_size=2)(x)
x = layers.BatchNormalization()(x)
x = layers.Conv3D(filters=128, kernel_size=3, activation='relu')(x)
x = layers.MaxPool3D(pool_size=2)(x)
x = layers.BatchNormalization()(x)
x = layers.Conv3D(filters=256, kernel_size=3, activation='relu')(x)
x = layers.MaxPool3D(pool_size=2)(x)
x = layers.BatchNormalization()(x)
x = layers.GlobalAveragePooling3D()(x)
features = layers.Dense(128)(x)
output = tf.nn.l2_normalize(x, axis=-1)

model = keras.Model(inputs=[input], outputs=[output])
feature_extractor = keras.Model(inputs=[input], outputs=[features])

# + id="bfe4e7b6"
model.summary()

# + [markdown] id="81e74dd6"
# ## Prepare data

# + id="c2e4f85f"
categories = ['bathtub', 'bed', 'chair', 'desk', 'dresser',
              'monitor', 'night_stand', 'sofa', 'table', 'toilet']

# + id="e9f92c4b"
train_pattern = DATA_DIR +'/train/*.npy'

train_list = glob(train_pattern)

# + id="22d76cda"
cat_re = re.compile(r'.+/(.+?)_[0-9]+\.npy')
train_labels = [cat_re.match(item)[1] for item in train_list]
train_ids = [categories.index(cat) for cat in train_labels]


# + id="7898f33c"
def get_augmentation(patch_size):
    return Compose([
        Rotate((-15, 15), (-15, 15), (-15, 15), interpolation=1, p=0.5),
        ElasticTransform((0, 0.25), interpolation=1, p=0.1),
        Flip(0, p=0.5),
        Flip(1, p=0.5),
        Flip(2, p=0.5),
        RandomRotate90((0, 1), p=0.5),
    ], p=1.0)

aug = get_augmentation((IMAGE_SIZE, IMAGE_SIZE, IMAGE_SIZE))


# + id="iEjzCg498DL5"
class OrderedStream(keras.utils.Sequence):
    def __init__(self, file_list, ids_list, batch_size=32):
        self.file_list = file_list
        self.ids_list = ids_list
        self.batch_size = batch_size
        self.num_files = len(file_list)

    def __len__(self):
        return int(np.ceil(self.num_files / float(self.batch_size)))

    def __getitem__(self, idx):
        x = []
        from_idx = idx * self.batch_size
        to_idx = min((idx + 1) * self.batch_size, self.num_files)
        for file_path in self.file_list[from_idx : to_idx]:
            x.append(self._read_npy_file(file_path))
        y = self.ids_list[from_idx : to_idx]
        return np.array(x), np.array(y)

    def _read_npy_file(self, path):
        data = np.load(path)
        data = np.expand_dims(data, axis=-1)
        return data.astype(np.float32)


# + id="jJKI68MAGpBC"
ordered_stream = OrderedStream(train_list, train_ids, batch_size=BATCH_SIZE)


# + id="YIaEDsBhwbeh"
class AugmentedStream(keras.utils.Sequence):
    def __init__(self, file_list, aug, num_samples=4, k=4):
        self.file_list = file_list
        self.num_files = len(file_list)
        self.file_indices = np.random.permutation(np.arange(self.num_files))
        self.aug = aug
        self.num_samples = num_samples
        self.k = k
        self.num_batchs = self.num_files // self.num_samples

    def __len__(self):
        return int(np.ceil(self.num_files / float(self.num_samples)))

    def __getitem__(self, idx):
        x, y = [], []
        from_idx = idx * self.num_samples
        to_idx = min((idx + 1) * self.num_samples, self.num_files)
        sample_indices = self.file_indices[from_idx : to_idx]
        for idx in sample_indices:
            data = {'image': self._read_npy_file(self.file_list[idx])}
            for i in range(self.k):
                aug_data = self.aug(**data)
                x.append(aug_data['image'])
                y.append(idx)
        return np.array(x), np.array(y)

    def _read_npy_file(self, path):
        data = np.load(path)
        data = np.expand_dims(data, axis=-1)
        return data.astype(np.float32)


# + id="LsooR5n802ek"
aug_stream = AugmentedStream(train_list, aug, num_samples=NUM_SAMPLES, k=K)

# + [markdown] id="cda1de99"
# ## Train model

# + id="757969d3"
loss_fn = tfa.losses.TripletSemiHardLoss(margin=0.1)
optimizer = tf.keras.optimizers.Adam(learning_rate=1.0e-4)

model.compile(loss=loss_fn, optimizer=optimizer)

model.fit(aug_stream, epochs=EPOCHS)

# + id="FYeMdaS8pen0"
model.save(os.path.join(EXPERIMENT_DIR, 'saved_models', 'whole_model'))
feature_extractor.save(os.path.join(EXPERIMENT_DIR, 'saved_models', 'feature_extractor'))

# + [markdown] id="KkMo941A_5QI"
# ## Copy experiment directory to Google Drive

# + id="_1Z05B5SWVcg"
# !cp -r $EXPERIMENT_DIR /content/drive/MyDrive/

# + id="aTV89t8gz5D1"


StarcoderdataJupyterScriptsDedupFiltered