Spaces:
Running
Running
| # -*- coding: utf-8 -*- | |
| """Utils for display input and output""" | |
| import matplotlib.pyplot as plt | |
| from PIL import Image | |
| def open_image(image_path: str) -> Image.Image: | |
| """Open an image from a path.""" | |
| return Image.open(image_path) | |
| def display_baselines(image, baselines, boundaries=None): | |
| """Display baselines and boundaries on an image.""" | |
| fig, ax = plt.subplots(figsize=(10, 10)) | |
| ax.imshow(image, cmap='gray') | |
| ax.axis('off') | |
| for idx, baseline in enumerate(baselines): | |
| baseline_x = [point[0] for point in baseline] | |
| baseline_y = [point[1] for point in baseline] | |
| ax.plot(baseline_x, baseline_y, color='blue', linewidth=0.7) | |
| if boundaries: | |
| for boundary in boundaries: | |
| boundary_x = [point[0] for point in boundary] | |
| boundary_y = [point[1] for point in boundary] | |
| ax.plot(boundary_x, boundary_y, color='red', linestyle='--', linewidth=1) | |
| fig_special, ax_special = plt.subplots(figsize=(10, 10)) | |
| ax_special.set_xlim(0, image.size[0]) | |
| ax_special.set_ylim(0, image.size[1]) | |
| ax_special.invert_yaxis() | |
| for idx, baseline in enumerate(baselines): | |
| baseline_x = [point[0] for point in baseline] | |
| baseline_y = [point[1] for point in baseline] | |
| ax_special.plot(baseline_x, baseline_y, color='blue', linewidth=0.7) | |
| ax_special.text(baseline_x[0], baseline_y[0], str(idx), fontsize=10, color='red') | |
| return fig, fig_special | |
| def display_baselines_with_text(image, baselines, lines): | |
| """Display baselines with text on an image.""" | |
| fig_special, ax_special = plt.subplots(figsize=(10, 10)) | |
| ax_special.set_xlim(0, image.size[0]) | |
| ax_special.set_ylim(0, image.size[1]) | |
| ax_special.invert_yaxis() | |
| for idx, group in enumerate(zip(lines, baselines)): | |
| baseline_x = [point[0] for point in group[1]] | |
| baseline_y = [point[1] for point in group[1]] | |
| ax_special.text(baseline_x[0], baseline_y[0], f"{str(idx)}: {group[0]}", fontsize=5.5, color='black') | |
| ax_special.axis('off') | |
| return fig_special | |
| def prepare_segments(seg_obj): | |
| """Prepare baselines and boundaries for display.""" | |
| baselines = [] | |
| boundaries = [] | |
| for line in seg_obj.lines: | |
| baselines.append(line.baseline) | |
| boundaries.append(line.boundary) | |
| return baselines, boundaries |