Update visualization

app.py

@@ -14,6 +14,7 @@ from tempfile import TemporaryDirectory
 import xlsxwriter
 import matplotlib.pyplot as plt
 import matplotlib.patches as patches
+from matplotlib.patches import Patch
 
 from paddleocr import PaddleOCR
 # import pytesseract
@@ -218,8 +219,9 @@ def visualize_structure(pil_img, structure_result):
     height = image.shape[0]
     # print(width, height)
 
-    fig, ax = plt.subplots(1)
-    ax.imshow(pil_img, interpolation='lanczos')
+    plt.imshow(pil_img, interpolation='lanczos')
+    plt.gcf().set_size_inches(20, 20)
+    ax = plt.gca()
 
     for i, result in enumerate(structure_result):
         class_id = int(result[5])
@@ -250,47 +252,90 @@ def visualize_structure(pil_img, structure_result):
             # Hatch
             rect = patches.Rectangle(bbox[:2], bbox[2]-bbox[0], bbox[3]-bbox[1],
                                      linewidth=1, alpha=0.4,
-                                     edgecolor=color,facecolor='none',
+                                     edgecolor=color, facecolor='none',
                                      linestyle='--',hatch=hatch)
             ax.add_patch(rect)
             # Edge
             rect = patches.Rectangle(bbox[:2], bbox[2]-bbox[0], bbox[3]-bbox[1],
                                      linewidth=linewidth,
-                                     edgecolor=color,facecolor='none',
+                                     edgecolor=color, facecolor='none',
                                      linestyle="--")
             ax.add_patch(rect)
 
+    plt.xticks([], [])
+    plt.yticks([], [])
+
+    legend_elements = []
+    for class_name in structure_class_names:
+        color, alpha, linewidth, hatch = get_bbox_decorations('recognition', structure_class_map[class_name])
+        legend_elements.append(
+            Patch(facecolor=color, edgecolor=color, label=class_name, hatch=hatch, alpha=alpha)
+        )
+
+    plt.legend(handles=legend_elements, bbox_to_anchor=(0.5, -0.02), loc='upper center', borderaxespad=0,
+                    fontsize=10, ncol=3)
+    plt.gcf().set_size_inches(10, 10)
     plt.axis('off')
     img_buf = io.BytesIO()
-    plt.savefig(img_buf, bbox_inches='tight', dpi=1000)
+    plt.savefig(img_buf, bbox_inches='tight', dpi=150)
+    plt.close()
 
     return PIL.Image.open(img_buf)
 
 
 def visualize_cells(pil_img, cells):
-    fig, ax = plt.subplots(1)
-    ax.imshow(pil_img, interpolation='lanczos')
+    plt.imshow(pil_img, interpolation="lanczos")
+    plt.gcf().set_size_inches(20, 20)
+    ax = plt.gca()
 
-    for i, cell in enumerate(cells):
+    for cell in cells:
         bbox = cell['bbox']
-        if cell['header']:
+
+        if cell['column header']:
+            facecolor = (1, 0, 0.45)
+            edgecolor = (1, 0, 0.45)
             alpha = 0.3
+            linewidth = 2
+            hatch='//////'
+        elif cell['projected row header']:
+            facecolor = (0.95, 0.6, 0.1)
+            edgecolor = (0.95, 0.6, 0.1)
+            alpha = 0.3
+            linewidth = 2
+            hatch='//////'
         else:
-            alpha = 0.125
-        rect = patches.Rectangle(bbox[:2], bbox[2]-bbox[0], bbox[3]-bbox[1], linewidth=1,
-                                edgecolor='none',facecolor="magenta", alpha=alpha)
+            facecolor = (0.3, 0.74, 0.8)
+            edgecolor = (0.3, 0.7, 0.6)
+            alpha = 0.3
+            linewidth = 2
+            hatch='\\\\\\\\\\\\'
+
+        rect = patches.Rectangle(bbox[:2], bbox[2]-bbox[0], bbox[3]-bbox[1], linewidth=linewidth,
+                                    edgecolor='none',facecolor=facecolor, alpha=0.1)
         ax.add_patch(rect)
-        rect = patches.Rectangle(bbox[:2], bbox[2]-bbox[0], bbox[3]-bbox[1], linewidth=1,
-                                edgecolor="magenta",facecolor='none',linestyle="--",
-                                alpha=0.08, hatch='///')
+        rect = patches.Rectangle(bbox[:2], bbox[2]-bbox[0], bbox[3]-bbox[1], linewidth=linewidth,
+                                    edgecolor=edgecolor,facecolor='none',linestyle='-', alpha=alpha)
         ax.add_patch(rect)
-        rect = patches.Rectangle(bbox[:2], bbox[2]-bbox[0], bbox[3]-bbox[1], linewidth=1,
-                                edgecolor="magenta",facecolor='none',linestyle="--")
+        rect = patches.Rectangle(bbox[:2], bbox[2]-bbox[0], bbox[3]-bbox[1], linewidth=0,
+                                    edgecolor=edgecolor,facecolor='none',linestyle='-', hatch=hatch, alpha=0.2)
         ax.add_patch(rect)
 
+    plt.xticks([], [])
+    plt.yticks([], [])
+
+    legend_elements = [Patch(facecolor=(0.3, 0.74, 0.8), edgecolor=(0.3, 0.7, 0.6),
+                                label='Data cell', hatch='\\\\\\\\\\\\', alpha=0.3),
+                        Patch(facecolor=(1, 0, 0.45), edgecolor=(1, 0, 0.45),
+                                label='Column header cell', hatch='//////', alpha=0.3),
+                        Patch(facecolor=(0.95, 0.6, 0.1), edgecolor=(0.95, 0.6, 0.1),
+                                label='Projected row header cell', hatch='//////', alpha=0.3)]
+    plt.legend(handles=legend_elements, bbox_to_anchor=(0.5, -0.02), loc='upper center', borderaxespad=0,
+                    fontsize=10, ncol=3)
+    plt.gcf().set_size_inches(10, 10)
     plt.axis('off')
     img_buf = io.BytesIO()
-    plt.savefig(img_buf, bbox_inches='tight', dpi=1000)
+    plt.savefig(img_buf, bbox_inches='tight', dpi=150)
+    plt.close()
 
     return PIL.Image.open(img_buf)