def test_imdenormalize(self):
norm_img = (self.img[:, :, ::-1] - self.mean) / self.std
rgb_baseline = (norm_img * self.std + self.mean)
bgr_baseline = rgb_baseline[:, :, ::-1]
img = mmcv.imdenormalize(norm_img, self.mean, self.std)
assert np.allclose(img, bgr_baseline)
img = mmcv.imdenormalize(norm_img, self.mean, self.std, to_bgr=False)
assert np.allclose(img, rgb_baseline)
def tensor2imgs(tensor, mean=(0, 0, 0), std=(1, 1, 1), to_rgb=True):
"""Convert tensor to 3-channel images.
Args:
tensor (torch.Tensor): Tensor that contains multiple images, shape (
N, C, H, W).
mean (tuple[float], optional): Mean of images. Defaults to (0, 0, 0).
std (tuple[float], optional): Standard deviation of images.
Defaults to (1, 1, 1).
to_rgb (bool, optional): Whether the tensor was converted to RGB
format in the first place. If so, convert it back to BGR.
Defaults to True.
Returns:
list[np.ndarray]: A list that contains multiple images.
"""
if torch is None:
raise RuntimeError('pytorch is not installed')
assert torch.is_tensor(tensor) and tensor.ndim == 4
assert len(mean) == 3
assert len(std) == 3
num_imgs = tensor.size(0)
mean = np.array(mean, dtype=np.float32)
std = np.array(std, dtype=np.float32)
imgs = []
for img_id in range(num_imgs):
img = tensor[img_id, ...].cpu().numpy().transpose(1, 2, 0)
img = mmcv.imdenormalize(
img, mean, std, to_bgr=to_rgb).astype(np.uint8)
imgs.append(np.ascontiguousarray(img))
return imgs
def tensor2imgs(tensor, mean=(0, 0, 0), std=(1, 1, 1), to_rgb=True):
"""Convert tensor to images.
Args:
tensor (torch.Tensor): Tensor that contains multiple images
mean (tuple[float], optional): Mean of images. Defaults to (0, 0, 0).
std (tuple[float], optional): Standard deviation of images.
Defaults to (1, 1, 1).
to_rgb (bool, optional): Whether convert the images to RGB format.
Defaults to True.
Returns:
list[np.ndarray]: A list that contains multiple images.
"""
num_imgs = tensor.size(0)
mean = np.array(mean, dtype=np.float32)
std = np.array(std, dtype=np.float32)
imgs = []
for img_id in range(num_imgs):
img = tensor[img_id, ...].cpu().numpy().transpose(1, 2, 0)
img = mmcv.imdenormalize(img, mean, std,
to_bgr=to_rgb).astype(np.uint8)
if len(img.shape) == 2:
img = img[:, :, np.newaxis]
imgs.append(np.ascontiguousarray(img))
return imgs
def tensor2grayimgs(tensor, mean=(127, ), std=(127, ), **kwargs):
"""Convert tensor to 1-channel gray images.
Args:
tensor (torch.Tensor): Tensor that contains multiple images, shape (
N, C, H, W).
mean (tuple[float], optional): Mean of images. Defaults to (127).
std (tuple[float], optional): Standard deviation of images.
Defaults to (127).
Returns:
list[np.ndarray]: A list that contains multiple images.
"""
assert torch.is_tensor(tensor) and tensor.ndim == 4
assert tensor.size(1) == len(mean) == len(std) == 1
num_imgs = tensor.size(0)
mean = np.array(mean, dtype=np.float32)
std = np.array(std, dtype=np.float32)
imgs = []
for img_id in range(num_imgs):
img = tensor[img_id, ...].cpu().numpy().transpose(1, 2, 0)
img = mmcv.imdenormalize(img, mean, std, to_bgr=False).astype(np.uint8)
imgs.append(np.ascontiguousarray(img))
return imgs
def backward_D(self):
fake = round(random.uniform(0.0, 0.1), 1)
real = round(random.uniform(0.9, 1.0), 1)
# fake = 0.0
# real = 1.0
# label
threashold = random.uniform(0.0, 1.0)
if threashold < 0.05:
r_fake = real
f_real = fake
else:
r_fake = fake
f_real = real
self.set_requires_grad(self.discriminator, True)
pred_fake = self.discriminator(self.p_t_1.detach())
loss_D_fake = self.loss(pred_fake, r_fake)
pred_real = self.discriminator(self.g_t_1)
loss_D_real = self.loss(pred_real, f_real)
loss_D = (loss_D_fake + loss_D_real) * 0.5
if self.key % self.p_s == 0:
self.key = 1
b_s = self.p_t_1.shape[0]
# plt.figure(figsize=(6, 2))
big_img = None
for i in range(0, b_s):
img_g = self.g_t_1[i].squeeze(0).cpu().permute(1, 2, 0).numpy()
img_p = self.p_t_1[i].squeeze(0).detach().cpu().permute(
1, 2, 0).numpy()
img_g = mmcv.imdenormalize(img_g, 127.5, 127.5, to_bgr=False)
img_p = mmcv.imdenormalize(img_p, 127.5, 127.5, to_bgr=False)
img = np.hstack([img_g, img_p])
if big_img is None:
big_img = img
else:
big_img = np.vstack([big_img, img])
# plt.subplot(i + 1, 1)
# plt.plot(img_g)
# plt.subplot(i + 1, 2)
# plt.plot(img_p)
name = str(int(time.time())) + '.jpg'
mmcv.imwrite(big_img, os.path.join('./img', name))
self.key += 1
# plt.savefig()
# plt.show()
return {'loss_D': loss_D}
def tensor2imgs(tensor, mean=(0, 0, 0), std=(1, 1, 1), to_rgb=True):
num_imgs = tensor.size(0)
mean = np.array(mean, dtype=np.float32)
std = np.array(std, dtype=np.float32)
imgs = []
for img_id in range(num_imgs):
img = tensor[img_id, ...].cpu().numpy().transpose(1, 2, 0)
img = mmcv.imdenormalize(
img, mean, std, to_bgr=to_rgb).astype(np.uint8)
imgs.append(np.ascontiguousarray(img))
return imgs
def tensor2video_snaps(tensor, mean=(0, 0, 0), std=(1, 1, 1), to_rgb=True):
num_videos = tensor.size(0)
num_frames = tensor.size(2)
mean = np.array(mean, dtype=np.float32)
std = np.array(std, dtype=np.float32)
video_snaps = []
for vid_id in range(num_videos):
img = tensor[vid_id, :, num_frames //
2, ...].cpu().numpy().transpose(1, 2, 0)
img = mmcv.imdenormalize(
img, mean, std, to_bgr=to_rgb).astype(np.uint8)
video_snaps.append(np.ascontiguousarray(img))
return video_snaps
def main():
args = parse_args()
os.makedirs(args.output, exist_ok=True)
cfg = Config.fromfile(args.config)
dataset = get_dataset(cfg.data.train)
for i in tqdm(np.random.randint(0, len(dataset), 500)):
data = dataset[i]
img = data['img'].data.numpy().transpose(1, 2, 0)
masks = data['gt_masks'].data.transpose(1, 2, 0).astype(bool)
bboxes = data['gt_bboxes'].data.numpy()
img = mmcv.imdenormalize(img, mean=cfg.img_norm_cfg.mean, std=cfg.img_norm_cfg.std, to_bgr=False)
img = draw_masks(img, masks).astype(np.uint8)
draw_bounding_boxes_on_image_array(img, bboxes, use_normalized_coordinates=False, thickness=5)
cv2.imwrite(osp.join(args.output, f'{i}_{np.random.randint(0, 10000)}.jpg'), img[..., ::-1])
def tensor2img3D(tensor,
slice_num=85,
mean=(123.675, 116.28, 103.53),
std=(58.395, 57.12, 57.375),
to_rgb=True):
mean = np.array(mean, dtype=np.float32)
std = np.array(std, dtype=np.float32)
# get the first image
imgs = tensor[0, ...].cpu().numpy().transpose(2, 3, 1, 0)
# get the first slice
img = imgs[:, :, slice_num, :]
img = mmcv.imdenormalize(img, mean, std, to_bgr=to_rgb).astype(np.uint8)
img = np.ascontiguousarray(img)
return img
def _imrenormalize(img, img_norm_cfg, new_img_norm_cfg):
"""Re-normalize the image."""
img_norm_cfg = img_norm_cfg.copy()
new_img_norm_cfg = new_img_norm_cfg.copy()
for k, v in img_norm_cfg.items():
if (k == 'mean' or k == 'std') and not isinstance(v, np.ndarray):
img_norm_cfg[k] = np.array(v, dtype=img.dtype)
# reverse cfg
if 'to_rgb' in img_norm_cfg:
img_norm_cfg['to_bgr'] = img_norm_cfg['to_rgb']
img_norm_cfg.pop('to_rgb')
for k, v in new_img_norm_cfg.items():
if (k == 'mean' or k == 'std') and not isinstance(v, np.ndarray):
new_img_norm_cfg[k] = np.array(v, dtype=img.dtype)
img = mmcv.imdenormalize(img, **img_norm_cfg)
img = mmcv.imnormalize(img, **new_img_norm_cfg)
return img
def tensor2img3DNPrint(tensor,
slice_num=85,
bboxes=np.array([[0, 0, 0, 0]]),
mean=(123.675, 116.28, 103.53),
std=(58.395, 57.12, 57.375),
to_rgb=True):
mean = np.array(mean, dtype=np.float32)
std = np.array(std, dtype=np.float32)
# get the first image
imgs = tensor[0, ...].cpu().numpy().transpose(2, 3, 1, 0)
# get the first slice
img = imgs[:, :, slice_num, :]
img = mmcv.imdenormalize(img, mean, std, to_bgr=False)
img = np.ascontiguousarray(img)
mmcv.imshow_bboxes(img, bboxes)
return img
def draw_heatmap(rois, roi_feature, tag, img, img_name):
mean = [123.675, 116.28, 103.53]
std = [58.395, 57.12, 57.375]
img = img.squeeze(0)
img = img.cpu().numpy().transpose(1, 2, 0)
img = mmcv.imdenormalize(img, np.array(mean, dtype=np.float32),
np.array(std, dtype=np.float32), True)
roi_img = []
for roi_id in range(len(rois)):
x1 = int(rois[roi_id, 0])
x2 = int(rois[roi_id, 2])
y1 = int(rois[roi_id, 1])
y2 = int(rois[roi_id, 3])
roi_img.append(img[y1:y2 + 1, x1:x2 + 1, :])
roi_fea = []
for j in range(roi_feature.size(0)):
roi_fea.append(roi_feature[j])
assert len(roi_img) == len(roi_fea)
for i in range(len(roi_img)):
roi = roi_img[i]
if roi.shape[0] > 50:
heatmap = roi_fea[i].detach().cpu().numpy()
heatmap = np.max(heatmap, axis=0)
# heatmap = F.relu(torch.sum(roi_fea[i], dim=0))
# heatmap = heatmap.detach().cpu().numpy()
heatmap = np.maximum(heatmap, 0)
heatmap /= np.max(heatmap)
heatmap = cv2.resize(heatmap, (roi.shape[1], roi.shape[0]))
heatmap = np.uint8(255 * heatmap)
heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
superposed_img = heatmap * 0.4 + roi
# brightness = v_channel*0.4 + roi
superposed_img = cv2.resize(
superposed_img, (2 * heatmap.shape[1], 2 * heatmap.shape[0]))
# v_channel = cv2.resize(v_channel, (2 * heatmap.shape[1], 2 * heatmap.shape[0]))
roi = cv2.resize(roi, (2 * heatmap.shape[1], 2 * heatmap.shape[0]))
save_path = visual_path + img_name + '_{}'.format(
tag[i]) + '_{}.jpg'.format(i)
save_img = visual_path + img_name + '_{}'.format(
tag[i]) + '_{}_original.jpg'.format(i)
cv2.imwrite(save_path, superposed_img)
cv2.imwrite(save_img, roi)
def tensor2imgs(tensor, mean=None, std=None, to_rgb=True):
"""Convert tensor to 3-channel images or 1-channel gray images.
Args:
tensor (torch.Tensor): Tensor that contains multiple images, shape (
N, C, H, W). :math:`C` can be either 3 or 1.
mean (tuple[float], optional): Mean of images. If None,
(0, 0, 0) will be used for tensor with 3-channel,
while (0, ) for tensor with 1-channel. Defaults to None.
std (tuple[float], optional): Standard deviation of images. If None,
(1, 1, 1) will be used for tensor with 3-channel,
while (1, ) for tensor with 1-channel. Defaults to None.
to_rgb (bool, optional): Whether the tensor was converted to RGB
format in the first place. If so, convert it back to BGR.
For the tensor with 1 channel, it must be False. Defaults to True.
Returns:
list[np.ndarray]: A list that contains multiple images.
"""
if torch is None:
raise RuntimeError('pytorch is not installed')
assert torch.is_tensor(tensor) and tensor.ndim == 4
channels = tensor.size(1)
assert channels in [1, 3]
if mean is None:
mean = (0, ) * channels
if std is None:
std = (1, ) * channels
assert (channels == len(mean) == len(std) == 3) or \
(channels == len(mean) == len(std) == 1 and not to_rgb)
num_imgs = tensor.size(0)
mean = np.array(mean, dtype=np.float32)
std = np.array(std, dtype=np.float32)
imgs = []
for img_id in range(num_imgs):
img = tensor[img_id, ...].cpu().numpy().transpose(1, 2, 0)
img = mmcv.imdenormalize(
img, mean, std, to_bgr=to_rgb).astype(np.uint8)
imgs.append(np.ascontiguousarray(img))
return imgs
def single_test(model, data_loader, cfg):
model.eval()
results = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
with torch.no_grad():
result, seg, img, meta = model(return_loss=False, **data)
#result = model(return_loss=False, **data)
#'''
for frame in range(8):
for idx_in_batch in range(seg.shape[0]):
seg0 = seg.argmax(1)[idx_in_batch, frame, :, :]
img0 = mmcv.imdenormalize(
img[idx_in_batch, :,
frame * 4, :, :].transpose([1, 2, 0]),
mean=np.array(cfg.img_norm_cfg.mean).reshape(1, 1, 3),
std=np.array(cfg.img_norm_cfg.std).reshape(1, 1, 3),
to_bgr=cfg.img_norm_cfg.to_rgb)
out_dir = os.path.join('outputs_ntucentercrop_1028',
meta[0]['img_path'],
'setting_%02d' % idx_in_batch)
if not os.path.isdir(out_dir):
os.makedirs(out_dir)
mmcv.imwrite(
img0, os.path.join(out_dir, 'img_%05d.png' % (frame)))
mmcv.imwrite(
seg0 * 255,
os.path.join(out_dir, 'seg_%05d.png' % (frame)))
#'''
results.append(result)
batch_size = data['img_group_0'].data[0].size(0)
for _ in range(batch_size):
prog_bar.update()
return results
def main():
# base configs
data_root = '/media/' + getpass.getuser(
) + '/Data/DoubleCircle/datasets/kaist-rgbt-encoder/'
# img_norm_cfg = dict(
# mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
# img_norm_cfg_t = dict(
# mean=[123.675, 123.675, 123.675], std=[58.395, 58.395, 58.395], to_rgb=False)
img_norm_cfg = dict(mean=[0, 0, 0], std=[255, 255, 255], to_rgb=True)
img_norm_cfg_t = dict(mean=[0, 0, 0], std=[147, 147, 147], to_rgb=False)
imgs_per_gpu = 16
workers_per_gpu = 2
max_epoch = 50
base_lr = 1e-2
# train and test dataset
train = dict(ann_file=data_root + 'annotations-pkl/train-all.pkl',
img_prefix=data_root + 'images/',
img_scale=1.0,
img_norm_cfg=img_norm_cfg,
img_norm_cfg_t=img_norm_cfg_t,
size_divisor=None,
flip_ratio=0.5,
with_mask=False,
with_crowd=True,
with_label=True)
test = dict(ann_file=data_root + 'annotations-pkl/test-all-rgb.pkl',
img_prefix=data_root + 'images/',
img_scale=1.0,
img_norm_cfg=img_norm_cfg,
img_norm_cfg_t=img_norm_cfg_t,
size_divisor=None,
flip_ratio=0,
with_mask=False,
with_crowd=True,
with_label=True)
dataset_train = CoderKaistDataset(**train)
dataset_test = CoderKaistDataset(**test)
# train and test data loader
data_loaders_train = build_dataloader(dataset_train,
imgs_per_gpu,
workers_per_gpu,
num_gpus=1,
dist=False)
data_loaders_test = build_dataloader(dataset_test,
imgs_per_gpu,
workers_per_gpu,
num_gpus=1,
dist=False)
# MINST dataset
# im_tfs = tfs.Compose([
# tfs.ToTensor(),
# tfs.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) # 标准化
# ])
#
# train_set = MNIST('./mnist', transform=im_tfs, download=True)
# train_data = DataLoader(train_set, batch_size=128, shuffle=True)
# train
net = AutoEncoder()
net.init_weights()
net.cuda()
# loss_fn = torch.nn.MSELoss(size_average=False)
loss_fn = torch.nn.MSELoss(reduction='elementwise_mean')
optimizer = optim.Adam(net.parameters(), lr=base_lr, weight_decay=0.0001)
print('Start training...\n')
# for e in range(max_epoch):
# for im in data_loaders_train:
# if torch.cuda.is_available():
# input = im['img_thermal_in'].cuda()
# input = Variable(input)
# # 前向传播
# code, output = net(input)
# loss = loss_fn(output, input)
# # 反向传播
# optimizer.zero_grad()
# loss.backward()
# optimizer.step()
#
# if (e + 1) % 1 == 0: # 每 1 次，将生成的图片保存一下
# print('epoch: {}, Loss: {:.4f}'.format(e + 1, loss.data))
# output = output.cpu().data
# target = im['img_thermal_in']
# pic = np.zeros((output.shape[0], output.shape[2], output.shape[3], output.shape[1]), dtype=np.uint8)
# target_pic = np.zeros((output.shape[0], output.shape[2], output.shape[3], output.shape[1]),
# dtype=np.uint8)
# mean = np.array(img_norm_cfg['mean'], dtype=np.float32)
# std = np.array(img_norm_cfg['std'], dtype=np.float32)
# for idx in range(imgs_per_gpu):
# img = output[idx, ...].numpy().transpose(1, 2, 0).astype(np.float32)
# pic[idx, :, :, :] = mmcv.imdenormalize(
# img, mean=mean, std=std, to_bgr=False).astype(np.uint8)
# target_img = target[idx, ...].numpy().transpose(1, 2, 0).astype(np.float32)
# target_pic[idx, :, :, :] = mmcv.imdenormalize(
# target_img, mean=mean, std=std, to_bgr=False).astype(np.uint8)
# if not os.path.exists('../../work_dirs/autoencoder'):
# os.mkdir('../../work_dirs/autoencoder')
# save_images(torch.from_numpy(pic.transpose((0, 3, 1, 2))),
# '../../work_dirs/autoencoder/image_{}.png'.format(e + 1))
# save_images(torch.from_numpy(target_pic.transpose(0, 3, 1, 2)),
# '../../work_dirs/autoencoder/target_image_{}.png'.format(e + 1))
# # update learn rate
# adjust_learning_rate(optimizer, base_lr, e)
# # save checkpoint
# filename = '../../work_dirs/autoencoder/epoch_{}.pth'.format(e + 1)
# save_checkpoint(net, filename=filename)
iter_epoch = len(data_loaders_train)
for e in range(max_epoch):
# training phase
net.train()
loss_iter = 0.0
for i, data_batch in enumerate(data_loaders_train):
code, decode_rgb, decode_thermal = net(
data_batch['img_rgb_in'].cuda(),
data_batch['img_thermal_in'].cuda())
data_batch['img_rgb_out'] = data_batch['img_rgb_out'].view(
(-1, 3, 128, 160))
data_batch['img_thermal_out'] = data_batch['img_thermal_out'].view(
(-1, 3, 128, 160))
loss_rgb = loss_fn(decode_rgb.cpu(), data_batch['img_rgb_out'])
loss_thermal = loss_fn(decode_thermal.cpu(),
data_batch['img_thermal_out'])
loss = loss_rgb + loss_thermal
loss_iter += loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i + 1) % 50 == 0:
print(
'Epoch:{}|{},Iteration:[{}|{}],Learning Rate:{},Loss:{:.4f}'
.format(e + 1, max_epoch, i + 1, len(data_loaders_train),
optimizer.param_groups[0]['lr'], loss_iter))
loss_iter = 0.0
# update learn rate
adjust_learning_rate(optimizer, base_lr, e + 1)
# save checkpoint
filename = '../../work_dirs/autoencoder/epoch_{}.pth'.format(e + 1)
save_checkpoint(net, filename=filename)
# evaluation phase
if (e + 1) % 1 == 0: # 每 1 次，将生成的图片保存一下
output_rgb = decode_rgb.cpu().data
target_rgb = data_batch['img_rgb_out']
output_thermal = decode_thermal.cpu().data
tartget_thermal = data_batch['img_thermal_out']
pic_rgb = np.zeros((output_rgb.shape[0], output_rgb.shape[2],
output_rgb.shape[3], output_rgb.shape[1]),
dtype=np.uint8)
target_pic_rgb = np.zeros(
(output_rgb.shape[0], output_rgb.shape[2], output_rgb.shape[3],
output_rgb.shape[1]),
dtype=np.uint8)
pic_thermal = np.zeros((output_rgb.shape[0], output_rgb.shape[2],
output_rgb.shape[3], output_rgb.shape[1]),
dtype=np.uint8)
target_pic_thermal = np.zeros(
(output_rgb.shape[0], output_rgb.shape[2], output_rgb.shape[3],
output_rgb.shape[1]),
dtype=np.uint8)
mean_rgb = np.array(img_norm_cfg['mean'], dtype=np.float32)
std_rgb = np.array(img_norm_cfg['std'], dtype=np.float32)
mean_thermal = np.array(img_norm_cfg_t['mean'], dtype=np.float32)
std_thermal = np.array(img_norm_cfg_t['std'], dtype=np.float32)
for idx in range(output_rgb.shape[0]):
# for rgb
img = output_rgb[idx,
...].numpy().transpose(1, 2,
0).astype(np.float32)
pic_rgb[idx, :, :, :] = mmcv.imdenormalize(
img, mean=mean_rgb, std=std_rgb,
to_bgr=False).astype(np.uint8)
target_img = target_rgb[idx, ...].numpy().transpose(
1, 2, 0).astype(np.float32)
target_pic_rgb[idx, :, :, :] = mmcv.imdenormalize(
target_img, mean=mean_rgb, std=std_rgb,
to_bgr=False).astype(np.uint8)
# for thermal
img_t = output_thermal[idx, ...].numpy().transpose(
1, 2, 0).astype(np.float32)
pic_thermal[idx, :, :, :] = mmcv.imdenormalize(
img_t, mean=mean_thermal, std=std_thermal,
to_bgr=False).astype(np.uint8)
target_img_t = tartget_thermal[idx, ...].numpy().transpose(
1, 2, 0).astype(np.float32)
target_pic_thermal[idx, :, :, :] = mmcv.imdenormalize(
target_img_t,
mean=mean_thermal,
std=std_thermal,
to_bgr=False).astype(np.uint8)
if not os.path.exists('../../work_dirs/autoencoder'):
os.mkdir('../../work_dirs/autoencoder')
save_images(
torch.from_numpy(pic_rgb.transpose((0, 3, 1, 2))),
'../../work_dirs/autoencoder/image_rgb_{}.png'.format(e + 1))
save_images(
torch.from_numpy(target_pic_rgb.transpose(0, 3, 1, 2)),
'../../work_dirs/autoencoder/target_image_rgb_{}.png'.format(
e + 1))
save_images(
torch.from_numpy(pic_thermal.transpose((0, 3, 1, 2))),
'../../work_dirs/autoencoder/image_thermal_{}.png'.format(e +
1))
save_images(
torch.from_numpy(target_pic_thermal.transpose(0, 3, 1, 2)),
'../../work_dirs/autoencoder/target_image_thermal_{}.png'.
format(e + 1))
def main(args):
if args.model.endswith('.onnx'):
backend = 'onnx'
elif args.model.endswith('.xml'):
backend = 'openvino'
else:
raise ValueError('Unknown model type.')
cfg = mmcv.Config.fromfile(args.config)
if args.update_config:
cfg.merge_from_dict(args.update_config)
cfg.model.pretrained = None
cfg.data.test.test_mode = True
if backend == 'openvino':
assert cfg.data.test.pipeline[1]['type'] == 'MultiScaleFlipAug'
normalize_idx = [i for i, v in enumerate(cfg.data.test.pipeline[1]['transforms']) if v['type'] == 'Normalize'][0]
cfg.data.test.pipeline[1]['transforms'][normalize_idx]['mean'] = [0.0, 0.0, 0.0]
cfg.data.test.pipeline[1]['transforms'][normalize_idx]['std'] = [1.0, 1.0, 1.0]
cfg.data.test.pipeline[1]['transforms'][normalize_idx]['to_rgb'] = False
print(cfg.data.test)
if args.video is not None and args.show:
dataset = VideoDataset(int(args.video), cfg.data)
data_loader = iter(dataset)
wait_key = 1
else:
dataset = build_dataset(cfg.data.test)
data_loader = build_dataloader(
dataset,
samples_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=False,
shuffle=False)
wait_key = -1
# Valid classes + background.
classes_num = len(dataset.CLASSES) + 1
if backend == 'openvino':
from mmdet.utils.deployment.openvino_backend import DetectorOpenVINO
model = DetectorOpenVINO(args.model,
args.model[:-3] + 'bin',
mapping_file_path=args.model[:-3] + 'mapping',
cfg=cfg,
classes=dataset.CLASSES)
else:
from mmdet.utils.deployment.onnxruntime_backend import ModelONNXRuntime
model = ModelONNXRuntime(args.model, cfg=cfg, classes=dataset.CLASSES)
results = []
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
im_data = data['img'][0].cpu().numpy()
try:
result = model(im_data)
result = postprocess(
result,
data['img_metas'][0].data[0],
num_classes=classes_num,
rescale=not args.show)
except Exception as ex:
print(f'\nException raised while processing item {i}:')
print(ex)
with_mask = hasattr(model.pt_model, 'with_mask') and model.pt_model.with_mask
result = empty_result(
num_classes=classes_num,
with_mask=with_mask)
results.append(result)
if args.show:
img_meta = data['img_metas'][0].data[0][0]
norm_cfg = img_meta['img_norm_cfg']
mean = np.array(norm_cfg['mean'], dtype=np.float32)
std = np.array(norm_cfg['std'], dtype=np.float32)
display_image = im_data[0].transpose(1, 2, 0)
display_image = mmcv.imdenormalize(display_image, mean, std, to_bgr=norm_cfg['to_rgb']).astype(np.uint8)
display_image = np.ascontiguousarray(display_image)
h, w, _ = img_meta['img_shape']
display_image = display_image[:h, :w, :]
model.show(display_image, result, score_thr=args.score_thr, wait_time=wait_key)
batch_size = data['img'][0].size(0)
for _ in range(batch_size):
prog_bar.update()
if args.out:
print(f'\nwriting results to {args.out}')
mmcv.dump(results, args.out)
if args.eval:
kwargs = cfg.get('evaluation', {})
kwargs.pop('interval', None)
kwargs.pop('gpu_collect', None)
kwargs['metric'] = args.eval
dataset.evaluate(results, **kwargs)
def main():
# base configs
data_root = '/media/' + getpass.getuser(
) + '/Data/DoubleCircle/datasets/kaist-rgbt-encoder/'
# img_norm_cfg = dict(
# mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
# img_norm_cfg_t = dict(
# mean=[123.675, 123.675, 123.675], std=[58.395, 58.395, 58.395], to_rgb=False)
img_norm_cfg = dict(mean=[0, 0, 0], std=[255, 255, 255], to_rgb=True)
img_norm_cfg_t = dict(mean=[0, 0, 0], std=[147, 147, 147], to_rgb=False)
imgs_per_gpu = 4
workers_per_gpu = 2
# test dataset
test = dict(ann_file=data_root + 'annotations-pkl/test-all-rgb.pkl',
img_prefix=data_root + 'images/',
img_scale=1.5,
img_norm_cfg=img_norm_cfg,
img_norm_cfg_t=img_norm_cfg_t,
size_divisor=None,
flip_ratio=0,
with_mask=False,
with_crowd=True,
with_label=True,
test_mode=True)
dataset_test = CoderKaistDataset(**test)
data_loaders_test = build_dataloader(dataset_test,
imgs_per_gpu,
workers_per_gpu,
num_gpus=1,
dist=False,
shuffle=False)
# train
net = AutoEncoder()
load_checkpoint(net, '../../work_dirs/autoencoder/epoch_50.pth')
net.cuda()
loss_fn = torch.nn.MSELoss(reduction='elementwise_mean')
net.eval()
with torch.no_grad():
for i, data_batch in enumerate(data_loaders_test):
code, decode_rgb, decode_thermal = net(
data_batch['img_rgb_in'].cuda(),
data_batch['img_thermal_in'].cuda())
loss_rgb = loss_fn(decode_rgb.cpu(), data_batch['img_rgb_out'])
loss_thermal = loss_fn(decode_thermal.cpu(),
data_batch['img_thermal_out'])
loss_total = loss_thermal + loss_rgb
print('Evaluation:[{}|{}],Loss:{}\n'.format(
i, len(data_loaders_test), loss_total))
output_rgb = decode_rgb.cpu().data
target_rgb = data_batch['img_rgb_out']
output_thermal = decode_thermal.cpu().data
tartget_thermal = data_batch['img_thermal_out']
pic_rgb = np.zeros((output_rgb.shape[0], output_rgb.shape[2],
output_rgb.shape[3], output_rgb.shape[1]),
dtype=np.uint8)
target_pic_rgb = np.zeros(
(output_rgb.shape[0], output_rgb.shape[2], output_rgb.shape[3],
output_rgb.shape[1]),
dtype=np.uint8)
pic_thermal = np.zeros((output_rgb.shape[0], output_rgb.shape[2],
output_rgb.shape[3], output_rgb.shape[1]),
dtype=np.uint8)
target_pic_thermal = np.zeros(
(output_rgb.shape[0], output_rgb.shape[2], output_rgb.shape[3],
output_rgb.shape[1]),
dtype=np.uint8)
mean_rgb = np.array(img_norm_cfg['mean'], dtype=np.float32)
std_rgb = np.array(img_norm_cfg['std'], dtype=np.float32)
mean_thermal = np.array(img_norm_cfg_t['mean'], dtype=np.float32)
std_thermal = np.array(img_norm_cfg_t['std'], dtype=np.float32)
for idx in range(output_rgb.shape[0]):
# for rgb
img = output_rgb[idx,
...].numpy().transpose(1, 2,
0).astype(np.float32)
pic_rgb[idx, :, :, :] = mmcv.imdenormalize(
img, mean=mean_rgb, std=std_rgb,
to_bgr=False).astype(np.uint8)
target_img = target_rgb[idx, ...].numpy().transpose(
1, 2, 0).astype(np.float32)
target_pic_rgb[idx, :, :, :] = mmcv.imdenormalize(
target_img, mean=mean_rgb, std=std_rgb,
to_bgr=False).astype(np.uint8)
# for thermal
img_t = output_thermal[idx, ...].numpy().transpose(
1, 2, 0).astype(np.float32)
pic_thermal[idx, :, :, :] = mmcv.imdenormalize(
img_t, mean=mean_thermal, std=std_thermal,
to_bgr=False).astype(np.uint8)
target_img_t = tartget_thermal[idx, ...].numpy().transpose(
1, 2, 0).astype(np.float32)
target_pic_thermal[idx, :, :, :] = mmcv.imdenormalize(
target_img_t,
mean=mean_thermal,
std=std_thermal,
to_bgr=False).astype(np.uint8)
if not os.path.exists('../../work_dirs/autoencoder/test_rgb'):
os.mkdir('../../work_dirs/autoencoder/test_rgb')
save_images(
torch.from_numpy(pic_rgb.transpose((0, 3, 1, 2))),
'../../work_dirs/autoencoder/test_rgb/image_rgb_{}.png'.format(
i + 1))
save_images(
torch.from_numpy(target_pic_rgb.transpose(0, 3, 1, 2)),
'../../work_dirs/autoencoder/test_rgb/target_image_rgb_{}.png'.
format(i + 1))
save_images(
torch.from_numpy(pic_thermal.transpose((0, 3, 1, 2))),
'../../work_dirs/autoencoder/test_rgb/image_thermal_{}.png'.
format(i + 1))
save_images(
torch.from_numpy(target_pic_thermal.transpose(0, 3, 1, 2)),
'../../work_dirs/autoencoder/test_rgb/target_image_thermal_{}.png'
.format(i + 1))