✨ Add torch to accelerate process

README.md

@@ -1,4 +1,4 @@
-Convert image/gif/video to ascii art.
+Convert image/gif/video to ascii art. If you inputs have plenty of frames, you can specify `--fast` to use torch accelerating the peocess.
 
 ## Screen Shot
 
@@ -22,20 +22,24 @@ img2art --help
 result:
 
 ```
- Usage: img2art [OPTIONS] SOURCE                                                                                                 
-                                                                                                                                 
-╭─ Arguments ───────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
-│ *    source      TEXT  Path to image [default: None] [required]                                                               │
-╰───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
-╭─ Options ─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
-│ --with-color    --no-with-color                                  Whether use color [default: no-with-color]                   │
-│ --scale                            FLOAT                         Scale applied to image [default: 1.0]                        │
-│ --threshold                        INTEGER                       Threshold applied to image, default to OSTU [default: -1]    │
-│ --save-raw                         TEXT                          Whether to save the raw data [default: None]                 │
-│ --bg-color                         <INTEGER INTEGER INTEGER>...  Backgound color, (-1, -1, -1) for none [default: -1, -1, -1] │
-│ --help                                                           Show this message and exit.                                  │
-╰───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
-
+ Usage: img2art [OPTIONS] SOURCE
+
+╭─ Arguments ────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
+│ *    source      TEXT  Path to image [default: None] [required]                                                                │
+╰────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
+╭─ Options ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
+│ --with-color    --no-with-color                                  Whether use color [default: no-with-color]                    │
+│ --scale                            FLOAT                         Scale applied to image [default: 1.0]                         │
+│ --threshold                        INTEGER                       Threshold applied to image, default to OSTU [default: -1]     │
+│ --save-raw                         TEXT                          Whether to save the raw data [default: None]                  │
+│ --bg-color                         <INTEGER INTEGER INTEGER>...  Backgound color, (-1, -1, -1) for none [default: -1, -1, -1]  │
+│ --fast          --no-fast                                        Whether use torch to accelerate when you inputs have plenty   │
+│                                                                  of frames.                                                    │
+│                                                                  [default: no-fast]                                            │
+│ --chunk-size                       INTEGER                       Chunk size of Videos or Gifs when using torch.                │
+│                                                                  [default: 1024]                                               │
+│ --help                                                           Show this message and exit.                                   │
+╰────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
 ```
 
 ```

img2art/cli.py

@@ -35,7 +35,8 @@ def main(
             callback=_generate_check_func(lambda x: not osp.exists(x)),
         ),
     ],
-    with_color: Annotated[bool, typer.Option(help="Whether use color")] = False,
+    with_color: Annotated[bool, typer.Option(
+        help="Whether use color")] = False,
     scale: Annotated[
         float,
         typer.Option(
@@ -61,8 +62,19 @@ def main(
             callback=_generate_check_func(_is_rgb),
         ),
     ] = (-1, -1, -1),
+    fast: Annotated[
+        bool,
+        typer.Option(
+            help="Whether use torch to accelerate when you inputs have plenty of frames."
+        ),
+    ] = False,
+    chunk_size: Annotated[
+        int, typer.Option(
+            help="Chunk size of Videos or Gifs when using torch.")
+    ] = 1024,
 ):
-    convert(source, with_color, scale, threshold, save_raw, bg_color)
+    convert(source, with_color, scale, threshold,
+            save_raw, bg_color, fast, chunk_size)
 
 
 def launch():

img2art/utils.py

@@ -1,4 +1,5 @@
 import os.path as osp
+from math import ceil
 from time import sleep
 from typing import List, Optional, Tuple
 
@@ -21,7 +22,6 @@ def apply_threshold(data: np.ndarray, threshold: int) -> np.ndarray:
         threshold = cv2.threshold(
             data, threshold, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU
         )[0]
-        print(threshold)
     return data > threshold
 
 
@@ -39,6 +39,55 @@ def convert2braille(bool_map: np.ndarray, y: int, x: int, h: int, w: int) -> str
     return BRAILLE[idx]
 
 
+def _fast_convert2braille(bool_maps: List[np.ndarray], chunk_size: int):
+    try:
+        import torch
+    except:
+        raise ModuleNotFoundError("Can not find torch")
+    with torch.no_grad():
+        ch = len(bool_maps)
+        inp = torch.tensor(np.array([bool_maps])).type(torch.float32)
+
+        if ch <= chunk_size:
+            loader = [inp]
+        else:
+            loader = torch.chunk(inp, ceil(ch / chunk_size))
+
+        kernel = (
+            torch.tensor([2**i for i in range(7, -1, -1)])
+            .reshape((1, 1, 4, 2))
+            .type(torch.float32)
+        )
+        if torch.cuda.is_available():
+            inp = inp.cuda()
+            kernel = kernel.cuda()
+        res = []
+        for x in loader:
+            ker = kernel.repeat(x.shape[1], 1, 1, 1)
+            res.append(
+                torch.nn.functional.conv2d(
+                    x, ker, padding=0, stride=(4, 2), groups=x.shape[1]
+                )
+            )
+        res = torch.concat(res, 1)
+        if torch.cuda.is_available():
+            res = res.cpu()
+        index_maps = res.numpy().astype(np.uint8)
+        index_maps = np.split(index_maps, ch, 1)
+        return [i[0, 0] for i in index_maps]
+
+
+def _get_braille(index_map: np.ndarray) -> List:
+    h, w = index_map.shape
+    raw = []
+    for i in range(h):
+        str_list = []
+        for j in range(w):
+            str_list.append(BRAILLE[int(index_map[i, j])])
+        raw.append("".join(str_list))
+    return raw
+
+
 def render(
     color_mat: np.ndarray,
     y: int,
@@ -106,7 +155,8 @@ def _apple_color(raw, bg_color, bgr_data):
     h = oh // 4
     w = ow // 2
     color_raw = [[] for _ in range(len(raw))]
-    colors = [cv2.resize(i, (w, h), interpolation=cv2.INTER_LINEAR) for i in bgr_data]
+    colors = [cv2.resize(i, (w, h), interpolation=cv2.INTER_LINEAR)
+              for i in bgr_data]
     colors = [cv2.cvtColor(i, cv2.COLOR_BGR2RGB) for i in colors]
     for idx, r in enumerate(raw):
         for i, line in enumerate(r):
@@ -117,17 +167,15 @@ def _apple_color(raw, bg_color, bgr_data):
     return color_raw
 
 
-def _fuse(fg, bg):
-    return fg + bg
-
-
 def convert(
     source: str,
     with_color: bool = False,
     scale: float = 1.0,
     threshold: int = -1,
     save_raw: Optional[str] = None,
     bg_color: Optional[Tuple[int, int, int]] = None,
+    fast: bool = False,
+    chunk_size: int = False,
 ):
     ext = osp.splitext(source)[1][1:]
     try:
@@ -146,8 +194,12 @@ def convert(
 
     gray_data = [cv2.cvtColor(i, cv2.COLOR_BGR2GRAY) for i in bgr_data]
     bool_maps = [apply_threshold(i, threshold) for i in gray_data]
-
-    raw = _apply_convert(bool_maps)
+
+    if fast:
+        index_maps = _fast_convert2braille(bool_maps, chunk_size)
+        raw = [_get_braille(x) for x in index_maps]
+    else:
+        raw = _apply_convert(bool_maps)
     color_raw = [[] for _ in range(len(gray_data))]
 
     if not with_color:

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "img2art"
-version = "0.2.0"
+version = "0.3.0"
 description = ""
 authors = ["Asthestarsfalll <1186454801@qq.com>"]
 license = "MIT"

test.py

@@ -0,0 +1,31 @@
+import numpy as np
+import torch
+
+from img2art.utils import fast_convert2braille
+
+BRAILLE = (
+    "⠀⢀⡀⣀⠠⢠⡠⣠⠄⢄⡄⣄⠤⢤⡤⣤"
+    "⠐⢐⡐⣐⠰⢰⡰⣰⠔⢔⡔⣔⠴⢴⡴⣴⠂⢂⡂⣂⠢⢢⡢⣢⠆⢆⡆⣆⠦⢦⡦⣦⠒⢒⡒⣒⠲⢲⡲⣲⠖⢖⡖⣖⠶⢶⡶⣶"
+    "⠈⢈⡈⣈⠨⢨⡨⣨⠌⢌⡌⣌⠬⢬⡬⣬⠘⢘⡘⣘⠸⢸⡸⣸⠜⢜⡜⣜⠼⢼⡼⣼⠊⢊⡊⣊⠪⢪⡪⣪⠎⢎⡎⣎⠮⢮⡮⣮⠚⢚⡚⣚⠺⢺⡺⣺⠞⢞⡞⣞⠾⢾⡾⣾"
+    "⠁⢁⡁⣁⠡⢡⡡⣡⠅⢅⡅⣅⠥⢥⡥⣥⠑⢑⡑⣑⠱⢱⡱⣱⠕⢕⡕⣕⠵⢵⡵⣵⠃⢃⡃⣃⠣⢣⡣⣣⠇⢇⡇⣇⠧⢧⡧⣧⠓⢓⡓⣓⠳⢳⡳⣳⠗⢗⡗⣗⠷⢷⡷⣷"
+    "⠉⢉⡉⣉⠩⢩⡩⣩⠍⢍⡍⣍⠭⢭⡭⣭⠙⢙⡙⣙⠹⢹⡹⣹⠝⢝⡝⣝⠽⢽⡽⣽⠋⢋⡋⣋⠫⢫⡫⣫⠏⢏⡏⣏⠯⢯⡯⣯⠛⢛⡛⣛⠻⢻⡻⣻⠟⢟⡟⣟⠿⢿⡿⣿"
+)
+
+# bool_map = np.array([[1, 1], [1, 1], [1, 1], [1, 1]]).astype(np.bool_)
+bool_map = np.array([[0, 0], [0, 0], [0, 0], [0, 1]]).astype(np.bool_)
+print(bool_map)
+index_map = fast_convert2braille(bool_map)
+print(index_map)
+
+
+a = torch.tensor(bool_map).type(torch.float32).reshape(1, 1, 4, 2)
+kernel = (
+    np.array([2**i for i in range(7, -1, -1)])
+    .reshape((4, 2))
+    .astype(np.float32)
+    .reshape(1, 1, 4, 2)
+)
+w = torch.tensor(kernel).type(torch.float32)
+b = torch.nn.functional.conv2d(a, w, padding=0)
+print(b)
+print(BRAILLE[int(b.item())])