feat: auto downsample if it is oversize

.gitignore

@@ -23,6 +23,8 @@ tmp/*
 *.zip
 *.mp4
 *.csv
+*.jpeg
+*.webp
 
 !__assets__/lr_inputs/*
 !__assets__/*

app.py

@@ -54,11 +54,12 @@ def inference(img_path, model_name):
 
 
         # In default, we will automatically use crop to match 4x size
-        super_resolved_img = super_resolve_img(generator, img_path, output_path=None, weight_dtype=weight_dtype, crop_for_4x=True)
+        super_resolved_img = super_resolve_img(generator, img_path, output_path=None, weight_dtype=weight_dtype, downsample_threshold=720, crop_for_4x=True)
         store_name = str(time.time()) + ".png"
         save_image(super_resolved_img, store_name)
         outputs = cv2.imread(store_name)
         outputs = cv2.cvtColor(outputs, cv2.COLOR_RGB2BGR)
+        os.remove(store_name)
         
         return outputs
     
@@ -78,6 +79,8 @@ if __name__ == '__main__':
 
     APISR aims at restoring and enhancing low-quality low-resolution anime images and video sources with various degradations from real-world scenarios.
 
+    ### Note: Due to memory restriction, all images whose short side is over 720 pixel will be downsampled to 720 pixel with the same aspect ratio.  E.g., 1920x1080 -> 1280x720
+            
     If APISR is helpful, please help star the GitHub Repo. Thanks! 
     """
 
@@ -112,6 +115,7 @@ if __name__ == '__main__':
                     ["__assets__/lr_inputs/image-00440.png"],
                     ["__assets__/lr_inputs/image-00164.jpg"],
                     ["__assets__/lr_inputs/img_eva.jpeg"],
+                    ["__assets__/lr_inputs/naruto.jpg"],
                 ],
                 [input_image],
             )

test_code/inference.py

@@ -19,19 +19,30 @@ from test_code.test_utils import load_grl, load_rrdb, load_cunet
 
 
 @torch.no_grad      # You must add these time, else it will have Out of Memory
-def super_resolve_img(generator, input_path, output_path=None, weight_dtype=torch.float32, crop_for_4x=True):
+def super_resolve_img(generator, input_path, output_path=None, weight_dtype=torch.float32, downsample_threshold=720, crop_for_4x=True):
     ''' Super Resolve a low resolution image
     Args:
         generator (torch):              the generator class that is already loaded
         input_path (str):               the path to the input lr images
         output_path (str):              the directory to store the generated images
         weight_dtype (bool):            the weight type (float32/float16)
+        downsample_threshold (int):     the threshold of height/width (short side) to downsample the input
         crop_for_4x (bool):             whether we crop the lr images to match 4x scale (needed for some situation)
     '''
     print("Processing image {}".format(input_path))
     
     # Read the image and do preprocess
     img_lr = cv2.imread(input_path)
+    h, w, c = img_lr.shape
+
+
+    # Downsample if needed
+    short_side = min(h, w)
+    if downsample_threshold != -1 and short_side > downsample_threshold:
+        resize_ratio = short_side / downsample_threshold
+        img_lr = cv2.resize(img_lr, (int(w/resize_ratio), int(h/resize_ratio)), interpolation = cv2.INTER_LINEAR)
+
+
     # Crop if needed
     if crop_for_4x:
         h, w, _ = img_lr.shape
@@ -139,12 +150,4 @@ if __name__ == "__main__":
         # In default, we will automatically use crop to match 4x size
         super_resolve_img(generator, input_dir, output_path, weight_dtype, crop_for_4x=True)
 
-    
-
-
-
-
-
-
-