Added Generic XSeg model usage

ETChush · May 17, 2021 · e230891 · e230891
1 parent 97ec56d
commit e230891
Showing 1 changed file with 23 additions and 10 deletions.
diff --git a/DFL_Colab.ipynb b/DFL_Colab.ipynb
@@ -14,8 +14,8 @@
         "WTuyUxgdLA13",
         "avAcSL_uvtq_"
       ],
-      "machine_shape": "hm",
-      "toc_visible": true
+      "toc_visible": true,
+      "machine_shape": "hm"
     },
     "kernelspec": {
       "name": "python3",
@@ -480,7 +480,7 @@
         "* Resize faceset to your model resolution. Since Colab doesn't have a powerful CPU, resizing samples during training increases iteration time. Faceset resize reduces iteration time by about 2x times. Don't forget to keep save original faceset on your PC.\n",
         "* Pack or unpack facesets with DFL packing tool.\n",
         "* Apply or remove trained XSeg model to the extracted faces.\n",
-        "\n"
+        "* Recommended for use, Generic XSeg model for auto segmentation.\n"
       ]
     },
     {
@@ -619,7 +619,7 @@
         "#@title Pack/Unpack aligned faceset\n",
         "\n",
         "Folder = \"data_src\" #@param [\"data_src\", \"data_dst\"]\n",
-        "Mode = \"pack\" #@param [\"pack\", \"unpack\"]\n",
+        "Mode = \"unpack\" #@param [\"pack\", \"unpack\"]\n",
         "\n",
         "cmd = \"/content/DeepFaceLab/main.py util --input-dir /content/workspace/\" + \\\n",
         "      f\"{Folder}/aligned --{Mode}-faceset\"\n",
@@ -639,12 +639,25 @@
         "#@title Apply or remove XSeg mask to the faces\n",
         "Mode = \"Apply mask\" #@param [\"Apply mask\", \"Remove mask\"]\n",
         "Data = \"data_src\" #@param [\"data_src\", \"data_dst\"]\n",
+        "GenericXSeg = True #@param {type:\"boolean\"}\n",
+        "\n",
+        "from pathlib import Path\n",
+        "mode_arg = 'apply' if Mode == \"Apply mask\" else 'remove'\n",
+        "\n",
+        "if GenericXSeg and not Path('/content/GenericXSeg').exists():\n",
+        "  print('Downloading Generic XSeg model ... ')\n",
+        "  xseg_link = 'https://github.com/chervonij/DFL-Colab/releases/download/GenericXSeg/GenericXSeg.zip'\n",
+        "  !mkdir /content/GenericXSeg\n",
+        "  !wget -q --no-check-certificate -r $xseg_link -O /content/GenericXSeg.zip\n",
+        "  !unzip -q /content/GenericXSeg.zip -d /content/GenericXSeg/\n",
+        "  !rm /content/GenericXSeg.zip\n",
+        "\n",
+        "main_path = '/content/DeepFaceLab/main.py'\n",
+        "data_path = f'/content/workspace/{Data}/aligned'\n",
+        "model_path = '/content/workspace/model' if not GenericXSeg else '/content/GenericXSeg'\n",
         "\n",
-        "main_path = \"/content/DeepFaceLab/main.py \"\n",
-        "data_path = \"/content/workspace/\"+Data+\"/aligned \"\n",
-        "mode_arg = \"apply \" if Mode == \"Apply mask\" else \"remove \"\n",
-        "cmd = main_path+\"xseg \"+mode_arg+\"--input-dir \"+data_path\n",
-        "cmd += \"--model-dir /content/workspace/model\" if mode_arg == \"apply \" else \"\"\n",
+        "cmd = f'{main_path} xseg {mode_arg} --input-dir {data_path} '\n",
+        "cmd += f'--model-dir {model_path}' if mode_arg == 'apply' else ''\n",
         "\n",
         "!python $cmd"
       ],
@@ -781,4 +794,4 @@
       "outputs": []
     }
   ]
-}
+}