Merge pull request #144 from amrzv/fix-gpen-notebook

Fix parser and futher arguments error
yangxy · Jun 28, 2022 · f04de33 · f04de33
2 parents 2661b0d + 5e99563
commit f04de33
Showing 1 changed file with 15 additions and 13 deletions.
diff --git a/GPEN.ipynb b/GPEN.ipynb
@@ -265,6 +265,7 @@
         "import __init_paths\n",
         "from face_enhancement import FaceEnhancement\n",
         "\n",
+        "\n",
         "def display(img1, img2):\n",
         "    fig = plt.figure(figsize=(25, 10))\n",
         "    ax1 = fig.add_subplot(1, 2, 1) \n",
@@ -276,6 +277,7 @@
         "    ax1.imshow(img1)\n",
         "    ax2.imshow(img2)\n",
         "\n",
+        "\n",
         "if __name__==\"__main__\":\n",
         "    parser = argparse.ArgumentParser()\n",
         "    parser.add_argument('--model', type=str, default='GPEN-BFR-512', help='GPEN model')\n",
@@ -296,30 +298,30 @@
         "    parser.add_argument('--indir', type=str, default='examples/imgs', help='input folder')\n",
         "    parser.add_argument('--outdir', type=str, default='results/outs-bfr', help='output folder')\n",
         "    parser.add_argument('--ext', type=str, default='.jpg', help='extension of output')\n",
-        "    args = parser.parse_args() \n",
+        "    args = parser.parse_args('')\n",
+        "\n",
+        "    # model = {'name':'GPEN-BFR-512', 'sr_model':'realesrnet', 'sr_scale': 2, 'size':512, 'channel_multiplier':2, 'narrow':1}\n",
+        "    # model = {'name':'GPEN-BFR-256', 'sr_model':'realesrnet', 'size':256, 'channel_multiplier':1, 'narrow':0.5}\n",
         "\n",
-        "    #model = {'name':'GPEN-BFR-512', 'sr_model':'realesrnet', 'sr_scale': 2, 'size':512, 'channel_multiplier':2, 'narrow':1}\n",
-        "    #model = {'name':'GPEN-BFR-256', 'sr_model':'realesrnet', 'size':256, 'channel_multiplier':1, 'narrow':0.5}\n",
-        "    \n",
         "    os.makedirs(args.outdir, exist_ok=True)\n",
         "\n",
         "    faceenhancer = FaceEnhancement(args, in_size=args.in_size, model=args.model, use_sr=args.use_sr, device='cuda' if args.use_cuda else 'cpu')\n",
         "\n",
-        "    file = os.path.join(indir, 'Solvay_conference_1927.png')\n",
+        "    file = os.path.join(args.indir, 'Solvay_conference_1927.png')\n",
         "    filename = os.path.basename(file)\n",
-        "    img = cv2.imread(file, cv2.IMREAD_COLOR) # BGR\n",
-        "    #im = cv2.resize(im, (0,0), fx=2, fy=2) #optional\n",
+        "    img = cv2.imread(file, cv2.IMREAD_COLOR)  # BGR\n",
+        "    # img = cv2.resize(img, (0,0), fx=2, fy=2) #optional\n",
         "\n",
         "    img_out, orig_faces, enhanced_faces = faceenhancer.process(img, aligned=False)\n",
-        "    \n",
+        "\n",
         "    img = cv2.resize(img, img_out.shape[:2][::-1])\n",
-        "    cv2.imwrite(os.path.join(outdir, '.'.join(filename.split('.')[:-1])+'_COMP.jpg'), np.hstack((img, img_out)))\n",
-        "    cv2.imwrite(os.path.join(outdir, '.'.join(filename.split('.')[:-1])+'_GPEN.jpg'), img_out)\n",
-        "        \n",
+        "    cv2.imwrite(os.path.join(args.outdir, '.'.join(filename.split('.')[:-1])+'_COMP.jpg'), np.hstack((img, img_out)))\n",
+        "    cv2.imwrite(os.path.join(args.outdir, '.'.join(filename.split('.')[:-1])+'_GPEN.jpg'), img_out)\n",
+        "\n",
         "    for m, (ef, of) in enumerate(zip(enhanced_faces, orig_faces)):\n",
         "        of = cv2.resize(of, ef.shape[:2])\n",
-        "        cv2.imwrite(os.path.join(outdir, '.'.join(filename.split('.')[:-1])+'_face%02d'%m+'.jpg'), np.hstack((of, ef)))\n",
-        "    \n",
+        "        cv2.imwrite(os.path.join(args.outdir, '.'.join(filename.split('.')[:-1])+'_face%02d'%m+'.jpg'), np.hstack((of, ef)))\n",
+        "\n",
         "    display(img, img_out)"
       ],
       "execution_count": null,