from mmengine.fileio import FileClient
from mmengine.dist import master_only
from einops import rearrange
import torch
import mmcv
import numpy as np
import os.path as osp
import cv2
from typing import Optional, Sequence
import torch.nn as nn
from mmdet.apis import inference_detector
from mmcv.transforms import Compose
from mmdet.engine import DetVisualizationHook
from mmdet.registry import HOOKS
from mmdet.structures import DetDataSample
from utils.io_utils import find_all_imgs, square_pad_resize, imglist2grid
def inference_detector(
model: nn.Module,
imgs,
test_pipeline
):
if isinstance(imgs, (list, tuple)):
is_batch = True
else:
imgs = [imgs]
is_batch = False
if len(imgs) == 0:
return []
test_pipeline = test_pipeline.copy()
if isinstance(imgs[0], np.ndarray):
# Calling this method across libraries will result
# in module unregistered error if not prefixed with mmdet.
test_pipeline[0].type = 'mmdet.LoadImageFromNDArray'
test_pipeline = Compose(test_pipeline)
result_list = []
for img in imgs:
# prepare data
if isinstance(img, np.ndarray):
# TODO: remove img_id.
data_ = dict(img=img, img_id=0)
else:
# TODO: remove img_id.
data_ = dict(img_path=img, img_id=0)
# build the data pipeline
data_ = test_pipeline(data_)
data_['inputs'] = [data_['inputs']]
data_['data_samples'] = [data_['data_samples']]
# forward the model
with torch.no_grad():
results = model.test_step(data_)[0]
result_list.append(results)
if not is_batch:
return result_list[0]
else:
return result_list
@HOOKS.register_module()
class InstanceSegVisualizationHook(DetVisualizationHook):
def __init__(self, visualize_samples: str = '',
read_rgb: bool = False,
draw: bool = False,
interval: int = 50,
score_thr: float = 0.3,
show: bool = False,
wait_time: float = 0.,
test_out_dir: Optional[str] = None,
file_client_args: dict = dict(backend='disk')):
super().__init__(draw, interval, score_thr, show, wait_time, test_out_dir, file_client_args)
self.vis_samples = []
if osp.exists(visualize_samples):
self.channel_order = channel_order = 'rgb' if read_rgb else 'bgr'
samples = find_all_imgs(visualize_samples, abs_path=True)
for imgp in samples:
img = mmcv.imread(imgp, channel_order=channel_order)
img, _, _, _ = square_pad_resize(img, 640)
self.vis_samples.append(img)
def before_val(self, runner) -> None:
total_curr_iter = runner.iter
self._visualize_data(total_curr_iter, runner)
return super().before_val(runner)
# def after_val_iter(self, runner: Runner, batch_idx: int, data_batch: dict,
# outputs: Sequence[DetDataSample]) -> None:
# """Run after every ``self.interval`` validation iterations.
# Args:
# runner (:obj:`Runner`): The runner of the validation process.
# batch_idx (int): The index of the current batch in the val loop.
# data_batch (dict): Data from dataloader.
# outputs (Sequence[:obj:`DetDataSample`]]): A batch of data samples
# that contain annotations and predictions.
# """
# # if self.draw is False:
# # return
# if self.file_client is None:
# self.file_client = FileClient(**self.file_client_args)
# # There is no guarantee that the same batch of images
# # is visualized for each evaluation.
# total_curr_iter = runner.iter + batch_idx
# # # Visualize only the first data
# # img_path = outputs[0].img_path
# # img_bytes = self.file_client.get(img_path)
# # img = mmcv.imfrombytes(img_bytes, channel_order='rgb')
# if total_curr_iter % self.interval == 0 and self.vis_samples:
# self._visualize_data(total_curr_iter, runner)
@master_only
def _visualize_data(self, total_curr_iter, runner):
tgt_size = 384
runner.model.eval()
outputs = inference_detector(runner.model, self.vis_samples, test_pipeline=runner.cfg.test_pipeline)
vis_results = []
for img, output in zip(self.vis_samples, outputs):
vis_img = self.add_datasample(
'val_img',
img,
data_sample=output,
show=self.show,
wait_time=self.wait_time,
pred_score_thr=self.score_thr,
draw_gt=False,
step=total_curr_iter)
vis_results.append(cv2.resize(vis_img, (tgt_size, tgt_size), interpolation=cv2.INTER_AREA))
drawn_img = imglist2grid(vis_results, tgt_size)
if drawn_img is None:
return
drawn_img = cv2.cvtColor(drawn_img, cv2.COLOR_BGR2RGB)
visualizer = self._visualizer
visualizer.set_image(drawn_img)
visualizer.add_image('val_img', drawn_img, total_curr_iter)
@master_only
def add_datasample(
self,
name: str,
image: np.ndarray,
data_sample: Optional['DetDataSample'] = None,
draw_gt: bool = True,
draw_pred: bool = True,
show: bool = False,
wait_time: float = 0,
# TODO: Supported in mmengine's Viusalizer.
out_file: Optional[str] = None,
pred_score_thr: float = 0.3,
step: int = 0) -> np.ndarray:
image = image.clip(0, 255).astype(np.uint8)
visualizer = self._visualizer
classes = visualizer.dataset_meta.get('classes', None)
palette = visualizer.dataset_meta.get('palette', None)
gt_img_data = None
pred_img_data = None
if data_sample is not None:
data_sample = data_sample.cpu()
if draw_gt and data_sample is not None:
gt_img_data = image
if 'gt_instances' in data_sample:
gt_img_data = visualizer._draw_instances(image,
data_sample.gt_instances,
classes, palette)
if 'gt_panoptic_seg' in data_sample:
assert classes is not None, 'class information is ' \
'not provided when ' \
'visualizing panoptic ' \
'segmentation results.'
gt_img_data = visualizer._draw_panoptic_seg(
gt_img_data, data_sample.gt_panoptic_seg, classes)
if draw_pred and data_sample is not None:
pred_img_data = image
if 'pred_instances' in data_sample:
pred_instances = data_sample.pred_instances
pred_instances = pred_instances[
pred_instances.scores > pred_score_thr]
pred_img_data = visualizer._draw_instances(image, pred_instances,
classes, palette)
if 'pred_panoptic_seg' in data_sample:
assert classes is not None, 'class information is ' \
'not provided when ' \
'visualizing panoptic ' \
'segmentation results.'
pred_img_data = visualizer._draw_panoptic_seg(
pred_img_data, data_sample.pred_panoptic_seg.numpy(),
classes)
if gt_img_data is not None and pred_img_data is not None:
drawn_img = np.concatenate((gt_img_data, pred_img_data), axis=1)
elif gt_img_data is not None:
drawn_img = gt_img_data
elif pred_img_data is not None:
drawn_img = pred_img_data
else:
# Display the original image directly if nothing is drawn.
drawn_img = image
return drawn_img