report_utils.py

import json
from typing import Dict


def prepare_report(scheduler_name: str, results: dict):
    image_grid = results["images"]
    scores = results["scores"]
    img_str = ""

    image_name = f"{scheduler_name}_images.png"
    image_grid.save(image_name)
    img_str = img_str = f"![img_grid_{scheduler_name}](/file=./{image_name})\n"

    report_str = f"""
\n\n## {scheduler_name}

### Sample images

{img_str}

### Scores

{scores}
\n\n
    """

    return report_str


def add_psnr_ssim_to_report(
    original_scheduler_name: str, ssim_scores: Dict = None, psnr_scores: Dict = None
) -> str:
    current_str = ""
    if ssim_scores is not None:
        current_str += f"""
\n\n
## SSIM

SSIM computed w.r.t the images generated with {original_scheduler_name}:\n\n {json.dumps(ssim_scores, indent=6)}
"""
    if psnr_scores is not None:
        current_str += f"""
\n\n
## PSNR

PSNR computed w.r.t the images generated with {original_scheduler_name}:\n\n {json.dumps(psnr_scores, indent=6)}
"""
    return current_str