[feat] visualization support for NN part 2 (#16)

This is the implementation of visualization support for NN Part 2. The functions are adapted from Detectron2 and JDE. To enable the function, simply add ++codec.save_visualization=True to the evaluation scripts. To disable it, just keep the scripts unchanged.

Author: jianwensong

cfgs/pipeline/remote_inference.yaml

@@ -33,3 +33,7 @@ evaluation:
     bypass: False
     dump: True
     evaluation_dir: "${codec.output_dir}/evaluation"
+visualization:
+    save_visualization: "${codec.save_visualization}"
+    visualization_dir: "${codec.output_dir}/visualization"
+    threshold: 0  # only for detectron2, 0 means default setting of detectron2

cfgs/pipeline/split_inference.yaml

@@ -44,3 +44,7 @@ evaluation:
     bypass: False
     dump: True
     evaluation_dir: "${codec.output_dir}/evaluation"
+visualization:
+    save_visualization: "${codec.save_visualization}"
+    visualization_dir: "${codec.output_dir}/visualization"
+    threshold: 0  # only for detectron2, 0 means default setting of detectron2

compressai_vision/evaluators/evaluators.py

@@ -29,15 +29,19 @@
 
 import json
 import math
+import os
 from collections import defaultdict
 from pathlib import Path
 from typing import Optional
 
+import cv2
 import motmetrics as mm
 import numpy as np
 import pandas as pd
 import torch
+from detectron2.data import MetadataCatalog
 from detectron2.evaluation import COCOEvaluator
+from detectron2.utils.visualizer import Visualizer
 from jde.utils.io import unzip_objs
 from mmpose.datasets.datasets import BaseCocoStyleDataset
 from mmpose.datasets.transforms import PackPoseInputs
@@ -95,6 +99,29 @@ def reset(self):
     def digest(self, gt, pred):
         return self._evaluator.process(gt, pred)
 
+    def save_visualization(self, gt, pred, output_dir, threshold):
+        gt_image = gt[0]["image"]
+        if torch.is_floating_point(gt_image):
+            gt_image = (gt_image * 255).clamp(0, 255).to(torch.uint8)
+            gt_image = gt_image[[2, 1, 0], ...]
+        gt_image = gt_image.permute(1, 2, 0).cpu().numpy()
+        gt_image = cv2.resize(gt_image, (gt[0]["width"], gt[0]["height"]))
+
+        img_id = gt[0]["image_id"]
+        metadata = MetadataCatalog.get(self.dataset_name)
+        instances = pred[0]["instances"].to("cpu")
+        if threshold:
+            keep = instances.scores >= threshold
+            instances = instances[keep]
+
+        v = Visualizer(gt_image[:, :, ::-1], metadata, scale=1)
+        out = v.draw_instance_predictions(
+            instances
+        )  # selected_instances for specific class
+        output_path = os.path.join(output_dir, f"{img_id}.jpg")
+        cv2.imwrite(output_path, out.get_image()[:, :, ::-1])
+        return
+
     def results(self, save_path: str = None):
         out = self._evaluator.evaluate()
 
@@ -259,6 +286,30 @@ def digest(self, gt, pred):
 
         return
 
+    def save_visualization(self, gt, pred, output_dir, threshold):
+        gt_image = gt[0]["image"]
+        if torch.is_floating_point(gt_image):
+            gt_image = (gt_image * 255).clamp(0, 255).to(torch.uint8)
+            gt_image = gt_image[[2, 1, 0], ...]
+        gt_image = gt_image.permute(1, 2, 0).cpu().numpy()
+        gt_image = cv2.resize(gt_image, (gt[0]["width"], gt[0]["height"]))
+
+        img_id = gt[0]["image_id"]
+        metadata = MetadataCatalog.get(self.dataset_name)
+        instances = pred[0]["instances"].to("cpu")
+
+        if threshold:
+            keep = instances.scores >= threshold
+            instances = instances[keep]
+
+        v = Visualizer(gt_image[:, :, ::-1], metadata, scale=1)
+        out = v.draw_instance_predictions(
+            instances
+        )  # selected_instances for specific class
+        output_path = os.path.join(output_dir, f"{img_id}.jpg")
+        cv2.imwrite(output_path, out.get_image()[:, :, ::-1])
+        return
+
     def _process_prediction(self, pred_dict):
         valid_cls = []
         valid_scores = []
@@ -425,6 +476,57 @@ def digest(self, gt, pred):
 
         self._predictions[int(gt[0]["image_id"])] = pred_list
 
+    def save_visualization(self, gt, pred, output_dir, threshold):
+        image_id = gt[0]["image_id"]
+        gt_image = gt[0]["image"].permute(1, 2, 0).cpu().numpy()
+        gt_image = (gt_image * 255).astype(np.uint8)
+        gt_image = cv2.resize(
+            cv2.cvtColor(gt_image, cv2.COLOR_RGB2BGR), (gt[0]["width"], gt[0]["height"])
+        )
+        online_im = self.plot_tracking(
+            gt_image, pred["tlwhs"], pred["ids"], frame_id=image_id
+        )
+        output_path = os.path.join(output_dir, f"{image_id}.png")
+        cv2.imwrite(output_path, online_im)
+        return
+
+    def plot_tracking(
+        self, image, tlwhs, obj_ids, scores=None, frame_id=0, fps=0.0, ids2=None
+    ):
+        im = np.ascontiguousarray(np.copy(image))
+        im_h, im_w = im.shape[:2]
+
+        text_scale = max(1, image.shape[1] / 1600.0)
+        text_thickness = 1 if text_scale > 1.1 else 1
+        line_thickness = max(1, int(image.shape[1] / 500.0))
+
+        for i, tlwh in enumerate(tlwhs):
+            x1, y1, w, h = tlwh
+            intbox = tuple(map(int, (x1, y1, x1 + w, y1 + h)))
+            obj_id = int(obj_ids[i])
+            id_text = "{}".format(int(obj_id))
+            if ids2 is not None:
+                id_text = id_text + ", {}".format(int(ids2[i]))
+            color = self.get_color(abs(obj_id))
+            cv2.rectangle(
+                im, intbox[0:2], intbox[2:4], color=color, thickness=line_thickness
+            )
+            cv2.putText(
+                im,
+                id_text,
+                (intbox[0], intbox[1] + 30),
+                cv2.FONT_HERSHEY_PLAIN,
+                text_scale,
+                (0, 0, 255),
+                thickness=text_thickness,
+            )
+        return im
+
+    def get_color(self, idx):
+        idx = idx * 3
+        color = ((37 * idx) % 255, (17 * idx) % 255, (29 * idx) % 255)
+        return color
+
     def results(self, save_path: str = None):
         out = self.mot_eval()
 

compressai_vision/pipelines/base.py

@@ -38,6 +38,7 @@
 
 import torch
 import torch.nn as nn
+from omegaconf.errors import InterpolationResolutionError
 from torch import Tensor
 
 from compressai_vision.codecs.utils import (
@@ -81,6 +82,15 @@ def __init__(
         self.bitstream_name = self.configs["codec"]["bitstream_name"]
         self._output_ext = ".h5"
 
+        try:
+            vis_flag = self.configs["visualization"].save_visualization
+        except InterpolationResolutionError:
+            vis_flag = False
+        if vis_flag:
+            self.vis_dir = self.configs["visualization"].visualization_dir
+            self.vis_threshold = self.configs["visualization"].get("threshold", None)
+            self._create_folder(self.vis_dir)
+
         self.codec_output_dir = Path(self.configs["codec"]["codec_output_dir"])
         self.is_mac_calculation = self.configs["codec"]["measure_complexity"]
         self._create_folder(self.codec_output_dir)

compressai_vision/pipelines/fo_vcm/conversion/detectron2.py

@@ -27,8 +27,7 @@
 # OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 # ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-"""From 51 dataset into Detectron2-compatible dataset
-"""
+"""From 51 dataset into Detectron2-compatible dataset"""
 from math import floor
 
 # import cv2

compressai_vision/pipelines/fo_vcm/patch.py

@@ -27,8 +27,7 @@
 # OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 # ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-"""monkey-patching for https://github.com/voxel51/fiftyone/issues/2096
-"""
+"""monkey-patching for https://github.com/voxel51/fiftyone/issues/2096"""
 import csv
 
 # import importhook # this module simply ...ks up everything (at least torch imports)

compressai_vision/pipelines/remote_inference/image_remote_inference.py

@@ -155,6 +155,11 @@ def __call__(
             end = time_measure()
             timing["nn_task"].append((end - start))
 
+            if getattr(self, "vis_dir", None) and hasattr(
+                evaluator, "save_visualization"
+            ):
+                evaluator.save_visualization(d, pred, self.vis_dir, self.vis_threshold)
+
             evaluator.digest(d, pred)
 
             out_res = d[0].copy()

compressai_vision/pipelines/remote_inference/video_remote_inference.py

@@ -184,6 +184,11 @@ def __call__(
             end = time_measure()
             timing["nn_task"].append((end - start))
 
+            if getattr(self, "vis_dir", None) and hasattr(
+                evaluator, "save_visualization"
+            ):
+                evaluator.save_visualization(d, pred, self.vis_dir, self.vis_threshold)
+
             evaluator.digest(d, pred)
 
             out_res = d[0].copy()

compressai_vision/pipelines/split_inference/image_split_inference.py

@@ -211,6 +211,12 @@ def __call__(
 
             if evaluator:
                 evaluator.digest(d, pred)
+                if getattr(self, "vis_dir", None) and hasattr(
+                    evaluator, "save_visualization"
+                ):
+                    evaluator.save_visualization(
+                        d, pred, self.vis_dir, self.vis_threshold
+                    )
 
             out_res = d[0].copy()
             del (

compressai_vision/pipelines/split_inference/video_split_inference.py

@@ -29,6 +29,7 @@
 
 
 import os
+from itertools import repeat
 from typing import Dict, List, Tuple, TypeVar
 
 import torch
@@ -277,7 +278,14 @@ def __call__(
         self.logger.info("Processing NN-Part2...")
         output_list = []
 
-        for e, ftensors in enumerate(tqdm(dec_ftensors_list)):
+        if getattr(self, "vis_dir", None):
+            dec_ftensors_list = zip(dec_ftensors_list, dataloader)
+        else:
+            dec_ftensors_list = zip(dec_ftensors_list, repeat(None))
+
+        for e, (ftensors, d) in enumerate(
+            tqdm(dec_ftensors_list, total=len(dataloader))
+        ):
             data = {k: v.to(self.device_nn_part2) for k, v in ftensors.items()}
             dec_features["data"] = data
             dec_features["file_name"] = file_names[e]
@@ -302,7 +310,12 @@ def __call__(
 
             if evaluator:
                 evaluator.digest(gt_inputs[e], pred)
-
+                if getattr(self, "vis_dir", None) and hasattr(
+                    evaluator, "save_visualization"
+                ):
+                    evaluator.save_visualization(
+                        d, pred, self.vis_dir, self.vis_threshold
+                    )
             out_res = dec_features.copy()
             del (out_res["data"], out_res["org_input_size"])
 
@@ -343,7 +356,7 @@ def __call__(
 
     @staticmethod
     def _feature_tensor_list_to_dict(
-        data: List[Dict[str, Tensor]]
+        data: List[Dict[str, Tensor]],
     ) -> Dict[str, Tensor]:
         """
         Converts a list of feature tensors into a dictionary format.