combine_and_split_datasets.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "moving for control method: canny\n",
      "  18243 train examples\n",
      "  184 test examples\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 18243/18243 [00:15<00:00, 1170.64it/s]\n",
      "100%|██████████| 184/184 [00:00<00:00, 1277.13it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "moving for control method: fullspec\n",
      "  18243 train examples\n",
      "  184 test examples\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 18243/18243 [00:06<00:00, 2689.05it/s]\n",
      "100%|██████████| 184/184 [00:00<00:00, 2852.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "moving for control method: sobel\n",
      "  18243 train examples\n",
      "  184 test examples\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 18243/18243 [00:10<00:00, 1669.15it/s]\n",
      "100%|██████████| 184/184 [00:00<00:00, 1607.21it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "moving for control method: sobeldenoise\n",
      "  18243 train examples\n",
      "  184 test examples\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 18243/18243 [00:08<00:00, 2075.91it/s]\n",
      "100%|██████████| 184/184 [00:00<00:00, 2078.06it/s]\n"
     ]
    }
   ],
   "source": [
    "import json, os\n",
    "import numpy as np\n",
    "import random\n",
    "import shutil\n",
    "from tqdm import tqdm\n",
    "\n",
    "# task: combine gtzan and slakh, and separate out some holdout examples for testing\n",
    "all_data_dirs = [\"gtzan-preprocessed/\",\"slakh-preprocessed/\"]\n",
    "control_methods = [\"canny\", \"fullspec\", \"sobel\", \"sobeldenoise\"]\n",
    "\n",
    "train_data_dir = \"train-data/\"\n",
    "val_data_dir = \"val-data/\"\n",
    "\n",
    "os.makedirs(train_data_dir, exist_ok=True)\n",
    "os.makedirs(val_data_dir, exist_ok=True)\n",
    "for control_method in control_methods:\n",
    "    os.makedirs(val_data_dir + \"source-\" + control_method + \"/\", exist_ok=True)\n",
    "    os.makedirs(train_data_dir + \"source-\" + control_method + \"/\", exist_ok=True)\n",
    "os.makedirs(val_data_dir + \"target/\", exist_ok=True)\n",
    "os.makedirs(train_data_dir + \"target/\", exist_ok=True)\n",
    "\n",
    "all_data = {}\n",
    "for control_method in control_methods:\n",
    "    all_data[control_method] = []\n",
    "    for rootdir in all_data_dirs:\n",
    "        with open(os.path.join(rootdir, 'prompt-'+control_method+'.json'), 'rt') as f:\n",
    "            for line in f:\n",
    "                all_data[control_method].append(json.loads(line))\n",
    "\n",
    "# get random train / test split\n",
    "random.seed(269364) \n",
    "n_total = len(all_data[control_methods[0]])\n",
    "n_test = int(n_total * 0.01)\n",
    "inds = list(np.arange(n_total))\n",
    "\n",
    "test_inds = random.sample(inds, n_test)\n",
    "test_inds = np.sort(test_inds)\n",
    "train_inds = np.array(inds)[[inds[i] not in test_inds for i in range(n_total)]]\n",
    "\n",
    "# move data over for each control method\n",
    "for c_num, control_method in enumerate(control_methods):\n",
    "\n",
    "    train_data = np.array(all_data[control_method])[list(train_inds)]\n",
    "    test_data = np.array(all_data[control_method])[list(test_inds)]\n",
    "\n",
    "    print(f\"moving for control method: {control_method}\")\n",
    "    print(f\"  {len(train_data)} train examples\")\n",
    "    print(f\"  {len(test_data)} test examples\")\n",
    "\n",
    "    # copy over train data\n",
    "    for i in tqdm(range(len(train_data))):\n",
    "        # only move target for first control method\n",
    "        if c_num == 0:\n",
    "            new_target_file = os.path.join(train_data_dir, \"target\", train_data[i][\"target\"].split('/')[-1])\n",
    "            shutil.copy(train_data[i][\"target\"], new_target_file)\n",
    "        # copy control method\n",
    "        new_file = os.path.join(train_data_dir, \"source-\"+control_method, train_data[i][\"source\"].split('/')[-1])\n",
    "        shutil.copy(train_data[i][\"source\"], new_file)\n",
    "\n",
    "    # write train prompt\n",
    "    with open(os.path.join(train_data_dir,\"prompt-\"+control_method+\".json\"), 'w') as outfile:\n",
    "        for i in range(len(train_data)):\n",
    "            packet = {\n",
    "                \"source\": str(os.path.join(train_data_dir, \"source-\"+control_method, train_data[i][\"source\"].split('/')[-1])),\n",
    "                \"target\":  str(os.path.join(train_data_dir, \"target\", train_data[i][\"target\"].split('/')[-1])),\n",
    "                \"prompt\": str(train_data[i][\"prompt\"])\n",
    "            }\n",
    "            json.dump(packet, outfile)\n",
    "            outfile.write('\\n')\n",
    "        outfile.close()\n",
    "\n",
    "    # copy over val data\n",
    "    for i in tqdm(range(len(test_data))):\n",
    "        # only move target for first control method\n",
    "        if c_num == 0:\n",
    "            new_target_file = os.path.join(val_data_dir, \"target\", test_data[i][\"target\"].split('/')[-1])\n",
    "            shutil.copy(test_data[i][\"target\"], new_target_file)\n",
    "        # copy control method\n",
    "        new_file = os.path.join(val_data_dir, \"source-\"+control_method, test_data[i][\"source\"].split('/')[-1])\n",
    "        shutil.copy(test_data[i][\"source\"], new_file)\n",
    "\n",
    "    # write val prompt\n",
    "    with open(os.path.join(val_data_dir,\"prompt-\"+control_method+\".json\"), 'w') as outfile:\n",
    "        for i in range(len(test_data)):\n",
    "            packet = {\n",
    "                \"source\": str(os.path.join(val_data_dir, \"source-\"+control_method, test_data[i][\"source\"].split('/')[-1])),\n",
    "                \"target\":  str(os.path.join(val_data_dir, \"target\", test_data[i][\"target\"].split('/')[-1])),\n",
    "                \"prompt\": str(test_data[i][\"prompt\"])\n",
    "            }\n",
    "            json.dump(packet, outfile)\n",
    "            outfile.write('\\n')\n",
    "        outfile.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "t2i-adapter",
   "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.8.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}