utils/visualization.py

import matplotlib.pyplot as plt
from typing import List, Dict
import io
from PIL import Image

def plot_results(results: List[Dict], query: str, num_images: int = 3) -> Image.Image:
    """
    Plot search results as image grid.
    
    Returns:
        PIL Image of the plot
    """
    fig, axes = plt.subplots(1, num_images, figsize=(20, 7))
    if num_images == 1:
        axes = [axes]
    
    fig.suptitle(f"🔍 Query: {query}", fontsize=18, fontweight='bold')
    
    for i, ax in enumerate(axes):
        if i < len(results):
            result = results[i]
            img = plt.imread(result["image_path"])
            ax.imshow(img)
            ax.set_title(
                f"Rank {result['rank']} | Score: {result['score']:.4f}",
                fontsize=14
            )
        ax.axis('off')
    
    plt.tight_layout()
    
    # Convert to PIL Image
    buf = io.BytesIO()
    plt.savefig(buf, format='png', bbox_inches='tight', dpi=100)
    buf.seek(0)
    img_pil = Image.open(buf)
    plt.close()
    
    return img_pil