# Copyright (c) 2025 Hansheng Chen import sys import os import re import unicodedata import numpy as np import torch import torch.distributed as dist import mmcv from copy import deepcopy from concurrent.futures import ThreadPoolExecutor from mmcv.runner import HOOKS, get_dist_info from mmcv.fileio import FileClient from mmgen.models.architectures.common import get_module_device from mmgen.core import GenerativeEvalHook as _GenerativeEvalHook from lakonlab.utils.io_utils import save_image, save_video, load_images_parallel from lakonlab.runner.timer import default_timers from lakonlab.utils import gc_context from lakonlab.ui.media_viewer import write_html default_timers.add_timer('total time') _reserved = {"CON", "PRN", "AUX", "NUL", *(f"COM{i}" for i in range(1, 10)), *(f"LPT{i}" for i in range(1, 10))} _invalid = re.compile(r'[<>:"/\\|?*\:%]+') def flatten_list(lst): for item in lst: if isinstance(item, list): yield from flatten_list(item) # recurse into sub-list else: yield item def _safe_name(name: str, max_len: int = 160) -> str: s = unicodedata.normalize("NFKC", str(name)) s = "".join(c for c in s if 32 <= ord(c) != 127) # drop control chars s = _invalid.sub("", s).strip(" ._") # strip invalid chars if s.startswith("."): s = s.lstrip(".") # avoid hidden files if s.upper() in _reserved: s += "_" # avoid reserved names s = (s or "untitled")[:max_len].rstrip(" .") # truncate & clean return s or "untitled" def evaluate(model, dataloader, metrics=None, feed_batch_size=32, viz_dir=None, viz_num=None, sample_kwargs=dict(), fps=16, enable_timers=False, reuse_viz=False): has_metrics = metrics is not None and len(metrics) > 0 if has_metrics: for metric in metrics: if hasattr(metric, 'load_to_gpu'): metric.load_to_gpu() if enable_timers: default_timers.enable_all() batch_size = dataloader.batch_size rank, ws = get_dist_info() total_batch_size = batch_size * ws max_num_fakes = len(dataloader.dataset) if viz_dir is not None: html_entries = [] saved_data_ids = set() file_client = FileClient.infer_client(uri=viz_dir) executor = ThreadPoolExecutor(max_workers=(os.cpu_count() or 4) * 4) else: html_entries = saved_data_ids = file_client = executor = None if rank == 0: mmcv.print_log( f'Generate {max_num_fakes} fake samples for evaluation', 'mmgen') pbar = mmcv.ProgressBar(max_num_fakes) log_vars = dict() batch_size_list = [] for i, data in enumerate(dataloader): can_reuse = False loaded_imgs_tensor = None loaded_html_entries = None if viz_dir is not None and reuse_viz: batch_names = list(flatten_list(data['name'].data)) ids = list(flatten_list(data['ids'].data)) reuse_candidates = [] all_png_exist = True for name, data_id in zip(batch_names, ids): safe_name = _safe_name(name) filename = f'{data_id:09d}_{safe_name}.png' filepath = file_client.join_path(viz_dir, filename) if file_client.isfile(filepath): reuse_candidates.append((data_id, name, filename, filepath)) else: all_png_exist = False break if all_png_exist: filepaths = [] loaded_html_entries = [] for (data_id, name, filename, filepath) in reuse_candidates: filepaths.append(filepath) rel_filepath = file_client.join_path(os.path.basename(viz_dir), filename) loaded_html_entries.append((data_id, rel_filepath, name)) try: loaded_imgs_tensor = torch.from_numpy( np.stack(load_images_parallel(filepaths, file_client), axis=0) ).permute(0, 3, 1, 2).float() / 255.0 can_reuse = True except Exception as e: can_reuse = False loaded_imgs_tensor = None loaded_html_entries = None if can_reuse: outputs_dict = dict( pred_imgs=loaded_imgs_tensor, # [0,1] float32 num_samples=loaded_imgs_tensor.size(0)) html_entries.extend(loaded_html_entries) else: # fall back to generation sample_kwargs_ = deepcopy(sample_kwargs) with default_timers['total time']: outputs_dict = model.val_step(data, show_pbar=rank == 0, **sample_kwargs_) if viz_dir is not None: batch_names = list(flatten_list(data['name'].data)) for batch_id, data_id in enumerate(flatten_list(data['ids'].data)): if (viz_num is not None and data_id >= viz_num) or (data_id in saved_data_ids): continue name = batch_names[batch_id] safe_name = _safe_name(name) image_viz = (outputs_dict['pred_imgs'][batch_id] * 255).round().to(torch.uint8) if image_viz.dim() == 3: # image image_viz = image_viz.permute(1, 2, 0).cpu().numpy() filename = f'{data_id:09d}_{safe_name}.png' executor.submit( save_image, image_viz, file_client.join_path(viz_dir, filename), file_client) elif image_viz.dim() == 4: # video image_viz = image_viz.permute(1, 2, 3, 0).cpu().numpy() # (t, h, w, c) filename = f'{data_id:09d}_{safe_name}.mp4' executor.submit( save_video, image_viz, file_client.join_path(viz_dir, filename), file_client, fps) else: raise ValueError(f'Unsupported image dimension: {image_viz.dim()}') rel_filepath = file_client.join_path(os.path.basename(viz_dir), filename) html_entries.append((data_id, rel_filepath, name)) saved_data_ids.add(data_id) if 'log_vars' in outputs_dict: for k, v in outputs_dict['log_vars'].items(): if k in log_vars: log_vars[k].append(outputs_dict['log_vars'][k]) else: log_vars[k] = [outputs_dict['log_vars'][k]] batch_size_list.append(outputs_dict['num_samples']) if has_metrics: pred_imgs = outputs_dict['pred_imgs'].split(feed_batch_size, dim=0) real_imgs = None if 'images' in data: real_imgs = data['images'] elif 'target_imgs' in outputs_dict: real_imgs = outputs_dict['target_imgs'] if real_imgs is not None: real_imgs = real_imgs.split(feed_batch_size, dim=0) requires_prompt = False for metric in metrics: requires_prompt |= getattr(metric, 'requires_prompt', False) if requires_prompt: prompts = list(flatten_list(data['name'].data)) # list of prompts prompts = [prompts[i:i + feed_batch_size] for i in range(0, len(prompts), feed_batch_size)] for metric in metrics: for batch_id, batch_imgs in enumerate(pred_imgs): if getattr(metric, 'requires_prompt', False): metric.feed( dict(imgs=batch_imgs * 2 - 1, prompts=prompts[batch_id]), 'fakes') if real_imgs is not None: metric.feed( dict(imgs=real_imgs[batch_id] * 2 - 1, prompts=prompts[batch_id]), 'reals') else: metric.feed(batch_imgs * 2 - 1, 'fakes') if real_imgs is not None: metric.feed(real_imgs[batch_id] * 2 - 1, 'reals') if rank == 0: pbar.update(total_batch_size) if ws > 1: device = get_module_device(model) batch_size_list = torch.tensor(batch_size_list, dtype=torch.float, device=device) batch_size_sum = torch.sum(batch_size_list) dist.all_reduce(batch_size_sum, op=dist.ReduceOp.SUM) for k, v in log_vars.items(): weigted_values = torch.tensor(log_vars[k], dtype=torch.float, device=device) * batch_size_list weigted_values_sum = torch.sum(weigted_values) dist.all_reduce(weigted_values_sum, op=dist.ReduceOp.SUM) log_vars[k] = float(weigted_values_sum / batch_size_sum) else: for k, v in log_vars.items(): log_vars[k] = np.average(log_vars[k], weights=batch_size_list) if viz_dir is not None: if ws > 1: gathered = [None for _ in range(ws)] dist.all_gather_object(gathered, html_entries) if rank == 0: html_entries = [e for sub in gathered for e in (sub or [])] if rank == 0: unique_entries = dict() for entry in html_entries: data_id = entry[0] if data_id not in unique_entries: unique_entries[data_id] = entry html_entries = list(unique_entries.values()) html_entries.sort(key=lambda item: item[0]) html_path = file_client.join_path(os.path.dirname(viz_dir), os.path.basename(viz_dir) + '.html') write_html(html_path, html_entries, file_client) executor.shutdown(wait=True) return log_vars @HOOKS.register_module(force=True) class GenerativeEvalHook(_GenerativeEvalHook): greater_keys = ['acc', 'top', 'AR@', 'auc', 'precision', 'mAP', 'is', 'test_ssim', 'test_psnr'] less_keys = ['loss', 'fid', 'kid', 'test_lpips'] _supported_best_metrics = ['fid', 'kid', 'is', 'test_ssim', 'test_psnr', 'test_lpips'] def __init__(self, *args, data='', viz_dir=None, feed_batch_size=32, viz_num=None, clear_reals=False, prefix='', metric_cpu_offload=False, **kwargs): super(GenerativeEvalHook, self).__init__(*args, **kwargs) self.data = data self.viz_dir = viz_dir self.file_client = FileClient.infer_client( uri=viz_dir) if viz_dir is not None else None self.feed_batch_size = feed_batch_size self.viz_num = viz_num self.clear_reals = clear_reals self.prefix = prefix self.metric_cpu_offload = metric_cpu_offload @torch.no_grad() def after_train_iter(self, runner): with gc_context(enable=True): interval = self.get_current_interval(runner) if not self.every_n_iters(runner, interval): return runner.model.eval() rank, ws = get_dist_info() if self.viz_dir is not None: viz_dir = self.file_client.join_path(self.viz_dir, str(runner.iter + 1)) if rank == 0: if self.file_client.exists(viz_dir): for name in self.file_client.list_dir_or_file(viz_dir): self.file_client.remove(self.file_client.join_path(viz_dir, name)) if ws > 1: dist.barrier() else: viz_dir = None log_vars = evaluate( runner.model, self.dataloader, self.metrics, self.feed_batch_size, viz_dir, self.viz_num, self.sample_kwargs) if len(runner.log_buffer.output) == 0: runner.log_buffer.clear() # a dirty walkround to change the line at the end of pbar if rank == 0: sys.stdout.write('\n') for metric in self.metrics: metric.summary() for name, val in metric._result_dict.items(): prefix_name = self.prefix + '_' + name if len(self.prefix) > 0 else name runner.log_buffer.output[self.data + '_' + prefix_name] = val # record best metric and save the best ckpt if self.save_best_ckpt and name in self.best_metric: self._save_best_ckpt(runner, val, name) for name, val in log_vars.items(): prefix_name = self.prefix + '_' + name if len(self.prefix) > 0 else name # print(self.data + '_' + prefix_name + ' = {}'.format(val)) runner.log_buffer.output[self.data + '_' + prefix_name] = val # record best metric and save the best ckpt if self.save_best_ckpt and name in self.best_metric: self._save_best_ckpt(runner, val, name) runner.log_buffer.ready = True runner.model.train() for metric in self.metrics: metric.clear(clear_reals=self.clear_reals) if self.metric_cpu_offload: if hasattr(metric, 'offload_to_cpu'): metric.offload_to_cpu() torch.cuda.empty_cache()