Spaces:
Running
Running
File size: 8,725 Bytes
a06fad0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 |
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) |