src/utils/plotting.py

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