coreml-community/ControlNet-v1-1-Annotators-cpu

1# Copyright (c) Facebook, Inc. and its affiliates.2import contextlib3import io4import itertools5import json6import logging7import numpy as np8import os9import tempfile10from collections import OrderedDict11from typing import Optional12from PIL import Image13from tabulate import tabulate14 15from annotator.oneformer.detectron2.data import MetadataCatalog16from annotator.oneformer.detectron2.utils import comm17from annotator.oneformer.detectron2.utils.file_io import PathManager18 19from .evaluator import DatasetEvaluator20 21logger = logging.getLogger(__name__)22 23 24class COCOPanopticEvaluator(DatasetEvaluator):25    """26    Evaluate Panoptic Quality metrics on COCO using PanopticAPI.27    It saves panoptic segmentation prediction in `output_dir`28 29    It contains a synchronize call and has to be called from all workers.30    """31 32    def __init__(self, dataset_name: str, output_dir: Optional[str] = None):33        """34        Args:35            dataset_name: name of the dataset36            output_dir: output directory to save results for evaluation.37        """38        self._metadata = MetadataCatalog.get(dataset_name)39        self._thing_contiguous_id_to_dataset_id = {40            v: k for k, v in self._metadata.thing_dataset_id_to_contiguous_id.items()41        }42        self._stuff_contiguous_id_to_dataset_id = {43            v: k for k, v in self._metadata.stuff_dataset_id_to_contiguous_id.items()44        }45 46        self._output_dir = output_dir47        if self._output_dir is not None:48            PathManager.mkdirs(self._output_dir)49 50    def reset(self):51        self._predictions = []52 53    def _convert_category_id(self, segment_info):54        isthing = segment_info.pop("isthing", None)55        if isthing is None:56            # the model produces panoptic category id directly. No more conversion needed57            return segment_info58        if isthing is True:59            segment_info["category_id"] = self._thing_contiguous_id_to_dataset_id[60                segment_info["category_id"]61            ]62        else:63            segment_info["category_id"] = self._stuff_contiguous_id_to_dataset_id[64                segment_info["category_id"]65            ]66        return segment_info67 68    def process(self, inputs, outputs):69        from panopticapi.utils import id2rgb70 71        for input, output in zip(inputs, outputs):72            panoptic_img, segments_info = output["panoptic_seg"]73            panoptic_img = panoptic_img.cpu().numpy()74            if segments_info is None:75                # If "segments_info" is None, we assume "panoptic_img" is a76                # H*W int32 image storing the panoptic_id in the format of77                # category_id * label_divisor + instance_id. We reserve -1 for78                # VOID label, and add 1 to panoptic_img since the official79                # evaluation script uses 0 for VOID label.80                label_divisor = self._metadata.label_divisor81                segments_info = []82                for panoptic_label in np.unique(panoptic_img):83                    if panoptic_label == -1:84                        # VOID region.85                        continue86                    pred_class = panoptic_label // label_divisor87                    isthing = (88                        pred_class in self._metadata.thing_dataset_id_to_contiguous_id.values()89                    )90                    segments_info.append(91                        {92                            "id": int(panoptic_label) + 1,93                            "category_id": int(pred_class),94                            "isthing": bool(isthing),95                        }96                    )97                # Official evaluation script uses 0 for VOID label.98                panoptic_img += 199 100            file_name = os.path.basename(input["file_name"])101            file_name_png = os.path.splitext(file_name)[0] + ".png"102            with io.BytesIO() as out:103                Image.fromarray(id2rgb(panoptic_img)).save(out, format="PNG")104                segments_info = [self._convert_category_id(x) for x in segments_info]105                self._predictions.append(106                    {107                        "image_id": input["image_id"],108                        "file_name": file_name_png,109                        "png_string": out.getvalue(),110                        "segments_info": segments_info,111                    }112                )113 114    def evaluate(self):115        comm.synchronize()116 117        self._predictions = comm.gather(self._predictions)118        self._predictions = list(itertools.chain(*self._predictions))119        if not comm.is_main_process():120            return121 122        # PanopticApi requires local files123        gt_json = PathManager.get_local_path(self._metadata.panoptic_json)124        gt_folder = PathManager.get_local_path(self._metadata.panoptic_root)125 126        with tempfile.TemporaryDirectory(prefix="panoptic_eval") as pred_dir:127            logger.info("Writing all panoptic predictions to {} ...".format(pred_dir))128            for p in self._predictions:129                with open(os.path.join(pred_dir, p["file_name"]), "wb") as f:130                    f.write(p.pop("png_string"))131 132            with open(gt_json, "r") as f:133                json_data = json.load(f)134            json_data["annotations"] = self._predictions135 136            output_dir = self._output_dir or pred_dir137            predictions_json = os.path.join(output_dir, "predictions.json")138            with PathManager.open(predictions_json, "w") as f:139                f.write(json.dumps(json_data))140 141            from panopticapi.evaluation import pq_compute142 143            with contextlib.redirect_stdout(io.StringIO()):144                pq_res = pq_compute(145                    gt_json,146                    PathManager.get_local_path(predictions_json),147                    gt_folder=gt_folder,148                    pred_folder=pred_dir,149                )150 151        res = {}152        res["PQ"] = 100 * pq_res["All"]["pq"]153        res["SQ"] = 100 * pq_res["All"]["sq"]154        res["RQ"] = 100 * pq_res["All"]["rq"]155        res["PQ_th"] = 100 * pq_res["Things"]["pq"]156        res["SQ_th"] = 100 * pq_res["Things"]["sq"]157        res["RQ_th"] = 100 * pq_res["Things"]["rq"]158        res["PQ_st"] = 100 * pq_res["Stuff"]["pq"]159        res["SQ_st"] = 100 * pq_res["Stuff"]["sq"]160        res["RQ_st"] = 100 * pq_res["Stuff"]["rq"]161 162        results = OrderedDict({"panoptic_seg": res})163        _print_panoptic_results(pq_res)164 165        return results166 167 168def _print_panoptic_results(pq_res):169    headers = ["", "PQ", "SQ", "RQ", "#categories"]170    data = []171    for name in ["All", "Things", "Stuff"]:172        row = [name] + [pq_res[name][k] * 100 for k in ["pq", "sq", "rq"]] + [pq_res[name]["n"]]173        data.append(row)174    table = tabulate(175        data, headers=headers, tablefmt="pipe", floatfmt=".3f", stralign="center", numalign="center"176    )177    logger.info("Panoptic Evaluation Results:\n" + table)178 179 180if __name__ == "__main__":181    from annotator.oneformer.detectron2.utils.logger import setup_logger182 183    logger = setup_logger()184    import argparse185 186    parser = argparse.ArgumentParser()187    parser.add_argument("--gt-json")188    parser.add_argument("--gt-dir")189    parser.add_argument("--pred-json")190    parser.add_argument("--pred-dir")191    args = parser.parse_args()192 193    from panopticapi.evaluation import pq_compute194 195    with contextlib.redirect_stdout(io.StringIO()):196        pq_res = pq_compute(197            args.gt_json, args.pred_json, gt_folder=args.gt_dir, pred_folder=args.pred_dir198        )199        _print_panoptic_results(pq_res)200