Fixes.

inference.py

@@ -4,6 +4,8 @@ import cv2
 import numpy as np
 import torch
 import architecture as arch
+from split_image import split_image
+from run_cmd import run_cmd
 
 def is_cuda():
     if torch.cuda.is_available():
@@ -37,19 +39,37 @@ 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
-img = cv2.imread(img_path, cv2.IMREAD_COLOR)
-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_PNG_COMPRESSION), 9])
+base_path = os.path.dirname(img_path)
+image_name = os.path.basename(img_path)
+
+# Split image
+run_cmd(f"split-image {img_path} 7 7 --output-dir {base_path} --quiet")
+
+for root, dirs, files in os.walk(base_path, topdown=True):
+    for x, name in enumerate(files):
+        file_path = os.path.join(root, name)
+
+        if file_path == img_path:
+            continue
+
+        # Read image
+        img = cv2.imread(file_path, cv2.IMREAD_COLOR)
+        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(f"Start upscaling tile {x}...")
+        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(f"Finished upscaling tile {x}, saving tile.")
+        cv2.imwrite(file_path, output)
+
+# Join all tiles
+run_cmd(f"cd {base_path} && split-image {image_name} 7 7 -r --quiet")
+
+# Open image and save as png with the ouput name
+img_out = cv2.imread(img_path);
+cv2.imwrite(output_dir, img_out, [int(cv2.IMWRITE_PNG_COMPRESSION), 5])

inference_manga_v2.py

@@ -6,8 +6,6 @@ import torch
 import architecture as arch
 from split_image import split_image
 from run_cmd import run_cmd
-import re
-from pathlib import Path
 
 def is_cuda():
     if torch.cuda.is_available():
@@ -35,17 +33,17 @@ for k, v in model.named_parameters():
     v.requires_grad = False
 model = model.to(device)
 
-base = os.path.dirname(img_path)
+base_path = os.path.dirname(img_path)
+image_name = os.path.basename(img_path)
 
 # Split image
-run_cmd(f"split-image {img_path} 7 7 --output-dir {base} --quiet")
+run_cmd(f"split-image {img_path} 7 7 --output-dir {base_path} --quiet")
 
-for root, dirs, files in os.walk(base, topdown=True):
+for root, dirs, files in os.walk(base_path, topdown=True):
     for x, name in enumerate(files):
         file_path = os.path.join(root, name)
 
         if file_path == img_path:
-            print(file_path)
             continue
 
         # Read image
@@ -64,9 +62,9 @@ for root, dirs, files in os.walk(base, topdown=True):
         cv2.imwrite(file_path, output)
 
 # Join all tiles
-run_cmd(f"cd {base} && split-image input.jpg 7 7 -r --quiet")
+run_cmd(f"cd {base_path} && split-image {image_name} 7 7 -r --quiet")
 
 # Open image and save as png with the ouput name
 img_out = cv2.imread(img_path);
-cv2.imwrite(output_dir, img_out, [int(cv2.IMWRITE_PNG_COMPRESSION), 9])
+cv2.imwrite(output_dir, img_out, [int(cv2.IMWRITE_PNG_COMPRESSION), 5])