initial commit :)

.gitignore

@@ -2,7 +2,7 @@
 __pycache__/
 *.py[cod]
 *$py.class
-
+*.pth
 # C extensions
 *.so
 
@@ -127,3 +127,5 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+.idea

README.md

@@ -0,0 +1,15 @@
+# Yet Another Anime Upscaler
+
+YAAU is a model to upscale images, trained on Danbooru.
+A lot of thanks to Gwern for hosting such an incredible resource.
+![original](./example/Miku.png)
+![result](./example/Result.png)
+## Installation
+
+you will need poetry. `./init_env.sh` will install dependencies.
+if you want to train the model yourself, you will need danbooru, or at least a part of it.
+If you want to use a pretrained model, you can download it in the releases.
+
+```
+## Usage
+`python ./super_res.py <source>`

init_env.sh

@@ -0,0 +1,2 @@
+poetry update
+poetry run pip install torch==1.7.1+cu101 torchvision==0.8.2+cu101 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html

nb/classifier_rating.ipynb

@@ -0,0 +1,272 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "derived-ratio",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd\n",
+    "from fastai.vision.all import *\n",
+    "from pathlib import Path\n",
+    "import itertools\n",
+    "import progressbar"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "charming-adapter",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "DANBOORU_PATH = Path('/home/lleonard/Documents/datasets/danbooru/0/danbooru2020/')\n",
+    "EXTENSIONS = ['png', 'jpg']"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "passing-russian",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "i = 0\n",
+    "def file_exists(images_id):\n",
+    "    with progressbar.ProgressBar(max_value=len(images_id)) as bar:\n",
+    "        def single_image_exists(image_id):\n",
+    "            global i\n",
+    "            last_4_digit = str(image_id % 1000).zfill(4)\n",
+    "            path_to_glob = str(DANBOORU_PATH / '512px' / last_4_digit / str(image_id)) + '.'\n",
+    "            globbed = list(itertools.chain(*[glob.glob(path_to_glob + ext) for ext in EXTENSIONS]))\n",
+    "            bar.update(i)\n",
+    "            i = i + 1\n",
+    "            if len(globbed) == 0:\n",
+    "                return False\n",
+    "            return True\n",
+    "\n",
+    "        return images_id.apply(single_image_exists)\n",
+    "    \n",
+    "    \n",
+    "\n",
+    "ratings = pd.read_csv('/home/lleonard/Documents/datasets/danbooru/0/danbooru2020/ratings_tags_0000.csv')\n",
+    "print(ratings)\n",
+    "ratings = ratings[ratings['id'] % 1000 < 50]\n",
+    "ratings = ratings[file_exists(ratings['id'])]\n",
+    "ratings = ratings[['id', 'rating', 'tags']]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "warming-knitting",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ratings = pd.read_csv('clean_0000.csv')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "processed-special",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ratings.groupby('rating').count()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "attached-british",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sampled_ratings = ratings[ratings['rating'] != 's'].append(ratings[ratings['rating'] == 's'].sample(15000))\n",
+    "sampled_ratings.groupby('rating').count()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "studied-boston",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ratings.to_csv('clean_0000.csv', index=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "mediterranean-contract",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "def get_x(row):\n",
+    "    image_id = str(row['id'])\n",
+    "    mod_1000 = str(row['id'] % 1000)\n",
+    "    path_to_glob = str(DANBOORU_PATH / '512px' / mod_1000.zfill(4) / str(image_id)) + '.'\n",
+    "    return list(itertools.chain(*[glob.glob(path_to_glob + ext) for ext in EXTENSIONS]))[0]\n",
+    "    \n",
+    "\n",
+    "def get_y(row):\n",
+    "    return row['rating']\n",
+    "\n",
+    "dblock = DataBlock(blocks=(ImageBlock, CategoryBlock),\n",
+    "                   item_tfms=Resize((224,224)),\n",
+    "                   batch_tfms=[*aug_transforms(),Normalize()],\n",
+    "                   get_x = get_x, get_y = get_y)\n",
+    "dblock.summary(sampled_ratings)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "killing-cowboy",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "dataloaders = dblock.dataloaders(sampled_ratings, bs=128)\n",
+    "dataloaders.show_batch()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "floating-ranch",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "learner = cnn_learner(dataloaders, models.vgg16_bn, metrics = accuracy).to_fp16()\n",
+    "learner.model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "pleased-cooking",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "learner.lr_find()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "novel-campbell",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "learner.fit_one_cycle(100, 3e-3)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "finished-orange",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "learner.unfreeze()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "genetic-citizenship",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "learner.fit_one_cycle(100, slice(3e-5, 3e-3))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "owned-crack",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "learner.save('danbooru_vgg_classifier')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "laden-breed",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "torch.cuda.empty_cache()\n",
+    "dataloaders = dblock.dataloaders(ratings, bs=64)\n",
+    "model = model.eval().to('cuda')\n",
+    "data = iter(dataloaders.valid)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "white-ocean",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "import torchvision\n",
+    "\n",
+    "item = next(data)\n",
+    "grid = torchvision.utils.make_grid(item[0])\n",
+    "show_image(grid)\n",
+    "preds = model(item[0])\n",
+    "print(preds[preds[:, 0].argmax(), :])\n",
+    "show_image(item[0][preds[:, 0].argmax()])\n",
+    "del preds\n",
+    "del item"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "vital-milton",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import fastai\n",
+    "fastai.__version__, torch.__version__"
+   ]
+  }
+ ],
+ "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.8.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}