feat: Add remove bg pipline

upcycling-kandidat · Apr 24, 2023 · 994c420 · 994c420
1 parent 953dd35
commit 994c420
Showing 1 changed file with 159 additions and 0 deletions.
diff --git a/remove-bg.ipynb b/remove-bg.ipynb
@@ -0,0 +1,159 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Install dependencies"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 42,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Requirement already satisfied: rembg[gpu] in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (2.0.34)\n",
+      "Requirement already satisfied: python-multipart>=0.0.5 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from rembg[gpu]) (0.0.6)\n",
+      "Requirement already satisfied: onnxruntime>=1.14.1 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from rembg[gpu]) (1.14.1)\n",
+      "Requirement already satisfied: imagehash>=4.3.1 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from rembg[gpu]) (4.3.1)\n",
+      "Requirement already satisfied: aiohttp>=3.8.1 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from rembg[gpu]) (3.8.4)\n",
+      "Requirement already satisfied: asyncer>=0.0.2 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from rembg[gpu]) (0.0.2)\n",
+      "Requirement already satisfied: numpy>=1.23.5 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from rembg[gpu]) (1.23.5)\n",
+      "Requirement already satisfied: pillow>=9.3.0 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from rembg[gpu]) (9.4.0)\n",
+      "Requirement already satisfied: scikit-image>=0.19.3 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from rembg[gpu]) (0.19.3)\n",
+      "Requirement already satisfied: pymatting>=1.1.8 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from rembg[gpu]) (1.1.8)\n",
+      "Requirement already satisfied: watchdog>=2.1.9 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from rembg[gpu]) (3.0.0)\n",
+      "Requirement already satisfied: tqdm>=4.64.1 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from rembg[gpu]) (4.65.0)\n",
+      "Requirement already satisfied: click>=8.1.3 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from rembg[gpu]) (8.1.3)\n",
+      "Requirement already satisfied: fastapi>=0.92.0 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from rembg[gpu]) (0.95.1)\n",
+      "Requirement already satisfied: pooch>=1.6.0 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from rembg[gpu]) (1.7.0)\n",
+      "Requirement already satisfied: filetype>=1.2.0 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from rembg[gpu]) (1.2.0)\n",
+      "Requirement already satisfied: opencv-python-headless>=4.6.0.66 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from rembg[gpu]) (4.7.0.72)\n",
+      "Requirement already satisfied: scipy>=1.9.3 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from rembg[gpu]) (1.10.1)\n",
+      "Requirement already satisfied: uvicorn>=0.20.0 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from rembg[gpu]) (0.21.1)\n",
+      "Requirement already satisfied: onnxruntime-gpu>=1.14.1 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from rembg[gpu]) (1.14.1)\n",
+      "Requirement already satisfied: charset-normalizer<4.0,>=2.0 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from aiohttp>=3.8.1->rembg[gpu]) (3.1.0)\n",
+      "Requirement already satisfied: multidict<7.0,>=4.5 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from aiohttp>=3.8.1->rembg[gpu]) (6.0.4)\n",
+      "Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from aiohttp>=3.8.1->rembg[gpu]) (4.0.2)\n",
+      "Requirement already satisfied: yarl<2.0,>=1.0 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from aiohttp>=3.8.1->rembg[gpu]) (1.9.1)\n",
+      "Requirement already satisfied: attrs>=17.3.0 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from aiohttp>=3.8.1->rembg[gpu]) (23.1.0)\n",
+      "Requirement already satisfied: aiosignal>=1.1.2 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from aiohttp>=3.8.1->rembg[gpu]) (1.3.1)\n",
+      "Requirement already satisfied: frozenlist>=1.1.1 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from aiohttp>=3.8.1->rembg[gpu]) (1.3.3)\n",
+      "Requirement already satisfied: anyio<4.0.0,>=3.4.0 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from asyncer>=0.0.2->rembg[gpu]) (3.6.2)\n",
+      "Requirement already satisfied: starlette<0.27.0,>=0.26.1 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from fastapi>=0.92.0->rembg[gpu]) (0.26.1)\n",
+      "Requirement already satisfied: pydantic!=1.7,!=1.7.1,!=1.7.2,!=1.7.3,!=1.8,!=1.8.1,<2.0.0,>=1.6.2 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from fastapi>=0.92.0->rembg[gpu]) (1.10.7)\n",
+      "Requirement already satisfied: PyWavelets in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from imagehash>=4.3.1->rembg[gpu]) (1.4.1)\n",
+      "Requirement already satisfied: protobuf in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from onnxruntime>=1.14.1->rembg[gpu]) (4.22.1)\n",
+      "Requirement already satisfied: coloredlogs in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from onnxruntime>=1.14.1->rembg[gpu]) (15.0.1)\n",
+      "Requirement already satisfied: sympy in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from onnxruntime>=1.14.1->rembg[gpu]) (1.11.1)\n",
+      "Requirement already satisfied: packaging in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from onnxruntime>=1.14.1->rembg[gpu]) (23.0)\n",
+      "Requirement already satisfied: flatbuffers in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from onnxruntime>=1.14.1->rembg[gpu]) (23.3.3)\n",
+      "Requirement already satisfied: platformdirs>=2.5.0 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from pooch>=1.6.0->rembg[gpu]) (3.2.0)\n",
+      "Requirement already satisfied: requests>=2.19.0 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from pooch>=1.6.0->rembg[gpu]) (2.28.2)\n",
+      "Requirement already satisfied: numba!=0.49.0 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from pymatting>=1.1.8->rembg[gpu]) (0.56.4)\n",
+      "Requirement already satisfied: tifffile>=2019.7.26 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from scikit-image>=0.19.3->rembg[gpu]) (2023.4.12)\n",
+      "Requirement already satisfied: networkx>=2.2 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from scikit-image>=0.19.3->rembg[gpu]) (3.0)\n",
+      "Requirement already satisfied: imageio>=2.4.1 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from scikit-image>=0.19.3->rembg[gpu]) (2.28.0)\n",
+      "Requirement already satisfied: h11>=0.8 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from uvicorn>=0.20.0->rembg[gpu]) (0.14.0)\n",
+      "Requirement already satisfied: sniffio>=1.1 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from anyio<4.0.0,>=3.4.0->asyncer>=0.0.2->rembg[gpu]) (1.3.0)\n",
+      "Requirement already satisfied: idna>=2.8 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from anyio<4.0.0,>=3.4.0->asyncer>=0.0.2->rembg[gpu]) (2.10)\n",
+      "Requirement already satisfied: llvmlite<0.40,>=0.39.0dev0 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from numba!=0.49.0->pymatting>=1.1.8->rembg[gpu]) (0.39.1)\n",
+      "Requirement already satisfied: setuptools in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from numba!=0.49.0->pymatting>=1.1.8->rembg[gpu]) (65.6.3)\n",
+      "Requirement already satisfied: typing-extensions>=4.2.0 in /home/yams/.local/lib/python3.9/site-packages (from pydantic!=1.7,!=1.7.1,!=1.7.2,!=1.7.3,!=1.8,!=1.8.1,<2.0.0,>=1.6.2->fastapi>=0.92.0->rembg[gpu]) (4.2.0)\n",
+      "Requirement already satisfied: certifi>=2017.4.17 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from requests>=2.19.0->pooch>=1.6.0->rembg[gpu]) (2022.12.7)\n",
+      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from requests>=2.19.0->pooch>=1.6.0->rembg[gpu]) (1.26.15)\n",
+      "Requirement already satisfied: humanfriendly>=9.1 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from coloredlogs->onnxruntime>=1.14.1->rembg[gpu]) (10.0)\n",
+      "Requirement already satisfied: mpmath>=0.19 in /home/yams/miniconda3/envs/chair-training/lib/python3.9/site-packages (from sympy->onnxruntime>=1.14.1->rembg[gpu]) (1.3.0)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Downloading data from 'https://github.com/danielgatis/rembg/releases/download/v0.0.0/isnet-general-use.onnx' to file '/home/yams/.u2net/isnet-general-use.onnx'.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Note: you may need to restart the kernel to use updated packages.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████| 179M/179M [00:00<00:00, 243GB/s]\n"
+     ]
+    },
+    {
+     "ename": "KeyboardInterrupt",
+     "evalue": "",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[42], line 21\u001b[0m\n\u001b[1;32m     19\u001b[0m \u001b[39mwith\u001b[39;00m \u001b[39mopen\u001b[39m(output_path, \u001b[39m'\u001b[39m\u001b[39mwb\u001b[39m\u001b[39m'\u001b[39m) \u001b[39mas\u001b[39;00m o:\n\u001b[1;32m     20\u001b[0m     \u001b[39minput\u001b[39m \u001b[39m=\u001b[39m i\u001b[39m.\u001b[39mread()\n\u001b[0;32m---> 21\u001b[0m     output \u001b[39m=\u001b[39m remove(\u001b[39minput\u001b[39;49m, session\u001b[39m=\u001b[39;49msession)\n\u001b[1;32m     22\u001b[0m     o\u001b[39m.\u001b[39mwrite(output)\n",
+      "File \u001b[0;32m~/miniconda3/envs/chair-training/lib/python3.9/site-packages/rembg/bg.py:144\u001b[0m, in \u001b[0;36mremove\u001b[0;34m(data, alpha_matting, alpha_matting_foreground_threshold, alpha_matting_background_threshold, alpha_matting_erode_size, session, only_mask, post_process_mask, bgcolor, *args, **kwargs)\u001b[0m\n\u001b[1;32m    141\u001b[0m \u001b[39mif\u001b[39;00m session \u001b[39mis\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m    142\u001b[0m     session \u001b[39m=\u001b[39m new_session(\u001b[39m\"\u001b[39m\u001b[39mu2net\u001b[39m\u001b[39m\"\u001b[39m, \u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m--> 144\u001b[0m masks \u001b[39m=\u001b[39m session\u001b[39m.\u001b[39;49mpredict(img, \u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m    145\u001b[0m cutouts \u001b[39m=\u001b[39m []\n\u001b[1;32m    147\u001b[0m \u001b[39mfor\u001b[39;00m mask \u001b[39min\u001b[39;00m masks:\n",
+      "File \u001b[0;32m~/miniconda3/envs/chair-training/lib/python3.9/site-packages/rembg/sessions/dis.py:16\u001b[0m, in \u001b[0;36mDisSession.predict\u001b[0;34m(self, img, *args, **kwargs)\u001b[0m\n\u001b[1;32m     13\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mpredict\u001b[39m(\u001b[39mself\u001b[39m, img: PILImage, \u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m List[PILImage]:\n\u001b[1;32m     14\u001b[0m     ort_outs \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39minner_session\u001b[39m.\u001b[39mrun(\n\u001b[1;32m     15\u001b[0m         \u001b[39mNone\u001b[39;00m,\n\u001b[0;32m---> 16\u001b[0m         \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mnormalize(img, (\u001b[39m0.485\u001b[39;49m, \u001b[39m0.456\u001b[39;49m, \u001b[39m0.406\u001b[39;49m), (\u001b[39m1.0\u001b[39;49m, \u001b[39m1.0\u001b[39;49m, \u001b[39m1.0\u001b[39;49m), (\u001b[39m1024\u001b[39;49m, \u001b[39m1024\u001b[39;49m)),\n\u001b[1;32m     17\u001b[0m     )\n\u001b[1;32m     19\u001b[0m     pred \u001b[39m=\u001b[39m ort_outs[\u001b[39m0\u001b[39m][:, \u001b[39m0\u001b[39m, :, :]\n\u001b[1;32m     21\u001b[0m     ma \u001b[39m=\u001b[39m np\u001b[39m.\u001b[39mmax(pred)\n",
+      "File \u001b[0;32m~/miniconda3/envs/chair-training/lib/python3.9/site-packages/rembg/sessions/base.py:28\u001b[0m, in \u001b[0;36mBaseSession.normalize\u001b[0;34m(self, img, mean, std, size, *args, **kwargs)\u001b[0m\n\u001b[1;32m     19\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mnormalize\u001b[39m(\n\u001b[1;32m     20\u001b[0m     \u001b[39mself\u001b[39m,\n\u001b[1;32m     21\u001b[0m     img: PILImage,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     26\u001b[0m     \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs\n\u001b[1;32m     27\u001b[0m ) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m Dict[\u001b[39mstr\u001b[39m, np\u001b[39m.\u001b[39mndarray]:\n\u001b[0;32m---> 28\u001b[0m     im \u001b[39m=\u001b[39m img\u001b[39m.\u001b[39;49mconvert(\u001b[39m\"\u001b[39;49m\u001b[39mRGB\u001b[39;49m\u001b[39m\"\u001b[39;49m)\u001b[39m.\u001b[39;49mresize(size, Image\u001b[39m.\u001b[39;49mLANCZOS)\n\u001b[1;32m     30\u001b[0m     im_ary \u001b[39m=\u001b[39m np\u001b[39m.\u001b[39marray(im)\n\u001b[1;32m     31\u001b[0m     im_ary \u001b[39m=\u001b[39m im_ary \u001b[39m/\u001b[39m np\u001b[39m.\u001b[39mmax(im_ary)\n",
+      "File \u001b[0;32m~/miniconda3/envs/chair-training/lib/python3.9/site-packages/PIL/Image.py:2192\u001b[0m, in \u001b[0;36mImage.resize\u001b[0;34m(self, size, resample, box, reducing_gap)\u001b[0m\n\u001b[1;32m   2184\u001b[0m             \u001b[39mself\u001b[39m \u001b[39m=\u001b[39m Image\u001b[39m.\u001b[39mreduce(\u001b[39mself\u001b[39m, factor, box\u001b[39m=\u001b[39mreduce_box)\n\u001b[1;32m   2185\u001b[0m         box \u001b[39m=\u001b[39m (\n\u001b[1;32m   2186\u001b[0m             (box[\u001b[39m0\u001b[39m] \u001b[39m-\u001b[39m reduce_box[\u001b[39m0\u001b[39m]) \u001b[39m/\u001b[39m factor_x,\n\u001b[1;32m   2187\u001b[0m             (box[\u001b[39m1\u001b[39m] \u001b[39m-\u001b[39m reduce_box[\u001b[39m1\u001b[39m]) \u001b[39m/\u001b[39m factor_y,\n\u001b[1;32m   2188\u001b[0m             (box[\u001b[39m2\u001b[39m] \u001b[39m-\u001b[39m reduce_box[\u001b[39m0\u001b[39m]) \u001b[39m/\u001b[39m factor_x,\n\u001b[1;32m   2189\u001b[0m             (box[\u001b[39m3\u001b[39m] \u001b[39m-\u001b[39m reduce_box[\u001b[39m1\u001b[39m]) \u001b[39m/\u001b[39m factor_y,\n\u001b[1;32m   2190\u001b[0m         )\n\u001b[0;32m-> 2192\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_new(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mim\u001b[39m.\u001b[39;49mresize(size, resample, box))\n",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
+     ]
+    }
+   ],
+   "source": [
+    "%pip install \"rembg[gpu]\"\n",
+    "\n",
+    "from pathlib import Path\n",
+    "from rembg import remove, new_session\n",
+    "\n",
+    "model_name = \"isnet-general-use\"\n",
+    "session = new_session(model_name=model_name)\n",
+    "\n",
+    "# Pass path to the folder containing images\n",
+    "path = Path('./datasets/komponenter').glob('*.JPG')\n",
+    "\n",
+    "for file in path: \n",
+    "    input_path = str(file)\n",
+    "    output_path = Path(f'{file.parent}/no-bg-{model_name}/{file.stem}-{model_name}-transparent.png')\n",
+    "\n",
+    "    # Create output folder if it doesn't exist\n",
+    "    output_path.parent.mkdir(exist_ok=True, parents=True)\n",
+    "\n",
+    "    with open(input_path, 'rb') as i:\n",
+    "        with open(output_path, 'wb') as o:\n",
+    "            input = i.read()\n",
+    "            output = remove(input, session=session)\n",
+    "            o.write(output)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "chair-training",
+   "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.16"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}