import itertools import os from typing import List, Optional import torch import numpy as np import tempfile from collections import OrderedDict from PIL import Image from tabulate import tabulate import json import contextlib import detectron2.utils.comm as comm from detectron2.utils.file_io import PathManager from detectron2.data import MetadataCatalog from detectron2.evaluation import COCOPanopticEvaluator from detectron2.utils.visualizer import ColorMode, Visualizer from detectron2.data import MetadataCatalog import io import math from PIL import Image from detectron2.solver.lr_scheduler import _get_warmup_factor_at_iter import logging logger = logging.getLogger(__name__) class TF2WarmupPolyLR(torch.optim.lr_scheduler._LRScheduler): """ Poly learning rate schedule used in TF DeepLab2. Reference: https://github.com/google-research/deeplab2/blob/main/trainer/trainer_utils.py#L23 """ def __init__( self, optimizer: torch.optim.Optimizer, max_iters: int, warmup_factor: float = 0.001, warmup_iters: int = 1000, warmup_method: str = "linear", last_epoch: int = -1, power: float = 0.9, constant_ending: float = 0.0, ): self.max_iters = max_iters self.warmup_factor = warmup_factor self.warmup_iters = warmup_iters self.warmup_method = warmup_method self.power = power self.constant_ending = constant_ending super().__init__(optimizer, last_epoch) def get_lr(self) -> List[float]: warmup_factor = _get_warmup_factor_at_iter( self.warmup_method, self.last_epoch, self.warmup_iters, self.warmup_factor ) if self.constant_ending > 0 and warmup_factor == 1.0: # Constant ending lr. if ( math.pow((1.0 - self.last_epoch / self.max_iters), self.power) < self.constant_ending ): return [base_lr * self.constant_ending for base_lr in self.base_lrs] if self.last_epoch < self.warmup_iters: return [ base_lr * warmup_factor for base_lr in self.base_lrs ] else: return [ base_lr * math.pow((1.0 - self.last_epoch / self.max_iters), self.power) for base_lr in self.base_lrs ] def _compute_values(self) -> List[float]: # The new interface return self.get_lr() class COCOPanopticEvaluatorwithVis(COCOPanopticEvaluator): """ COCO Panoptic Evaluator that supports saving visualizations. TODO(qihangyu): Note that original implementation will also write all predictions to a tmp folder and then run official evaluation script, we may also check how to copy from the tmp folder for visualization. """ def __init__(self, dataset_name: str, output_dir: Optional[str] = None, save_vis_num=0): super().__init__(dataset_name=dataset_name, output_dir=output_dir) self.metadata = MetadataCatalog.get("coco_2017_val_panoptic_with_sem_seg") self.output_dir = output_dir self.save_vis_num = save_vis_num def process(self, inputs, outputs): from panopticapi.utils import id2rgb cur_save_num = 0 for input, output in zip(inputs, outputs): panoptic_img, segments_info = output["panoptic_seg"] panoptic_seg = panoptic_img.cpu() panoptic_img = panoptic_seg.numpy() file_name = os.path.basename(input["file_name"]) file_name_png = os.path.splitext(file_name)[0] + ".png" if cur_save_num < self.save_vis_num: image = output["original_image"] image = image.permute(1, 2 ,0).cpu().numpy()#[:, :, ::-1] visualizer = Visualizer(image, self.metadata, instance_mode=ColorMode.IMAGE) vis_output = visualizer.draw_panoptic_seg_predictions( panoptic_seg, segments_info ) if not os.path.exists(os.path.join(self.output_dir, 'vis')): os.makedirs(os.path.join(self.output_dir, 'vis')) out_filename = os.path.join(self.output_dir, 'vis', file_name_png) vis_output.save(out_filename) cur_save_num += 1 if segments_info is None: # If "segments_info" is None, we assume "panoptic_img" is a # H*W int32 image storing the panoptic_id in the format of # category_id * label_divisor + instance_id. We reserve -1 for # VOID label, and add 1 to panoptic_img since the official # evaluation script uses 0 for VOID label. label_divisor = self._metadata.label_divisor segments_info = [] for panoptic_label in np.unique(panoptic_img): if panoptic_label == -1: # VOID region. continue pred_class = panoptic_label // label_divisor isthing = ( pred_class in self._metadata.thing_dataset_id_to_contiguous_id.values() ) segments_info.append( { "id": int(panoptic_label) + 1, "category_id": int(pred_class), "isthing": bool(isthing), } ) # Official evaluation script uses 0 for VOID label. panoptic_img += 1 with io.BytesIO() as out: Image.fromarray(id2rgb(panoptic_img)).save(out, format="PNG") segments_info = [self._convert_category_id(x) for x in segments_info] self._predictions.append( { "image_id": input["image_id"], "file_name": file_name_png, "png_string": out.getvalue(), "segments_info": segments_info, } ) def evaluate(self): comm.synchronize() self._predictions = comm.gather(self._predictions) self._predictions = list(itertools.chain(*self._predictions)) if not comm.is_main_process(): return # PanopticApi requires local files gt_json = PathManager.get_local_path(self._metadata.panoptic_json) gt_folder = PathManager.get_local_path(self._metadata.panoptic_root) with tempfile.TemporaryDirectory(prefix="panoptic_eval") as pred_dir: logger.info("Writing all panoptic predictions to {} ...".format(pred_dir)) for p in self._predictions: with open(os.path.join(pred_dir, p["file_name"]), "wb") as f: f.write(p.pop("png_string")) with open(gt_json, "r") as f: json_data = json.load(f) json_data["annotations"] = self._predictions output_dir = self._output_dir or pred_dir predictions_json = os.path.join(output_dir, "predictions.json") with PathManager.open(predictions_json, "w") as f: f.write(json.dumps(json_data)) from kmax_deeplab.evaluation.panoptic_evaluation import pq_compute with contextlib.redirect_stdout(io.StringIO()): pq_res = pq_compute( gt_json, PathManager.get_local_path(predictions_json), gt_folder=gt_folder, pred_folder=pred_dir, ) res = {} res["PQ"] = 100 * pq_res["All"]["pq"] res["SQ"] = 100 * pq_res["All"]["sq"] res["RQ"] = 100 * pq_res["All"]["rq"] res["PQ_th"] = 100 * pq_res["Things"]["pq"] res["SQ_th"] = 100 * pq_res["Things"]["sq"] res["RQ_th"] = 100 * pq_res["Things"]["rq"] res["PQ_st"] = 100 * pq_res["Stuff"]["pq"] res["SQ_st"] = 100 * pq_res["Stuff"]["sq"] res["RQ_st"] = 100 * pq_res["Stuff"]["rq"] results = OrderedDict({"panoptic_seg": res}) _print_panoptic_results(pq_res) return results def _print_panoptic_results(pq_res): headers = ["", "PQ", "SQ", "RQ", "#categories"] data = [] for name in ["All", "Things", "Stuff"]: row = [name] + [pq_res[name][k] * 100 for k in ["pq", "sq", "rq"]] + [pq_res[name]["n"]] data.append(row) table = tabulate( data, headers=headers, tablefmt="pipe", floatfmt=".3f", stralign="center", numalign="center" ) logger.info("Panoptic Evaluation Results:\n" + table)