Merge pull request #499 from dmMaze/main

remove accumulate_color of TextBlock
zyddnys · Nov 12, 2023 · 20d6a87 · 20d6a87
2 parents d1740d2 + be9dc8b
commit 20d6a87
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 34 deletions.
diff --git a/manga_translator/manga_translator.py b/manga_translator/manga_translator.py
@@ -492,7 +492,6 @@ async def _run_textline_merge(self, ctx: Context):
 
         for region in text_regions:
             if ctx.font_color_fg or ctx.font_color_bg:
-                region.accumulate_color = False
                 if ctx.font_color_bg:
                     region.adjust_bg_color = False
 

diff --git a/manga_translator/ocr/model_32px.py b/manga_translator/ocr/model_32px.py
@@ -125,8 +125,7 @@ async def _infer(self, image: np.ndarray, textlines: List[Quadrilateral], args:
                     cur_region.bg_b = bb
                 else:
                     cur_region.text.append(txt)
-                    cur_region.fg_colors += np.array([fr, fg, fb])
-                    cur_region.bg_colors += np.array([br, bg, bb])
+                    cur_region.update_font_colors(np.array([fr, fg, fb]), np.array([br, bg, bb]))
 
                 out_regions.append(cur_region)
 

diff --git a/manga_translator/ocr/model_48px.py b/manga_translator/ocr/model_48px.py
@@ -148,8 +148,7 @@ async def _infer(self, image: np.ndarray, textlines: List[Quadrilateral], verbos
                     cur_region.bg_b = bb
                 else:
                     cur_region.text.append(txt)
-                    cur_region.fg_colors += np.array([fr, fg, fb])
-                    cur_region.bg_colors += np.array([br, bg, bb])
+                    cur_region.update_font_colors(np.array([fr, fg, fb]), np.array([br, bg, bb]))
 
                 out_regions.append(cur_region)
 

diff --git a/manga_translator/ocr/model_48px_ctc.py b/manga_translator/ocr/model_48px_ctc.py
@@ -144,8 +144,7 @@ async def _infer(self, image: np.ndarray, textlines: List[Quadrilateral], args:
                     cur_region.bg_b = bb
                 else:
                     cur_region.text.append(txt)
-                    cur_region.fg_colors += np.array([fr, fg, fb])
-                    cur_region.bg_colors += np.array([br, bg, bb])
+                    cur_region.update_font_colors(np.array([fr, fg, fb]), np.array([br, bg, bb]))
 
                 out_regions.append(cur_region)
 

diff --git a/manga_translator/textline_merge/__init__.py b/manga_translator/textline_merge/__init__.py
@@ -203,6 +203,6 @@ async def dispatch(textlines: List[Quadrilateral], width: int, height: int, verb
         lines = [txtln.pts for txtln in txtlns]
         texts = [txtln.text for txtln in txtlns]
         region = TextBlock(lines, texts, font_size=font_size, angle=angle, prob=np.exp(total_logprobs),
-                           fg_color=fg_color, bg_color=bg_color, accumulate_color=False)
+                           fg_color=fg_color, bg_color=bg_color)
         text_regions.append(region)
     return text_regions
diff --git a/manga_translator/utils/textblock.py b/manga_translator/utils/textblock.py
@@ -58,7 +58,6 @@ def __init__(self, lines: List,
                  alignment: str = 'auto',
                  rich_text: str = "",
                  _bounding_rect: List = None,
-                 accumulate_color = True,
                  default_stroke_width = 0.2,
                  font_weight = 50,
                  target_lang: str = "",
@@ -81,7 +80,6 @@ def __init__(self, lines: List,
 
         self.translation = translation
 
-        # note they're accumulative rgb values of textlines
         self.fg_colors = fg_color
         self.bg_colors = bg_color
 
@@ -99,7 +97,6 @@ def __init__(self, lines: List,
         self._bounding_rect = _bounding_rect
         self.default_stroke_width = default_stroke_width
         self.font_weight = font_weight
-        self.accumulate_color = accumulate_color
         self.adjust_bg_color = True
 
         self.opacity = opacity
@@ -295,31 +292,24 @@ def is_bulleted_list(self):
                     bullet_type_idx = i
         return bullet_type_idx >= 0
 
-    def set_font_colors(self, fg_colors, bg_colors, accumulate=True):
-        self.accumulate_color = accumulate
-        num_lines = len(self.lines) if accumulate and len(self.lines) > 0 else 1
-        # set font color
-        fg_colors = np.array(fg_colors) * num_lines
-        self.fg_colors = fg_colors
-        # set stroke color  
-        bg_colors = np.array(bg_colors) * num_lines
-        self.bg_colors = bg_colors
+    def set_font_colors(self, fg_colors, bg_colors):
+        self.fg_colors = np.array(fg_colors)
+        self.bg_colors = np.array(bg_colors)
+
+    def update_font_colors(self, fg_colors: np.ndarray, bg_colors: np.ndarray):
+        nlines = len(self)
+        if nlines > 0:
+            self.fg_colors += fg_colors / nlines
+            self.bg_colors += bg_colors / nlines
 
     def get_font_colors(self, bgr=False):
-        num_lines = len(self.lines)
-
-        frgb = np.array(self.fg_colors)
-        brgb = np.array(self.bg_colors)
-        if self.accumulate_color:
-            if num_lines > 0:
-                frgb = (frgb / num_lines).astype(np.int32)
-                brgb = (brgb / num_lines).astype(np.int32)
-                if bgr:
-                    frgb, brgb = frgb[::-1], brgb[::-1]
-                else:
-                    frgb, brgb = frgb, brgb
-            else:
-                return [0, 0, 0], [0, 0, 0]
+
+        frgb = np.array(self.fg_colors).astype(np.int32)
+        brgb = np.array(self.bg_colors).astype(np.int32)
+
+        if bgr:
+            frgb = frgb[::-1]
+            brgb = brgb[::-1]
 
         if self.adjust_bg_color:
             fg_avg = np.mean(frgb)