train_net_utils.py

"""
Copyright (2023) Bytedance Ltd. and/or its affiliates

Licensed under the Apache License, Version 2.0 (the "License"); 
you may not use this file except in compliance with the License. 
You may obtain a copy of the License at 

    http://www.apache.org/licenses/LICENSE-2.0 

Unless required by applicable law or agreed to in writing, software 
distributed under the License is distributed on an "AS IS" BASIS, 
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
See the License for the specific language governing permissions and 
limitations under the License. 
"""
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)