[feat] visualization support for NN part 2

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

@@ -52,6 +52,11 @@
 from compressai_vision.registry import register_evaluator
 from compressai_vision.utils import time_measure, to_cpu
 
+from detectron2.utils.visualizer import Visualizer
+from detectron2.data import MetadataCatalog
+import cv2
+import os
+
 from .base_evaluator import BaseEvaluator
 from .tf_evaluation_utils import (
     DetectionResultFields,
@@ -94,6 +99,27 @@ 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 +285,28 @@ 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 +473,42 @@ 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., ids2=None):
+        im = np.ascontiguousarray(np.copy(image))
+        im_h, im_w = im.shape[:2]
+
+        text_scale = max(1, image.shape[1] / 1600.)
+        text_thickness = 1 if text_scale > 1.1 else 1
+        line_thickness = max(1, int(image.shape[1] / 500.))
+
+        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

@@ -35,7 +35,7 @@
 from pathlib import Path
 from typing import Callable, Dict
 from uuid import uuid4 as uuid
-
+from omegaconf.errors import InterpolationResolutionError 
 import torch
 import torch.nn as nn
 from torch import Tensor
@@ -81,6 +81,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/remote_inference/image_remote_inference.py

@@ -155,6 +155,9 @@ 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,9 @@ 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,8 @@ 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

@@ -35,7 +35,7 @@
 from torch import Tensor
 from torch.utils.data import DataLoader
 from tqdm import tqdm
-
+from itertools import repeat
 from compressai_vision.evaluators import BaseEvaluator
 from compressai_vision.model_wrappers import BaseWrapper
 from compressai_vision.registry import register_pipeline
@@ -277,7 +277,12 @@ 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 +307,8 @@ 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"])