MikkoLipsanen commited on
Commit
aca3760
·
verified ·
1 Parent(s): ed565bd

Create plotting_functions.py

Browse files
Files changed (1) hide show
  1. plotting_functions.py +39 -0
plotting_functions.py ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import matplotlib.pyplot as plt
2
+ from PIL import Image, ImageDraw
3
+
4
+ class PlotHTR:
5
+ def __init__(self, region_alpha = 0.3, line_alpha = 0.3):
6
+ self.region_colors = {'marginalia': (0,201,87), 'page-number': (0,0,255), 'paragraph': (238,18,137), 'case-start': (238,18,137)}
7
+ self.region_alpha = region_alpha
8
+ self.line_alpha = line_alpha
9
+
10
+ def plot_regions(self, region_data, image):
11
+ """Plots the detected text regions."""
12
+ img = image.copy()
13
+ draw = ImageDraw.Draw(img)
14
+ for region in region_data:
15
+ region_type = region['region_name']
16
+ img_h, img_w = region['img_shape']
17
+ region_polygon = list(map(tuple, region['region_coords']))
18
+
19
+ color = self.region_colors[region_type]
20
+
21
+ draw.polygon(region_polygon, fill = color, outline='black')
22
+ # Lower alpha value is applied to filled polygons to make them transparent
23
+ blended_img = Image.blend(image, img, self.region_alpha)
24
+ return blended_img
25
+
26
+ def plot_lines(self, region_data, image):
27
+ """Plots the detected text lines."""
28
+ img = image.copy()
29
+ draw = ImageDraw.Draw(img)
30
+ for region in region_data:
31
+ img_h, img_w = region['img_shape']
32
+ line_polygons = region['lines']
33
+ region_type = region['region_name']
34
+ color = self.region_colors[region_type]
35
+ for i, polygon in enumerate(line_polygons):
36
+ draw.polygon(polygon, fill = color)
37
+ # Lower alpha value is applied to filled polygons to make them transparent
38
+ blended_img = Image.blend(image, img, self.line_alpha)
39
+ return blended_img