{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "collaborative-barrel",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "zsh:1: command not found: wget\n",
      "tar: Error opening archive: Failed to open 'MNIST.tar.gz'\n",
      "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz\n",
      "Failed to download (trying next):\n",
      "HTTP Error 503: Service Unavailable\n",
      "\n",
      "Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz\n",
      "Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz to ./MNIST/raw/train-images-idx3-ubyte.gz\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d7cf3b755a764d92980348c842e3a172",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/9912422 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting ./MNIST/raw/train-images-idx3-ubyte.gz to ./MNIST/raw\n",
      "\n",
      "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to ./MNIST/raw/train-labels-idx1-ubyte.gz\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7aebe80330474aada7b6dc1a1e2838d6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/28881 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting ./MNIST/raw/train-labels-idx1-ubyte.gz to ./MNIST/raw\n",
      "\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to ./MNIST/raw/t10k-images-idx3-ubyte.gz\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4414efcf0b7945d29f7256ab77747dc4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1648877 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting ./MNIST/raw/t10k-images-idx3-ubyte.gz to ./MNIST/raw\n",
      "\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to ./MNIST/raw/t10k-labels-idx1-ubyte.gz\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2d18107291a84953bee7b21b4962c402",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/4542 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting ./MNIST/raw/t10k-labels-idx1-ubyte.gz to ./MNIST/raw\n",
      "\n",
      "Processing...\n",
      "Done!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.9/site-packages/torchvision/datasets/mnist.py:502: UserWarning: The given NumPy array is not writeable, and PyTorch does not support non-writeable tensors. This means you can write to the underlying (supposedly non-writeable) NumPy array using the tensor. You may want to copy the array to protect its data or make it writeable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at  ../torch/csrc/utils/tensor_numpy.cpp:143.)\n",
      "  return torch.from_numpy(parsed.astype(m[2], copy=False)).view(*s)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([10000, 1, 28, 28])\n"
     ]
    }
   ],
   "source": [
    "# Load data (MNIST)\n",
    "\n",
    "import torch\n",
    "import torchvision\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "!wget www.di.ens.fr/~lelarge/MNIST.tar.gz\n",
    "!tar -zxvf MNIST.tar.gz\n",
    "\n",
    "batch_size_train = 100000\n",
    "batch_size_test = 20000\n",
    "log_interval = 10\n",
    "\n",
    "\n",
    "# num_samples = 128 # number of training samples\n",
    "num_samples = 128\n",
    "num_samples_test = 100 # number of test samples\n",
    "\n",
    "\n",
    "old_dim = 28 # MNIST original dimension\n",
    "\n",
    "\n",
    "random_seed = 1\n",
    "torch.backends.cudnn.enabled = False\n",
    "torch.manual_seed(random_seed)\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    torchvision.datasets.MNIST(root = './', train=True, download=True,\n",
    "                             transform=torchvision.transforms.Compose([\n",
    "                               torchvision.transforms.ToTensor(),\n",
    "                               torchvision.transforms.Normalize(\n",
    "                                 (0.1307,), (0.3081,))\n",
    "                             ])),\n",
    "    batch_size=batch_size_train, shuffle=True)\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(\n",
    "    torchvision.datasets.MNIST(root = './', train=False, download=True,\n",
    "                             transform=torchvision.transforms.Compose([\n",
    "                               torchvision.transforms.ToTensor(),\n",
    "                               torchvision.transforms.Normalize(\n",
    "                                 (0.1307,), (0.3081,))\n",
    "                             ])),\n",
    "    batch_size=batch_size_test, shuffle=True)\n",
    "\n",
    "examples = enumerate(test_loader)\n",
    "batch_idx, (example_data, example_targets) = next(examples)\n",
    "\n",
    "print(example_data.shape)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "married-howard",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([10000, 1, 28, 28])\n",
      "torch.Size([10000])\n"
     ]
    }
   ],
   "source": [
    "examples = enumerate(test_loader)\n",
    "batch_idx, (example_data, example_targets) = next(examples)\n",
    "\n",
    "print(example_data.shape)\n",
    "print(example_targets.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "caring-acrylic",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loaded\n"
     ]
    }
   ],
   "source": [
    "# Create training data and test data for binary classification task\n",
    "\n",
    "from tqdm import tqdm \n",
    "training_data = enumerate(train_loader)\n",
    "test_data = enumerate(test_loader)\n",
    "(batch_id1, (data_tr_old, target_tr_old)) = next(training_data)\n",
    "(batch_id2, (data_test_old, target_test_old)) = next(test_data)\n",
    "print(\"loaded\")\n",
    "\n",
    "\n",
    "# Classes\n",
    "class1 = [0]\n",
    "class2 = [1]\n",
    "\n",
    "# Pad the image on the right and on the bottom\n",
    "def filter_digits(data_old, target_old, new_dim, num, threshold):\n",
    "    data = torch.zeros([num, 1, new_dim, new_dim])\n",
    "    target = torch.zeros([num])\n",
    "\n",
    "    num_samples_per_class = int(num / 2)\n",
    "    sample1 = 0\n",
    "    sample2 = 0\n",
    "    attempt = 0\n",
    "    sample = 0\n",
    "\n",
    "    while (sample1 < num_samples_per_class) or (sample2 < num_samples_per_class):\n",
    "        if attempt == threshold:\n",
    "            print(\"FAILED: need more samples in batch\")\n",
    "            return (data, target)\n",
    "        # Balance classes\n",
    "        is_sample_for_class1 = (target_old[attempt] in class1) and (sample1 < num_samples_per_class)\n",
    "        is_sample_for_class2 = (target_old[attempt] in class2) and (sample2 < num_samples_per_class)\n",
    "        if is_sample_for_class1 or is_sample_for_class2: \n",
    "            avg = torch.mean(data_old[attempt][0])\n",
    "            target[sample] = target_old[attempt]\n",
    "            for i in range(old_dim):\n",
    "                for j in range(old_dim):\n",
    "                    data[sample][0][i][j] = data_old[attempt][0][i][j] - avg\n",
    "            if is_sample_for_class1:\n",
    "                sample1 += 1\n",
    "            if is_sample_for_class2:\n",
    "                sample2 += 1\n",
    "          # Augment sample counts\n",
    "            sample += 1\n",
    "        attempt += 1\n",
    "\n",
    "    target.apply_(lambda x: 1 if (x in class1) else -1)\n",
    "\n",
    "    data = data.float()\n",
    "    target = target.float()\n",
    "    return (data, target)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "still-pizza",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Training data and test data for 28 x 28 images \n",
    "(data_tr, target_tr) = filter_digits(data_tr_old, target_tr_old, old_dim, num_samples, batch_size_train)\n",
    "print(\"Created training data\")\n",
    "(data_test, target_test) = filter_digits(data_test_old, target_test_old, old_dim, num_samples_test, batch_size_test)\n",
    "print(\"Created test data\")\n",
    "\n",
    "\n",
    "torch.save(data_tr, \"training-test-data/training_data.txt\")\n",
    "\n",
    "\n",
    "torch.save(target_tr, \"training-test-data/training_targets.txt\")\n",
    "\n",
    "\n",
    "torch.save(data_test, \"training-test-data/test_data.txt\")\n",
    "\n",
    "\n",
    "torch.save(target_test, \"training-test-data/test_targets.txt\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "defined-valley",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created training data\n",
      "Created test data\n"
     ]
    }
   ],
   "source": [
    "# Training data and test data for 112 x 112 padded images \n",
    "\n",
    "new_dim = 112 \n",
    "(data_tr, target_tr) = filter_digits(data_tr_old, target_tr_old, new_dim, num_samples, batch_size_train)\n",
    "print(\"Created training data\")\n",
    "(data_test, target_test) = filter_digits(data_test_old, target_test_old, new_dim, num_samples_test, batch_size_test)\n",
    "print(\"Created test data\")\n",
    "\n",
    "\n",
    "torch.save(data_tr, \"training-test-data/training_data_augmented.txt\")\n",
    "\n",
    "\n",
    "torch.save(target_tr, \"training-test-data/training_targets_augmented.txt\")\n",
    "\n",
    "\n",
    "torch.save(data_test, \"training-test-data/test_data_augmented.txt\")\n",
    "\n",
    "\n",
    "torch.save(target_test, \"training-test-data/test_targets_augmented.txt\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "frequent-fortune",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
