Support for more models.

app.py

@@ -26,14 +26,20 @@ def inference(img, size, type):
     run_cmd("mkdir " + OUTPUT_DIR)
     img.save(INPUT_DIR + "1.jpg", "JPEG")
 
-    run_cmd("python inference_manga.py "+ os.path.join(INPUT_DIR, "1.jpg") + " " + os.path.join(OUTPUT_DIR, "1_out.jpg"))
+    if type is "Manga v1":
+        run_cmd("python inference_manga_v2.py "+ os.path.join(INPUT_DIR, "1.jpg") + " " + os.path.join(OUTPUT_DIR, "1_out.jpg"))
+    else:
+        run_cmd("python inference.py "+ os.path.join(INPUT_DIR, "1.jpg") + " " + os.path.join(OUTPUT_DIR, "1_out.jpg") + " " + type)
+
     img_out = Image.open(os.path.join(OUTPUT_DIR, "1_out.jpg"))
+
     if size is "x2":
         img_out = img_out.resize((s//2 for s in img_out.size))
+
     return [img_out]
 
 input_image = gr.Image(type="pil", label="Input")
-upscale_type = gr.Radio(["Manga", "Anime", "General"], label="Select the type of picture you want to upscale:", value="Manga")
+upscale_type = gr.Radio(["Manga v2", "Manga v1", "Anime", "General"], label="Select the type of picture you want to upscale:", value="Manga v2")
 upscale_size = gr.Radio(["x4", "x2"], label="Upscale by:", value="x4")
 output_image = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery").style(grid=[2], height="auto")
 

inference.py

@@ -0,0 +1,56 @@
+import sys
+import os.path
+import cv2
+import numpy as np
+import torch
+import architecture as arch
+
+def is_cuda():
+    if torch.cuda.is_available():
+        return True
+    else:
+        return False
+
+model_type = sys.argv[3]
+
+if model_type is "Manga v1":
+    model_path = "4x_eula_digimanga_bw_v1_860k.pth"
+elif model_type is "Anime":
+    model_path = "4x-AnimeSharp.pth"
+else:
+    model_path = "4x-UniScaleV2_Sharp.pth"
+
+img_path = sys.argv[1]
+output_dir = sys.argv[2]
+device = torch.device('cuda' if is_cuda() else 'cpu')
+
+model = arch.RRDB_Net(3, 3, 64, 23, gc=32, upscale=4, norm_type=None, act_type='leakyrelu', mode='CNA', res_scale=1, upsample_mode='upconv')
+
+if is_cuda():
+    model.load_state_dict(torch.load(model_path), strict=True)
+else:
+    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')), strict=True)
+
+model.eval()
+
+for k, v in model.named_parameters():
+    v.requires_grad = False
+model = model.to(device)
+
+base = os.path.splitext(os.path.basename(img_path))[0]
+
+# read image
+print(img_path);
+img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
+img = img * 1.0 / 255
+img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()
+img_LR = img.unsqueeze(0)
+img_LR = img_LR.to(device)
+
+print('Start upscaling...')
+with torch.no_grad():
+    output = model(img_LR).data.squeeze().float().cpu().clamp_(0, 1).numpy()
+    output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))
+    output = (output * 255.0).round()
+print('Finished upscaling, saving image.')
+cv2.imwrite(output_dir, output, [int(cv2.IMWRITE_JPEG_QUALITY), 90])