CAD-AID / src /utils /plotting.py
Julia Jørstad
First version
452a352
import cv2
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import streamlit as st
def plot_results_streamlit(image_path, masks):
""" Plot results and return the image """
original_image = cv2.imread(image_path)
original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
fig, ax = plt.subplots(figsize=(8, 8))
ax.imshow(original_image)
if masks:
for mask in masks:
mask = np.array(mask, dtype=np.uint8)
ax.imshow(mask, cmap="jet", alpha=0.4)
ax.axis("off")
#ax.set_title("Segmentation Results")
return fig
def draw_ocr_boxes(image_np, df_ocr, color=(0, 255, 0), thickness=2):
"""
Draws bounding boxes around OCR-detected text on the image.
Args:
image_np (np.ndarray): The original image as NumPy array (RGB).
df_ocr (pd.DataFrame): DataFrame with 'text' and 'box' columns.
color (tuple): Color for the boxes (default green).
thickness (int): Thickness of the rectangle.
Returns:
np.ndarray: Image with OCR bounding boxes drawn.
"""
output_img = image_np.copy()
for _, row in df_ocr.iterrows():
if 'box' in row and isinstance(row['box'], (tuple, list)) and len(row['box']) == 4:
x1, y1, x2, y2 = map(int, row['box'])
cv2.rectangle(output_img, (x1, y1), (x2, y2), color, thickness)
# Optional: draw text label
cv2.putText(output_img, row['text'], (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX,
fontScale=0.5, color=color, thickness=1, lineType=cv2.LINE_AA)
return output_img