Refactor my_loader

alimoezzi · Aug 7, 2022 · b799154 · b799154
1 parent 72a19d7
commit b799154
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pokemon_dcgan_pytorch.ipynb b/pokemon_dcgan_pytorch.ipynb
@@ -499,7 +499,7 @@
    "outputs": [],
    "source": [
     "# New Dataset instance based on official artwork\n",
-    "def my_loader(path):\n",
+    "def transparent_white_bg_loader(path):\n",
     "    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)\n",
     "    alpha_channel = img[:, :, 3]\n",
     "    _, mask = cv2.threshold(alpha_channel, 254, 255, cv2.THRESH_BINARY)  # binarize mask\n",
@@ -508,7 +508,7 @@
     "    return new_img\n",
     "dataset = dset.ImageFolder(\n",
     "    root=config['dataset_path'],\n",
-    "    loader = my_loader,\n",
+    "    loader = transparent_white_bg_loader,\n",
     "    transform=transforms.Compose([\n",
     "        transforms.ToPILImage(),\n",
     "        transforms.Resize(config['image_size']),\n",