brunvelop/ComfyUI

1import datetime2import numpy as np3import os4from PIL import Image5import pytest6from pytest import fixture7from typing import Tuple, List8 9from cv2 import imread, cvtColor, COLOR_BGR2RGB10from skimage.metrics import structural_similarity as ssim11 12 13"""14This test suite compares images in 2 directories by file name15The directories are specified by the command line arguments --baseline_dir and --test_dir16 17"""18# ssim: Structural Similarity Index19# Returns a tuple of (ssim, diff_image)20def ssim_score(img0: np.ndarray, img1: np.ndarray) -> Tuple[float, np.ndarray]:21    score, diff = ssim(img0, img1, channel_axis=-1, full=True)22    # rescale the difference image to 0-255 range23    diff = (diff * 255).astype("uint8")24    return score, diff25    26# Metrics must return a tuple of (score, diff_image)27METRICS = {"ssim": ssim_score}28METRICS_PASS_THRESHOLD = {"ssim": 0.95}29 30 31class TestCompareImageMetrics:32    @fixture(scope="class")33    def test_file_names(self, args_pytest):34        test_dir = args_pytest['test_dir']35        fnames = self.gather_file_basenames(test_dir)  36        yield fnames37        del fnames38 39    @fixture(scope="class", autouse=True)40    def teardown(self, args_pytest):41        yield42        # Runs after all tests are complete43        # Aggregate output files into a grid of images44        baseline_dir = args_pytest['baseline_dir']45        test_dir = args_pytest['test_dir']46        img_output_dir = args_pytest['img_output_dir']47        metrics_file = args_pytest['metrics_file']48 49        grid_dir = os.path.join(img_output_dir, "grid")50        os.makedirs(grid_dir, exist_ok=True)51 52        for metric_dir in METRICS.keys():53            metric_path = os.path.join(img_output_dir, metric_dir)54            for file in os.listdir(metric_path):55                if file.endswith(".png"):56                    score = self.lookup_score_from_fname(file, metrics_file)57                    image_file_list = []58                    image_file_list.append([59                                            os.path.join(baseline_dir, file),60                                            os.path.join(test_dir, file),61                                            os.path.join(metric_path, file)62                                            ])63                    # Create grid64                    image_list = [[Image.open(file) for file in files] for files in image_file_list]65                    grid = self.image_grid(image_list)66                    grid.save(os.path.join(grid_dir, f"{metric_dir}_{score:.3f}_{file}"))67    68    # Tests run for each baseline file name69    @fixture()70    def fname(self, baseline_fname):71        yield baseline_fname72        del baseline_fname73    74    def test_directories_not_empty(self, args_pytest):75        baseline_dir = args_pytest['baseline_dir']76        test_dir = args_pytest['test_dir']77        assert len(os.listdir(baseline_dir)) != 0, f"Baseline directory {baseline_dir} is empty"78        assert len(os.listdir(test_dir)) != 0, f"Test directory {test_dir} is empty"79 80    def test_dir_has_all_matching_metadata(self, fname, test_file_names, args_pytest):81        # Check that all files in baseline_dir have a file in test_dir with matching metadata82        baseline_file_path = os.path.join(args_pytest['baseline_dir'], fname)83        file_paths = [os.path.join(args_pytest['test_dir'], f) for f in test_file_names]84        file_match = self.find_file_match(baseline_file_path, file_paths)85        assert file_match is not None, f"Could not find a file in {args_pytest['test_dir']} with matching metadata to {baseline_file_path}"86 87    # For a baseline image file, finds the corresponding file name in test_dir and 88    # compares the images using the metrics in METRICS89    @pytest.mark.parametrize("metric", METRICS.keys())90    def test_pipeline_compare(91        self,92        args_pytest,93        fname,94        test_file_names,95        metric,96    ):97        baseline_dir = args_pytest['baseline_dir']98        test_dir = args_pytest['test_dir']99        metrics_output_file = args_pytest['metrics_file']100        img_output_dir = args_pytest['img_output_dir']101        102        baseline_file_path = os.path.join(baseline_dir, fname)103 104        # Find file match105        file_paths = [os.path.join(test_dir, f) for f in test_file_names]106        test_file = self.find_file_match(baseline_file_path, file_paths)107 108        # Run metrics109        sample_baseline = self.read_img(baseline_file_path)110        sample_secondary = self.read_img(test_file)111        112        score, metric_img = METRICS[metric](sample_baseline, sample_secondary)113        metric_status = score > METRICS_PASS_THRESHOLD[metric]114 115        # Save metric values116        with open(metrics_output_file, 'a') as f:117            run_info = os.path.splitext(fname)[0]118            metric_status_str = "PASS ✅" if metric_status else "FAIL ❌"119            date_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")120            f.write(f"| {date_str} | {run_info} | {metric} | {metric_status_str} | {score} | \n")121 122        # Save metric image123        metric_img_dir = os.path.join(img_output_dir, metric)124        os.makedirs(metric_img_dir, exist_ok=True)125        output_filename = f'{fname}'126        Image.fromarray(metric_img).save(os.path.join(metric_img_dir, output_filename))127 128        assert score > METRICS_PASS_THRESHOLD[metric]129 130    def read_img(self, filename: str) -> np.ndarray:131        cvImg = imread(filename)132        cvImg = cvtColor(cvImg, COLOR_BGR2RGB)133        return cvImg134 135    def image_grid(self, img_list: list[list[Image.Image]]):136        # imgs is a 2D list of images137        # Assumes the input images are a rectangular grid of equal sized images138        rows = len(img_list)139        cols = len(img_list[0])140 141        w, h = img_list[0][0].size142        grid = Image.new('RGB', size=(cols*w, rows*h))143        144        for i, row in enumerate(img_list):145            for j, img in enumerate(row):146                grid.paste(img, box=(j*w, i*h))147        return grid148 149    def lookup_score_from_fname(self,150                                fname: str,151                                metrics_output_file: str152        ) -> float:153        fname_basestr = os.path.splitext(fname)[0]154        with open(metrics_output_file, 'r') as f:155            for line in f:156                if fname_basestr in line:157                    score = float(line.split('|')[5])158                    return score159        raise ValueError(f"Could not find score for {fname} in {metrics_output_file}")160 161    def gather_file_basenames(self, directory: str):162        files = []163        for file in os.listdir(directory):164            if file.endswith(".png"):165                files.append(file)166        return files167 168    def read_file_prompt(self, fname:str) -> str:169        # Read prompt from image file metadata170        img = Image.open(fname)171        img.load()172        return img.info['prompt']173    174    def find_file_match(self, baseline_file: str, file_paths: List[str]):175        # Find a file in file_paths with matching metadata to baseline_file176        baseline_prompt = self.read_file_prompt(baseline_file)177 178        # Do not match empty prompts179        if baseline_prompt is None or baseline_prompt == "":180            return None181 182        # Find file match183        # Reorder test_file_names so that the file with matching name is first184        # This is an optimization because matching file names are more likely 185        # to have matching metadata if they were generated with the same script186        basename = os.path.basename(baseline_file)187        file_path_basenames = [os.path.basename(f) for f in file_paths]188        if basename in file_path_basenames:189            match_index = file_path_basenames.index(basename)190            file_paths.insert(0, file_paths.pop(match_index))191 192        for f in file_paths:193            test_file_prompt = self.read_file_prompt(f)194            if baseline_prompt == test_file_prompt:195                return f