0.4.2 (#35)

* Fix NCNN BGR/RGB issue + update ncnn dep

* Fix NCNN stuff

backend/nodes/ncnn_nodes.py

@@ -96,16 +96,6 @@ def __init__(self):
         self.sub = "NCNN"
 
     def upscale(self, img: np.ndarray, net: tuple, input_name: str, output_name: str):
-        dtype_max = 1
-        try:
-            dtype_max = np.iinfo(img.dtype).max
-        except:
-            logger.debug("img dtype is not an int")
-
-        img = (img.astype("float32") / dtype_max * 255).astype(
-            np.uint8
-        )  # don't ask lol
-
         # Try/except block to catch errors
         try:
             vkdev = ncnn.get_gpu_device(0)
@@ -140,15 +130,14 @@ def run(self, net_tuple: tuple, img: np.ndarray) -> np.ndarray:
         # TODO: This can prob just be a shared function tbh
         # Transparency hack (white/black background difference alpha)
         if in_nc == 3 and c == 4:
+            # NCNN expects RGB
+            img = cv2.cvtColor(img, cv2.COLOR_BGRA2RGBA)
             # Ignore single-color alpha
             unique = np.unique(img[:, :, 3])
             if len(unique) == 1:
                 logger.info("Single color alpha channel, ignoring.")
                 output = self.upscale(img[:, :, :3], net, input_name, output_name)
-                output = np.dstack(
-                    (output, np.full(output.shape[:-1], (unique[0] * 255)))
-                )
-                output = np.clip(output.astype(np.float32) / 255, 0, 1)
+                output = np.dstack((output, np.full(output.shape[:-1], (unique[0]))))
             else:
                 img1 = np.copy(img[:, :, :3])
                 img2 = np.copy(img[:, :, :3])
@@ -158,8 +147,6 @@ def run(self, net_tuple: tuple, img: np.ndarray) -> np.ndarray:
 
                 output1 = self.upscale(img1, net, input_name, output_name)
                 output2 = self.upscale(img2, net, input_name, output_name)
-                output1 = np.clip(output1.astype(np.float32) / 255, 0, 1)
-                output2 = np.clip(output2.astype(np.float32) / 255, 0, 1)
                 alpha = 1 - np.mean(output2 - output1, axis=2)
                 output = np.dstack((output1, alpha))
         else:
@@ -176,12 +163,17 @@ def run(self, net_tuple: tuple, img: np.ndarray) -> np.ndarray:
             elif img.shape[2] == 3 and in_nc == 4:
                 logger.debug("Expanding image channels")
                 img = np.dstack((img, np.full(img.shape[:-1], 1.0)))
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
             output = self.upscale(img, net, input_name, output_name)
 
             if gray:
                 output = np.average(output, axis=2)
 
-            output = output.astype(np.float32) / 255
+        if output.ndim > 2:
+            if output.shape[2] == 4:
+                output = cv2.cvtColor(output, cv2.COLOR_BGRA2RGBA)
+            elif output.shape[2] == 3:
+                output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
 
         output = np.clip(output, 0, 1)
 

backend/nodes/utils/ncnn_auto_split.py

@@ -4,7 +4,6 @@
 import numpy as np
 from ncnn_vulkan import ncnn
 from sanic.log import logger
-from torch import Tensor
 
 
 def fix_dtype_range(img):
@@ -14,22 +13,26 @@ def fix_dtype_range(img):
     except:
         logger.debug("img dtype is not an int")
 
-    img = (np.clip(img.astype("float32") / dtype_max, 0, 1) * 255).astype(np.uint8)
+    img = (
+        (np.clip(img.astype("float32") / dtype_max, 0, 1) * 255)
+        .round()
+        .astype(np.uint8)
+    )
     return img
 
 
 # NCNN version of the 'auto_split_upscale' function
 def ncnn_auto_split_process(
     lr_img: np.ndarray,
     net,
-    overlap: int = 32,
+    overlap: int = 16,
     max_depth: int = None,
     current_depth: int = 1,
     input_name: str = "data",
     output_name: str = "output",
     blob_vkallocator=None,
     staging_vkallocator=None,
-) -> Tuple[Tensor, int]:
+) -> Tuple[np.ndarray, int]:
     # Original code: https://github.com/JoeyBallentine/ESRGAN/blob/master/utils/dataops.py
 
     # Prevent splitting from causing an infinite out-of-vram loop
@@ -45,18 +48,19 @@ def ncnn_auto_split_process(
         ex.set_staging_vkallocator(staging_vkallocator)
         # ex.set_light_mode(True)
         try:
+            lr_img_fix = fix_dtype_range(lr_img.copy())
             mat_in = ncnn.Mat.from_pixels(
-                lr_img.copy(),
-                ncnn.Mat.PixelType.PIXEL_BGR,
-                lr_img.shape[1],
-                lr_img.shape[0],
+                lr_img_fix,
+                ncnn.Mat.PixelType.PIXEL_RGB,
+                lr_img_fix.shape[1],
+                lr_img_fix.shape[0],
             )
             mean_vals = []
             norm_vals = [1 / 255.0, 1 / 255.0, 1 / 255.0]
             mat_in.substract_mean_normalize(mean_vals, norm_vals)
             ex.input(input_name, mat_in)
             _, mat_out = ex.extract(output_name)
-            result = fix_dtype_range(np.array(mat_out).transpose(1, 2, 0))
+            result = np.array(mat_out).transpose(1, 2, 0).astype(np.float32)
             del ex, mat_in, mat_out
             # # Clear VRAM
             # blob_vkallocator.clear()

backend/nodes/utils/pytorch_auto_split.py

@@ -12,7 +12,7 @@ def auto_split_process(
     lr_img: Tensor,
     model: torch.nn.Module,
     scale: int = 4,
-    overlap: int = 32,
+    overlap: int = 16,
     max_depth: int = None,
     current_depth: int = 1,
 ) -> Tuple[Tensor, int]:

package.json

@@ -1,7 +1,7 @@
 {
   "name": "chainner",
   "productName": "chaiNNer",
-  "version": "0.4.1",
+  "version": "0.4.2",
   "description": "A flowchart based image processing GUI",
   "main": ".webpack/main",
   "scripts": {

src/helpers/dependencies.js

@@ -16,5 +16,5 @@ export default (isNvidiaAvailable) => [{
 }, {
   name: 'NCNN',
   packageName: 'ncnn-vulkan',
-  version: '2022.3.14',
+  version: '2022.3.23',
 }];