plotting_functions.py

import matplotlib.pyplot as plt
from PIL import Image, ImageDraw

class PlotHTR:
    def __init__(self, region_alpha = 0.3, line_alpha = 0.3):
        self.polygon_color = (238,18,137)
        self.region_alpha = region_alpha 
        self.line_alpha = line_alpha 

    def plot_regions(self, region_data, image):
        """Plots the detected text regions."""
        img = image.copy()
        draw = ImageDraw.Draw(img)
        for region in region_data:
            img_h, img_w = region['img_shape']
            region_polygon = list(map(tuple, region['region_coords']))
            draw.polygon(region_polygon, fill = self.polygon_color, outline='black')
        # Lower alpha value is applied to filled polygons to make them transparent
        blended_img = Image.blend(image, img, self.region_alpha)
        return blended_img

    def plot_lines(self, region_data, image):
        """Plots the detected text lines."""
        img = image.copy()
        draw = ImageDraw.Draw(img)
        for region in region_data:
            img_h, img_w = region['img_shape']
            line_polygons = region['lines']
            print(f'{len(line_polygons)} text lines')
            for i, polygon in enumerate(line_polygons):
                line_polygon = list(map(tuple, polygon))
                draw.polygon(line_polygon, fill = self.polygon_color, outline='black')
        # Lower alpha value is applied to filled polygons to make them transparent
        blended_img = Image.blend(image, img, self.line_alpha)
        return blended_img