from typing import Dict, List, Tuple, TypedDict import datasets as ds import evaluate import numpy as np import numpy.typing as npt _DESCRIPTION = """\ Computes some average IoU metrics that are different to each other in previous works. """ _KWARGS_DESCRIPTION = """\ Args: layouts (`list` of `dict`): A list of dictionaries representing layouts including `list` of `bboxes` (float) and `list` of `categories` (int). Returns: dicrionaly: A set of average IoU scores. Examples: Example 1: Single processing >>> metric = evaluate.load("creative-graphic-design/layout-average-iou") >>> num_samples, num_categories = 24, 4 >>> layout = { >>> "bboxes": np.random.rand(num_samples, num_categories), >>> "categories": np.random.randint(0, num_categories, size=(num_samples,)), >>> } >>> metric.add(layouts=layout) >>> print(metric.compute()) Example 2: Batch processing >>> metric = evaluate.load("creative-graphic-design/layout-average-iou") >>> batch_size, num_samples, num_categories = 512, 24, 4 >>> layouts = [ >>> { >>> "bboxes": np.random.rand(num_samples, num_categories), >>> "categories": np.random.randint(0, num_categories, size=(num_samples,)), >>> } >>> for _ in range(batch_size) >>> ] >>> metric.add_batch(layouts=layouts) >>> print(metric.compute()) """ _CITATION = """\ @inproceedings{arroyo2021variational, title={Variational transformer networks for layout generation}, author={Arroyo, Diego Martin and Postels, Janis and Tombari, Federico}, booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition}, pages={13642--13652}, year={2021} } @inproceedings{kong2022blt, title={BLT: bidirectional layout transformer for controllable layout generation}, author={Kong, Xiang and Jiang, Lu and Chang, Huiwen and Zhang, Han and Hao, Yuan and Gong, Haifeng and Essa, Irfan}, booktitle={European Conference on Computer Vision}, pages={474--490}, year={2022}, organization={Springer} } """ def convert_xywh_to_ltrb( batch_bbox: npt.NDArray[np.float64], ) -> Tuple[ npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64], ]: xc, yc, w, h = batch_bbox x1 = xc - w / 2 y1 = yc - h / 2 x2 = xc + w / 2 y2 = yc + h / 2 return (x1, y1, x2, y2) class Layout(TypedDict): bboxes: npt.NDArray[np.float64] categories: npt.NDArray[np.int64] def compute_iou( bbox1: npt.NDArray[np.float64], bbox2: npt.NDArray[np.float64], generalized: bool = False, ) -> npt.NDArray[np.float64]: # shape: bbox1 (N, 4), bbox2 (N, 4) assert bbox1.shape[0] == bbox2.shape[0] assert bbox1.shape[1] == bbox1.shape[1] == 4 l1, t1, r1, b1 = convert_xywh_to_ltrb(bbox1.T) l2, t2, r2, b2 = convert_xywh_to_ltrb(bbox2.T) a1, a2 = (r1 - l1) * (b1 - t1), (r2 - l2) * (b2 - t2) # intersection l_max = np.maximum(l1, l2) r_min = np.minimum(r1, r2) t_max = np.maximum(t1, t2) b_min = np.minimum(b1, b2) cond = (l_max < r_min) & (t_max < b_min) ai = np.where(cond, (r_min - l_max) * (b_min - t_max), np.zeros_like(a1[0])) au = a1 + a2 - ai iou = ai / au if not generalized: return iou # outer region l_min = np.minimum(l1, l2) r_max = np.maximum(r1, r2) t_min = np.minimum(t1, t2) b_max = np.maximum(b1, b2) ac = (r_max - l_min) * (b_max - t_min) giou = iou - (ac - au) / ac return giou def compute_perceptual_iou( bbox1: npt.NDArray[np.float64], bbox2: npt.NDArray[np.float64], N: int = 32, ) -> npt.NDArray[np.float64]: """ Computes 'Perceptual' IoU [Kong+, BLT'22] """ # shape: bbox1 (N, 4), bbox2 (N, 4) assert bbox1.shape[0] == bbox2.shape[0] assert bbox1.shape[1] == bbox1.shape[1] == 4 l1, t1, r1, b1 = convert_xywh_to_ltrb(bbox1.T) l2, t2, r2, b2 = convert_xywh_to_ltrb(bbox2.T) a1 = (r1 - l1) * (b1 - t1) # intersection l_max = np.maximum(l1, l2) r_min = np.minimum(r1, r2) t_max = np.maximum(t1, t2) b_min = np.minimum(b1, b2) cond = (l_max < r_min) & (t_max < b_min) ai = np.where(cond, (r_min - l_max) * (b_min - t_max), np.zeros_like(a1[0])) unique_box_1 = np.unique(bbox1, axis=0) l1, t1, r1, b1 = [ (x * N).round().astype(np.int32).clip(0, N) for x in convert_xywh_to_ltrb(unique_box_1.T) ] canvas = np.zeros((N, N)) for left, top, right, bottom in zip(l1, t1, r1, b1): canvas[top:bottom, left:right] = 1 global_area_union = canvas.sum() / (N**2) return ai / global_area_union if global_area_union > 0.0 else np.zeros((1,)) def compute_average_iou(layout: Layout, perceptual: bool) -> float: bboxes = np.asarray(layout["bboxes"]) N = len(bboxes) if N in [0, 1]: return 0.0 # no overlap in principle ii, jj = np.meshgrid(range(N), range(N)) ii, jj = ii.flatten(), jj.flatten() is_non_diag = ii != jj # IoU for diag is always 1.0 ii, jj = ii[is_non_diag], jj[is_non_diag] iou = ( compute_perceptual_iou(bboxes[ii], bboxes[jj]) if perceptual else compute_iou(bboxes[ii], bboxes[jj]) ) # pick all pairs of overlapped objects cond = iou > np.finfo(np.float32).eps # to avoid very-small nonzero score = iou[cond].mean().item() if len(iou[cond]) > 0 else 0.0 return score class LayoutAverageIoU(evaluate.Metric): def _info(self) -> evaluate.EvaluationModuleInfo: return evaluate.EvaluationModuleInfo( description=_DESCRIPTION, citation=_CITATION, inputs_description=_KWARGS_DESCRIPTION, features=ds.Features( { "layouts": { "bboxes": ds.Sequence(ds.Sequence((ds.Value("float64")))), "categories": ds.Sequence(ds.Value("int64")), } } ), codebase_urls=[ "https://github.com/CyberAgentAILab/layout-dm/blob/main/src/trainer/trainer/helpers/metric.py#L399-L431", ], ) def _compute(self, *, layouts: List[Layout]) -> Dict[str, float]: scores_blt = [ compute_average_iou(layout, perceptual=True) for layout in layouts ] scores_vnt = [ compute_average_iou(layout, perceptual=False) for layout in layouts ] score_blt = np.mean(scores_blt).item() score_vnt = np.mean(scores_vnt).item() results = { "average-iou_BLT": score_blt, "average-iou_VTN": score_vnt, } return results