<jupyter_start><jupyter_text>Metric Learning 3D-CNN for clustering  for Google Colab<jupyter_code>from google.colab import drive
drive.mount('/content/drive')
!unzip -q /content/drive/MyDrive/data/mn10_64.zip
!pip install volumentations-3D tensorflow-addons tensorflow-determinism kaleido<jupyter_output><empty_output><jupyter_text>Import modules and set parameters<jupyter_code>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
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)
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
os.makedirs(EXPERIMENT_DIR, exist_ok=True)<jupyter_output><empty_output><jupyter_text>Define models<jupyter_code>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
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])
model.summary()<jupyter_output><empty_output><jupyter_text>Prepare data<jupyter_code>categories = ['bathtub', 'bed', 'chair', 'desk', 'dresser',
              'monitor', 'night_stand', 'sofa', 'table', 'toilet']
train_pattern = DATA_DIR +'/train/*.npy'

train_list = glob(train_pattern)
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]
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))
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)
ordered_stream = OrderedStream(train_list, train_ids, batch_size=BATCH_SIZE)
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)
aug_stream = AugmentedStream(train_list, aug, num_samples=NUM_SAMPLES, k=K)<jupyter_output><empty_output><jupyter_text>Train model<jupyter_code>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)
model.save(os.path.join(EXPERIMENT_DIR, 'saved_models', 'whole_model'))
feature_extractor.save(os.path.join(EXPERIMENT_DIR, 'saved_models', 'feature_extractor'))<jupyter_output><empty_output><jupyter_text>Copy experiment directory to Google Drive<jupyter_code>!cp -r $EXPERIMENT_DIR /content/drive/MyDrive/<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1340573561853046, 'fraction_numerical': 0.020872221279484134, 'mean_word_length': 4.0548885077186965, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '172791', 'license': 'Apache-2.0', 'n_tokens_mistral': 2230, 'n_tokens_neox': 2150, 'n_words': 382, 'path': 'MetricLearning3DWithVolumentation_3D.ipynb', 'repo_name': 'masa-ita/3d_clustering_experiments', 'size': '5893'}