ProteinNetModel.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████| 554/554 [00:05<00:00, 109.03it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(444, 2048, 3)\n",
      "(444, 2048, 3)\n",
      "(110, 2048, 3)\n",
      "(110, 2048, 3)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import open3d as o3d\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "# load proteins in the form of point cloud xyz format \n",
    "# and returns train and validation data according to validation_rate\n",
    "\n",
    "# load protein dataset\n",
    "input_directory = './protein_shape'\n",
    "\n",
    "def load_proteins (input_directory, validation_rate):\n",
    "\n",
    "    train_points = []\n",
    "    test_points = []\n",
    "    \n",
    "    file_list = os.listdir(input_directory)\n",
    "    \n",
    "    for file_name in tqdm(file_list):\n",
    "        \n",
    "        # read point cloud\n",
    "        full_path = os.path.join(input_directory, file_name)\n",
    "        pcd = o3d.io.read_point_cloud(full_path)\n",
    "        points = np.asarray(pcd.points)\n",
    "        \n",
    "        # add to list\n",
    "        train_points.append(points)\n",
    "    \n",
    "    \n",
    "    train_points = np.array(train_points)\n",
    "    validation_length = (int)(train_points.shape[0]*validation_rate)\n",
    "    train_length = train_points.shape[0]- validation_length\n",
    "    np.random.shuffle(train_points)\n",
    "    \n",
    "    train_points, test_points = train_points[:train_length,:,:], train_points[train_length:,:,:]\n",
    "    \n",
    "    return train_points, train_points, test_points, test_points\n",
    "\n",
    "\n",
    "train_points, train_labels, test_points, test_labels = load_proteins(input_directory, 0.2)\n",
    "\n",
    "print(train_points.shape)\n",
    "print(train_labels.shape)\n",
    "\n",
    "print(test_points.shape)\n",
    "print(test_labels.shape)\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"proteinnet\"\n",
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "input_1 (InputLayer)            [(None, 2048, 3)]    0                                            \n",
      "__________________________________________________________________________________________________\n",
      "conv1d (Conv1D)                 (None, 2048, 32)     128         input_1[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization (BatchNorma (None, 2048, 32)     128         conv1d[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "activation (Activation)         (None, 2048, 32)     0           batch_normalization[0][0]        \n",
      "__________________________________________________________________________________________________\n",
      "conv1d_1 (Conv1D)               (None, 2048, 64)     2112        activation[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_1 (BatchNor (None, 2048, 64)     256         conv1d_1[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "activation_1 (Activation)       (None, 2048, 64)     0           batch_normalization_1[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "conv1d_2 (Conv1D)               (None, 2048, 512)    33280       activation_1[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_2 (BatchNor (None, 2048, 512)    2048        conv1d_2[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "activation_2 (Activation)       (None, 2048, 512)    0           batch_normalization_2[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "global_max_pooling1d (GlobalMax (None, 512)          0           activation_2[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "dense (Dense)                   (None, 256)          131328      global_max_pooling1d[0][0]       \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_3 (BatchNor (None, 256)          1024        dense[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "activation_3 (Activation)       (None, 256)          0           batch_normalization_3[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "dense_1 (Dense)                 (None, 128)          32896       activation_3[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_4 (BatchNor (None, 128)          512         dense_1[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "activation_4 (Activation)       (None, 128)          0           batch_normalization_4[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "dense_2 (Dense)                 (None, 9)            1161        activation_4[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "reshape (Reshape)               (None, 3, 3)         0           dense_2[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "dot (Dot)                       (None, 2048, 3)      0           input_1[0][0]                    \n",
      "                                                                 reshape[0][0]                    \n",
      "==================================================================================================\n",
      "Total params: 204,873\n",
      "Trainable params: 202,889\n",
      "Non-trainable params: 1,984\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "# This code cell is adapted from https://keras.io/examples/vision/pointnet/\n",
    "import tensorflow as tf\n",
    "\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "\n",
    "\n",
    "#tf.random.set_random_seed(1234)\n",
    "\n",
    "def augment(points, labels):\n",
    "    # jitter points\n",
    "    points += tf.random.uniform(points.shape, -0.005, 0.005, dtype=tf.float64)\n",
    "    # shuffle points maybe it is not necessary\n",
    "    points = tf.random.shuffle(points)\n",
    "    labels = points\n",
    "    return points, labels\n",
    "\n",
    "\n",
    "def conv_bn(x, filters):\n",
    "    x = layers.Conv1D(filters, kernel_size=1, padding=\"valid\")(x)\n",
    "    x = layers.BatchNormalization(momentum=0.0)(x)\n",
    "    return layers.Activation(\"tanh\")(x)\n",
    "\n",
    "\n",
    "def dense_bn(x, filters):\n",
    "    x = layers.Dense(filters)(x)\n",
    "    x = layers.BatchNormalization(momentum=0.0)(x)\n",
    "    return layers.Activation(\"tanh\")(x)\n",
    "\n",
    "class OrthogonalRegularizer(keras.regularizers.Regularizer):\n",
    "    def __init__(self, num_features, l2reg=0.001):\n",
    "        self.num_features = num_features\n",
    "        self.l2reg = l2reg\n",
    "        self.eye = tf.eye(num_features)\n",
    "\n",
    "    def __call__(self, x):\n",
    "        x = tf.reshape(x, (-1, self.num_features, self.num_features))\n",
    "        xxt = tf.tensordot(x, x, axes=(2, 2))\n",
    "        xxt = tf.reshape(xxt, (-1, self.num_features, self.num_features))\n",
    "        return tf.reduce_sum(self.l2reg * tf.square(xxt - self.eye))\n",
    "    \n",
    "def tnet(inputs, num_features):\n",
    "    \n",
    "    # Initalise bias as the indentity matrix\n",
    "    bias = keras.initializers.Constant(np.eye(num_features).flatten())\n",
    "    reg = OrthogonalRegularizer(num_features)\n",
    "    \n",
    "    x = conv_bn(inputs, 32)\n",
    "    x = conv_bn(x, 64)\n",
    "    x = conv_bn(x, 512)\n",
    "    x = layers.GlobalMaxPooling1D()(x)\n",
    "    x = dense_bn(x, 256)\n",
    "    x = dense_bn(x, 128)\n",
    "    x = layers.Dense(\n",
    "        num_features * num_features,\n",
    "        kernel_initializer=\"zeros\",\n",
    "        bias_initializer=bias,\n",
    "        activity_regularizer=reg,\n",
    "    )(x)\n",
    "    feat_T = layers.Reshape((num_features, num_features))(x)\n",
    "    # Apply affine transformation to input features\n",
    "    return layers.Dot(axes=(2, 1))([inputs, feat_T])\n",
    "\n",
    "\n",
    "NUM_POINTS = 2048\n",
    "BATCH_SIZE = 16\n",
    "\n",
    "train_dataset = tf.data.Dataset.from_tensor_slices((train_points, train_labels))\n",
    "test_dataset = tf.data.Dataset.from_tensor_slices((test_points, test_labels))\n",
    "\n",
    "# data augmentation\n",
    "train_dataset = train_dataset.shuffle(len(train_points)).map(augment).batch(BATCH_SIZE)\n",
    "test_dataset = test_dataset.shuffle(len(test_points)).batch(BATCH_SIZE)\n",
    "\n",
    "\n",
    "inputs = keras.Input(shape=(NUM_POINTS, 3))\n",
    "x = tnet(inputs, 3)\n",
    "\n",
    "model = keras.Model(inputs=inputs, outputs=x, name=\"proteinnet\")\n",
    "model.summary()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run this cell to train a new model or skip it to use a model already trained\n",
    "model.compile(\n",
    "    loss=\"cosine_similarity\",\n",
    "    optimizer=keras.optimizers.Adam(learning_rate=0.001),\n",
    "    metrics=[\"accuracy\"],\n",
    ")\n",
    "\n",
    "model_checkpoint = tf.keras.callbacks.ModelCheckpoint('./proteinnet_weights.hdf5',\n",
    "                                   monitor='val_loss', save_weights_only=True, save_best_only=True)\n",
    "\n",
    "model.fit(train_dataset, epochs=50, validation_data=test_dataset, callbacks=[model_checkpoint])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4/4 [==============================] - 2s 198ms/step - loss: -0.9218 - accuracy: 0.9860\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[-0.9217734336853027, 0.9859996438026428]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Run this cell to test the perfomance of the model\n",
    "# Specify the model to use\n",
    "model.load_weights('./proteinnet_weights.hdf5')\n",
    "model.compile(\n",
    "    loss=\"cosine_similarity\",\n",
    "    optimizer=keras.optimizers.Adam(learning_rate=0.001),\n",
    "    metrics=[\"accuracy\"],\n",
    ")\n",
    "\n",
    "model.evaluate(test_points, test_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the trained model to extract features\n",
    "model_feature_512 = keras.Model(inputs=model.input, outputs=model.get_layer('global_max_pooling1d').output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████| 10/10 [00:00<00:00, 74.19it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(10, 2048, 3)\n",
      "(10, 512)\n"
     ]
    }
   ],
   "source": [
    "# read query proteins one by one and transform in features 512, 256, 128, save them in txt file\n",
    "input_directory = './queries_protein_shape'\n",
    "\n",
    "file_list = os.listdir(input_directory)\n",
    "\n",
    "file_list= sorted(file_list, key=lambda x: int(x[:-10]))\n",
    "\n",
    "list_points =[]\n",
    "\n",
    "for file_name in tqdm(file_list):\n",
    "    #print(file_name)\n",
    "    # read point cloud\n",
    "    full_path = os.path.join(input_directory, file_name)\n",
    "    pcd = o3d.io.read_point_cloud(full_path)\n",
    "    points = np.asarray(pcd.points)\n",
    "        \n",
    "    # add to list\n",
    "    list_points.append(points)\n",
    "    \n",
    "\n",
    "proteins = np.array(list_points)\n",
    "print(proteins.shape)\n",
    "\n",
    "\n",
    "#extract features\n",
    "queries_features_512 = model_feature_512.predict(proteins)\n",
    "print(queries_features_512.shape)\n",
    "\n",
    "#save queries features to txt files\n",
    "np.savetxt('./queries_features_512.txt', queries_features_512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 554/554 [00:05<00:00, 98.74it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(554, 2048, 3)\n",
      "(554, 512)\n"
     ]
    }
   ],
   "source": [
    "# read dataset proteins one by one and transform in features 512, save them in one txt file\n",
    "input_directory = './protein_shape'\n",
    "\n",
    "file_list = os.listdir(input_directory)\n",
    "\n",
    "file_list= sorted(file_list, key=lambda x: int(x[:-10]))\n",
    "\n",
    "list_points =[]\n",
    "\n",
    "for file_name in tqdm(file_list):\n",
    "    #print(file_name)\n",
    "    # read point cloud\n",
    "    full_path = os.path.join(input_directory, file_name)\n",
    "    pcd = o3d.io.read_point_cloud(full_path)\n",
    "    points = np.asarray(pcd.points)\n",
    "        \n",
    "    # add to list\n",
    "    list_points.append(points)\n",
    "    \n",
    "\n",
    "proteins = np.array(list_points)\n",
    "print(proteins.shape)\n",
    "\n",
    "\n",
    "#extract features\n",
    "features_512 = model_feature_512.predict(proteins)\n",
    "print(features_512.shape)\n",
    "\n",
    "#save dataset features to txt files\n",
    "np.savetxt('./features_512.txt', features_512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#calculate dissimilarity matrices using euclidean distances\n",
    "# for the ten queries to the 554 shapes\n",
    "# and save to binary file\n",
    "def dissimimarity_function(queries_features, features):\n",
    "    \n",
    "    dissim = []\n",
    "    \n",
    "    for i in range(queries_features.shape[0]):\n",
    "        \n",
    "        dist = (features - queries_features[i])**2\n",
    "        dist = np.sum(dist, axis=1)\n",
    "        dist = np.sqrt(dist)\n",
    "        #dist = dist.reshape(1,554)\n",
    "        # append distance of query i to all 554 proteins\n",
    "        dissim.append(dist)\n",
    "        \n",
    "    dissim = np.array(dissim)\n",
    "    return dissim\n",
    "\n",
    "\n",
    "dissim_512 = dissimimarity_function(queries_features_512, features_512)\n",
    "#print(dissim_512)\n",
    "\n",
    "\n",
    "# save binary format\n",
    "dissim_512.tofile('./dissim_512.bin')\n",
    "\n",
    "# save txt format\n",
    "# np.savetxt('dissim_512.txt', dissim_512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[0.7599654  0.39820912 0.43557584 ... 0.3895003  0.6007906  0.37562284]\n",
      " [0.65125334 0.44127426 0.43598667 ... 0.4557057  0.49995616 0.47706172]\n",
      " [0.65658325 0.29750183 0.41647333 ... 0.24150035 0.52480197 0.44314343]\n",
      " ...\n",
      " [0.7216959  0.48784533 0.43673837 ... 0.42315108 0.57508284 0.44344944]\n",
      " [0.6728263  0.19655116 0.3315216  ... 0.23727167 0.47231674 0.33636096]\n",
      " [0.7421873  0.45231634 0.38263497 ... 0.46039557 0.55437803 0.49546513]]\n"
     ]
    }
   ],
   "source": [
    "# load dissimilarity matrices\n",
    "dissim_512 = np.fromfile('./dissim_512.bin', dtype=np.float32) \n",
    "dissim_512 = dissim_512.reshape(10, 554)\n",
    "print(dissim_512)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "TensorFlow-CPU",
   "language": "python",
   "name": "tensorflow"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}