feat: Support Grayscale Image (#16)

Final2x · Jan 1, 2024 · d4c9575 · d4c9575
1 parent c1bf9ef
commit d4c9575
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 3 deletions.
diff --git a/assets/gray.jpg b/assets/gray.jpg
diff --git a/scripts/gen_config.py b/scripts/gen_config.py
@@ -25,8 +25,9 @@ def gen_config() -> None:
         "gpuid": gpuid,
         "inputpath": [
             "./1/1/4/5/1/4/1/9/1/9/8/1/0.jpg",
+            str(projectPATH / "assets" / "gray.jpg"),
             str(projectPATH / "assets" / "herta.jpg"),
-            str(projectPATH / "assets" / "herta.jpg"),
+            str(projectPATH / "assets" / "final2x-10.png"),
             str(projectPATH / "assets" / "final2x-10.png"),
             str(projectPATH / "assets" / "final2x-20.png"),
             str(projectPATH / "assets" / "final2x-40.png"),

diff --git a/src/Final2x_core/src/SRqueue.py b/src/Final2x_core/src/SRqueue.py
@@ -45,7 +45,12 @@ def SR_queue() -> None:
                 # The file may not be read correctly.
                 # In unix-like system, the Filename Extension is not important.
                 img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_UNCHANGED)
-                if img.shape[2] == 4:
+
+                if len(img.shape) == 2:
+                    logger.warning("Grayscale image detected, Convert to RGB image.")
+                    img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
+
+                elif img.shape[2] == 4:
                     logger.warning("4 channels image detected.")
                     PrintProgressLog().Total += PrintProgressLog().sr_n
                     # Extract alpha channel

diff --git a/tests/util.py b/tests/util.py
@@ -33,8 +33,9 @@ def CONFIG() -> tuple[int, str, str, str]:
         "gpuid": gpuid,
         "inputpath": [
             "./1/1/4/5/1/4/1/9/1/9/8/1/0.jpg",
+            str(projectPATH / "assets" / "gray.jpg"),
             str(projectPATH / "assets" / "herta.jpg"),
-            str(projectPATH / "assets" / "herta.jpg"),
+            str(projectPATH / "assets" / "final2x-10.png"),
             str(projectPATH / "assets" / "final2x-10.png"),
             str(projectPATH / "assets" / "final2x-20.png"),
             str(projectPATH / "assets" / "final2x-40.png"),