better image file handling

ESRGANer.py

@@ -103,6 +103,8 @@ class ESRGANer():
                 tile_idx = y * tiles_x + x + 1
                 input_tile = self.img[:, :, input_start_y_pad:input_end_y_pad, input_start_x_pad:input_end_x_pad]
 
+                print("Image processing started...")
+
                 # upscale tile
                 try:
                     with torch.no_grad():

inference.py

@@ -13,7 +13,7 @@ def is_cuda():
     else:
         return False
 
-model_type = sys.argv[3]
+model_type = sys.argv[2]
 
 if model_type == "Anime":
     model_path = "models/4x-AnimeSharp.pth"
@@ -22,8 +22,7 @@ if model_type == "Photo":
 else:
     model_path = "models/4x-UniScaleV2_Sharp.pth"
 
-img_path = sys.argv[1]
-output_dir = sys.argv[2]
+OUTPUT_PATH = sys.argv[1]
 device = torch.device('cuda' if is_cuda() else 'cpu')
 
 if model_type != "Photo":
@@ -45,7 +44,7 @@ for k, v in model.named_parameters():
 model = model.to(device)
 
 # Read image
-img = cv2.imread(img_path, cv2.IMREAD_COLOR)
+img = cv2.imread(OUTPUT_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)
@@ -57,4 +56,4 @@ output = upsampler.enhance(img_LR)
 output = output.squeeze().float().cpu().clamp_(0, 1).numpy()
 output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))
 output = (output * 255.0).round()
-cv2.imwrite(output_dir, output, [int(cv2.IMWRITE_PNG_COMPRESSION), 5])
+cv2.imwrite(OUTPUT_PATH, output, [int(cv2.IMWRITE_PNG_COMPRESSION), 5])

inference_manga_v2.py

@@ -12,8 +12,7 @@ def is_cuda():
         return False
 
 model_path = 'models/4x_eula_digimanga_bw_v2_nc1_307k.pth'
-img_path = sys.argv[1]
-output_dir = sys.argv[2]
+OUTPUT_PATH = sys.argv[1]
 device = torch.device('cuda' if is_cuda() else 'cpu')
 
 model = arch.RRDB_Net(1, 1, 64, 23, gc=32, upscale=4, norm_type=None, act_type='leakyrelu', mode='CNA', res_scale=1, upsample_mode='upconv')
@@ -32,7 +31,7 @@ for k, v in model.named_parameters():
 model = model.to(device)
 
 # Read image
-img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
+img = cv2.imread(OUTPUT_PATH, cv2.IMREAD_GRAYSCALE)
 img = img * 1.0 / 255
 img = torch.from_numpy(img[np.newaxis, :, :]).float()
 img_LR = img.unsqueeze(0)
@@ -44,4 +43,4 @@ output = upsampler.enhance(img_LR)
 output = output.squeeze(dim=0).float().cpu().clamp_(0, 1).numpy()
 output = np.transpose(output, (1, 2, 0))
 output = (output * 255.0).round()
-cv2.imwrite(output_dir, output, [int(cv2.IMWRITE_PNG_COMPRESSION), 5])
+cv2.imwrite(OUTPUT_PATH, output, [int(cv2.IMWRITE_PNG_COMPRESSION), 5])

process_image.py

@@ -1,9 +1,9 @@
 import os
 import gradio as gr
 from run_cmd import run_cmd
-from random import randint
 from PIL import Image
 import tempfile
+import uuid
 
 temp_path = tempfile.gettempdir()
 
@@ -11,29 +11,20 @@ def inference(img, size, type):
     if not img:
         raise Exception("No image!")
 
-    _id = randint(1, 10000)
-    INPUT_DIR = os.path.join(temp_path, f"input_image{str(_id)}")
-    OUTPUT_DIR = os.path.join(temp_path, f"output_image{str(_id)}")
-    img_in_path = os.path.join(INPUT_DIR, "input.jpg")
-    img_out_path = os.path.join(OUTPUT_DIR, f"output_{size}.png")
-    run_cmd(f"mkdir {INPUT_DIR}")
-    run_cmd(f"mkdir {OUTPUT_DIR}")
+    OUTPUT_PATH = os.path.join(temp_path, f"{str(uuid.uuid4())[0:12]}_{size}.png")
     
-    img.save(img_in_path, "PNG")
+    img.save(OUTPUT_PATH)
 
     if type == "Manga":
-        run_cmd(f"python inference_manga_v2.py {img_in_path} {img_out_path}")
+        run_cmd(f"python inference_manga_v2.py {OUTPUT_PATH}")
     else:
-        run_cmd(f"python inference.py {img_in_path} {img_out_path} {type}")
+        run_cmd(f"python inference.py {OUTPUT_PATH} {type}")
 
-    img_out = Image.open(img_out_path)
+    img_out = Image.open(OUTPUT_PATH)
 
     if size == "x2":
         img_out = img_out.resize((img_out.width // 2, img_out.height // 2), resample=Image.BICUBIC)
 
-    # Remove input and output image
-    run_cmd(f"rm -rf {INPUT_DIR}")
-
     img_out.thumbnail((600, 600), Image.ANTIALIAS)
 
-    return img_out, gr.File.update(value=img_out_path, visible=True)
+    return img_out, gr.File.update(value=OUTPUT_PATH, visible=True)