diff --git a/README.md b/README.md index 6678241..def903e 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -# Multi-Scale Atrous Spatial Pyramid Pooling for Larger and Denser receptive field in Semantic Segmentation +`# Multi-Scale Atrous Spatial Pyramid Pooling for Larger and Denser receptive field in Semantic Segmentation ## Introduction This code have simply upgraded code in the paper, "DenseASPP for Semantic Segmentation in Street Scenes". The paper don't have training code, only proposed model code so our team have implemented training code. But not simple training implementation code, our contribution is proposed. diff --git a/denseASPP.py b/denseASPP.py index 8bf2b1c..9604aac 100644 --- a/denseASPP.py +++ b/denseASPP.py @@ -10,7 +10,6 @@ class DenseASPP(nn.Module): * output_scale can only set as 8 or 16 """ def __init__(self, args, model_cfg, n_class=19, output_stride=8): - # def __init__(self, model_cfg, n_class=19, output_stride=8): super(DenseASPP, self).__init__() # bn_size = model_cfg['bn_size'] bn_size = args.batch_size diff --git a/miscellaneous/misc.py b/miscellaneous/misc.py index 5340b74..b3ac0ca 100755 --- a/miscellaneous/misc.py +++ b/miscellaneous/misc.py @@ -1,14 +1,14 @@ -import os -from math import ceil - +import torch.nn.functional as F import numpy as np import torch -import torch.nn.functional as F -from torch import nn -from torch.autograd import Variable +import math +import os from sklearn.metrics import confusion_matrix -import math +from torch.autograd import Variable +from tqdm import tqdm +from math import ceil +from torch import nn def check_mkdir(dir_name): if not os.path.exists(dir_name): @@ -404,4 +404,41 @@ def wrapper(self, x): outputs_all_scales += outputs return outputs_all_scales - return wrapper \ No newline at end of file + return wrapper + +def evaluate_model(args, val_dataloader, model, criterion): + eval_miou_sum = 0 + # sum_iou_class = 0 + eval_loss_sum = 0 + for sampled_eval in tqdm(val_dataloader.data): + with torch.no_grad(): + eval_image = sampled_eval['image'].cuda(args.gpu, non_blocking=True) + eval_gt = sampled_eval['gt'].cuda(args.gpu, non_blocking=True) + + eval_output = model(eval_image) + model.eval() + eval_loss = criterion(eval_output, eval_gt) + + eval_output = F.softmax(eval_output, dim=1) + eval_output = torch.argmax(eval_output, dim=1) + eval_output = eval_output.contiguous().view(-1) + eval_gt = eval_gt.contiguous().view(-1) + + iou_per_class = [] + for num_class in range(len(val_dataloader.class_names)): + true_class = (eval_output == num_class) + true_label = (eval_gt == num_class) + if true_label.long().sum().item() == 0: + iou_per_class.append(np.nan) + else: + intersect = torch.logical_and(true_class, true_label).sum().float().item() + union = torch.logical_or(true_class, true_label).sum().float().item() + + iou = (intersect + 1e-10) / (union + 1e-10) + iou_per_class.append(iou) + + eval_miou_sum += np.nanmean(iou_per_class) + # sum_iou_class += sum(iou_per_class) + eval_loss_sum += eval_loss + + return eval_loss_sum, eval_miou_sum \ No newline at end of file diff --git a/msaspp_dataloader.py b/msaspp_dataloader.py index 366e473..088c234 100644 --- a/msaspp_dataloader.py +++ b/msaspp_dataloader.py @@ -2,6 +2,7 @@ import random import numpy as np import torch +import torchvision from torchvision.transforms import functional as F from torch.utils.data import Dataset, DataLoader, dataloader @@ -130,6 +131,14 @@ def __getitem__(self, index): gt_copy[gt == key] = value gt = Image.fromarray(gt_copy.astype(np.uint8)) + if self.mode == 'eval': + image = np.array(image, dtype=np.float32) / 255.0 + gt = np.array(gt, dtype=np.float32) + + sample = {'image': image, 'gt': gt} + sample = self.transform(sample) + return data_name, sample + rescaled_image, rescaled_gt = self.resize_random_crop(image, gt, self.args.input_height, self.args.input_width) rescaled_image = np.array(rescaled_image, dtype=np.float32) / 255.0 @@ -139,33 +148,25 @@ def __getitem__(self, index): sample = {'image': image, 'gt': gt} sample = self.transform(sample) - if self.mode == 'eval': - return data_name, sample return sample - # elif self.mode == 'val': - # image = np.array(image, dtype=np.float32) / 255.0 - # gt = np.array(gt, dtype=np.float32) - - # sample = {'image': image, 'gt': gt} - # sample = self.transform(sample) - # return sample - - # def rotate_image(self, img, angle, flag=Image.BILINEAR): - # result = img.rotate(angle, resample=flag) - - # return result - def resize_random_crop(self, image, gt, height, width): scaling = random.uniform(0.5, 2.0) scale_w, scale_h = [int(i*scaling) for i in image.size] + resized_image = image.resize((scale_w, scale_h), Image.CUBIC) resized_gt = gt.resize((scale_w, scale_h), Image.NEAREST) - i, j, h, w = transforms.RandomCrop.get_params(resized_image, output_size=(height, width)) - crop_image = F.crop(resized_image, i, j, h, w) - crop_gt = F.crop(resized_gt, i, j, h, w) + x1 = random.randint(0, scale_w - width) + y1 = random.randint(0, scale_h - height) + + crop_image = resized_image.crop((x1, y1, x1 + width, y1 + height)) + crop_gt = resized_gt.crop((x1, y1, x1 + width, y1 + height)) + + # i, j, h, w = transforms.RandomCrop.get_params(resized_image, output_size=(height, width)) + # crop_image = F.crop(resized_image, i, j, h, w) + # crop_gt = F.crop(resized_gt, i, j, h, w) return crop_image, crop_gt @@ -185,13 +186,13 @@ def train_preprocess(self, image, gt): def augment_image(self, image): # gamma augmentation - # gamma = random.uniform(0.5, 2.0) - # image_aug = image ** gamma + gamma = random.uniform(0.9, 1.1) + image_aug = image ** gamma # brightness augmentation - brightness = random.uniform(-10, 10) - image_aug = image * brightness - + brightness = random.uniform(0.9, 1.1) + image_aug = image_aug * brightness + # color augmentation # colors = np.random.uniform(0.9, 1.1, size=3) # white = np.ones((image.shape[0], image.shape[1])) diff --git a/msaspp_inference.py b/msaspp_inference.py index bbf3cae..d6011b6 100644 --- a/msaspp_inference.py +++ b/msaspp_inference.py @@ -14,45 +14,35 @@ from tqdm import tqdm parser = argparse.ArgumentParser(description="msaspp image segmentation inference code") -parser.add_argument('--model_name', type=str, help='model name to be trained', default='denseaspp-v2') -parser.add_argument('--data_path', type=str, help='test data path', default=os.getcwd()) -parser.add_argument('--input_height', type=int, help='input height', default=512) -parser.add_argument('--input_width', type=int, help='input width', default=512) -parser.add_argument('--batch_size', type=int, help='train batch size', default=8) -parser.add_argument('--checkpoint_path', type=str, help='path to a specific checkpoint to load', default=os.path.join(os.getcwd(), 'log')) -parser.add_argument('--num_checkpoint', type=str, help='model to be saved after training', default='model-0029000_mIoU-0.486.pth') -parser.add_argument('--num_seed', type=int, help='random seed number', default=1) -parser.add_argument('--num_threads', type=int, help='number of threads to use for data loading', default=5) -parser.add_argument('--gpu', type=int, help='GPU id to use', default=0) +parser.add_argument('--model_name', type=str, help='model name to be trained', default='denseaspp-v3') +parser.add_argument('--data_path', type=str, help='test data path', default=os.getcwd()) +parser.add_argument('--input_height', type=int, help='input height', default=512) +parser.add_argument('--input_width', type=int, help='input width', default=512) +parser.add_argument('--batch_size', type=int, help='train batch size', default=8) +parser.add_argument('--log_directory', type=str, help='directory to save checkpoints and summaries', default=os.path.join(os.getcwd(), 'log')) +parser.add_argument('--num_checkpoint', type=str, help='model to be saved after training', default='model-0029500_mIoU-0.445.pth') +parser.add_argument('--num_seed', type=int, help='random seed number', default=1) +parser.add_argument('--num_threads', type=int, help='number of threads to use for data loading', default=5) +parser.add_argument('--gpu', type=int, help='GPU id to use', default=0) args = parser.parse_args() -def save_prediction(prediction): - command = 'mkdir ' + os.path.join(os.getcwd(), "prediction_image") - os.system(command) - - plt.figure(figsize=(40, 10)) - for i in range(19): - pred = prediction[:, i, :].permute(1, 2, 0) - pred = pred.detach().cpu().numpy() - plt.subplot(1, 19, i+1) - plt.imshow(pred) - plt.axis('off') - - plt.savefig(os.path.join(os.getcwd(), "prediction_image", "pred_imageset.png"), dpi=400, bbox_inches='tight') - print('Saving done') -def cover_colormap(image): +def check_directory(path): + if not os.path.exists(path): + os.mkdir(path) + +def fill_colormap(prediction): COLOR_MAP = [(128, 64, 128), (244, 35, 232), (70, 70, 70), (102, 102, 156), (190, 153, 153), (153, 153, 153), - (250, 170, 30), (220, 220, 0), (107, 142, 35), (152, 251, 152), (70, 130, 180), (220, 20, 60), - (255, 0, 0), (0, 0, 142), (0, 0, 70), (0, 60, 100), (0, 80, 100), (0, 0, 230), (119, 11, 32)] + (250, 170, 30), (220, 220, 0), (107, 142, 35), (152, 251, 152), (70, 130, 180), (220, 20, 60), + (255, 0, 0), (0, 0, 142), (0, 0, 70), (0, 60, 100), (0, 80, 100), (0, 0, 230), (119, 11, 32)] - image = image.data.cpu().numpy() - # image = image[0, :] - image = image.argmax(axis=1)[0] - row, col = image.shape + prediction = prediction.data.cpu().numpy() + # prediction = prediction[0, :] + prediction = prediction.argmax(axis=1)[0] + row, col = prediction.shape dst = np.zeros((row, col, 3), dtype=np.uint8) for i in range(19): - dst[image == i] = COLOR_MAP[i] + dst[prediction == i] = COLOR_MAP[i] dst = np.array(dst, dtype=np.uint8) dst = cv2.cvtColor(dst, cv2.COLOR_RGB2BGR) @@ -78,46 +68,48 @@ def get_iou(prediction, gt): return iou_per_class def test(args): - dataloader_val = msasppDataLoader(args, mode='val') + dataloader_eval = msasppDataLoader(args, mode='eval') model = DenseASPP(args, model_cfg=DenseASPP121.Model_CFG) torch.cuda.set_device(args.gpu) model = torch.nn.DataParallel(model, device_ids=[args.gpu]) - model_path = os.path.join(args.checkpoint_path, args.model_name, 'eval_model', args.num_checkpoint) + model_path = os.path.join(args.log_directory, args. model_name, 'eval_model', args.num_checkpoint) checkpoint = torch.load(model_path) model.load_state_dict(checkpoint['model']) model.eval() model.cuda(args.gpu) with torch.no_grad(): - for idx, sample in enumerate(dataloader_val.data): - + for idx, (name, sample) in enumerate(dataloader_eval.data): image = sample['image'].cuda(args.gpu) gt = sample['gt'].cuda(args.gpu) prediction = model(image) - + prediction = F.log_softmax(prediction, dim=1) iou_per_class = get_iou(prediction, gt) + if prediction.size()[0] < 1024: + prediction = F.upsample(prediction, size=(1024, 2048), mode='bilinear') mIoU = np.nanmean(iou_per_class) - print(mIoU) - # color_prediction = cover_colormap(prediction) + color_prediction = fill_colormap(prediction) + print("idx: {}, data name: {}, mIoU: {}".format(idx+1, *name, mIoU)) - # plt.suptitle("{}".format(*data_name)) - # plt.figure(figsize=(10, 5)) - # plt.subplot(1, 2, 1) - # plt.imshow(color_prediction) - # plt.axis('off') + check_directory(path=os.path.join(args.log_directory, args.model_name, 'prediction_image')) + plt.figure(figsize=(18, 5)) + plt.suptitle("name: {} - mIoU: {}".format(*name, mIoU)) + plt.subplot(1, 2, 1) + plt.imshow(color_prediction) + plt.axis('off') - # plt_img = image.squeeze().permute(1, 2, 0) - # plt_img = plt_img.data.cpu().numpy() - # plt.subplot(1, 2, 2) - # plt.imshow(plt_img) - # plt.axis('off') - - # plt.savefig('{}_color.png'.format(*data_name), dpi=400, bbox_inches='tight') - a=1 + plt_img = image.squeeze().permute(1, 2, 0) + plt_img = plt_img.data.cpu().numpy() + plt.subplot(1, 2, 2) + plt.imshow(plt_img) + plt.axis('off') + + plt.savefig(os.path.join(args.log_directory, args.model_name, 'prediction_image', '{}_color'.format(*name)), dpi=400, bbox_inches='tight') + if __name__ == "__main__": test(args) \ No newline at end of file diff --git a/msaspp_main.py b/msaspp_main.py index f2c1732..a4f15fb 100644 --- a/msaspp_main.py +++ b/msaspp_main.py @@ -21,7 +21,7 @@ parser = argparse.ArgumentParser(description='Multi-scale ASPP training') parser.add_argument('--mode', type=str, help='training and validation mode', default='train') -parser.add_argument('--model_name', type=str, help='model name to be trained', default='denseaspp-v3') +parser.add_argument('--model_name', type=str, help='model name to be trained', default='denseaspp-v5.1') # Dataset parser.add_argument('--data_path', type=str, help='training data path', default=os.getcwd()) @@ -64,9 +64,36 @@ def check_folder(path): if not os.path.exists(path): os.makedirs(path) +def fast_hist(pred, gtruth, num_classes): + # mask indicates pixels we care about + mask = (gtruth >= 0) & (gtruth < num_classes) + + # stretch ground truth labels by num_classes + # class 0 -> 0 + # class 1 -> 19 + # class 18 -> 342 + # + # TP at 0 + 0, 1 + 1, 2 + 2 ... + # + # TP exist where value == num_classes*class_id + class_id + # FP = row[class].sum() - TP + # FN = col[class].sum() - TP + hist = np.bincount(num_classes * gtruth[mask].astype(int) + pred[mask], + minlength=num_classes ** 2) + hist = hist.reshape(num_classes, num_classes) + return hist + +def eval_metrics(hist): + acc = np.diag(hist).sum() / hist.sum() + acc_cls = np.diag(hist) / hist.sum(axis=1) + acc_cls = np.nanmean(acc_cls) + divisor = hist.sum(axis=1) + hist.sum(axis=0) - np.diag(hist) + iu = np.diag(hist) / divisor + + return iu, acc, acc_cls def evaluate_model(val_dataloader, model, criterion): - eval_iou_sum_score = 0 + iou_acc = 0 eval_loss_sum = 0 for sampled_eval in tqdm(val_dataloader.data): with torch.no_grad(): @@ -75,31 +102,25 @@ def evaluate_model(val_dataloader, model, criterion): eval_output = model(eval_image) model.eval() - eval_loss = criterion(eval_output, eval_gt) - - eval_output = F.softmax(eval_output, dim=1) - eval_output = torch.argmax(eval_output, dim=1) - eval_output = eval_output.contiguous().view(-1) - eval_gt = eval_gt.contiguous().view(-1) - iou_per_class = [] - for num_class in range(len(val_dataloader.class_names)): - true_class = (eval_output == num_class) - true_label = (eval_gt == num_class) - if true_label.long().sum().item() == 0: - iou_per_class.append(np.nan) - else: - intersect = torch.logical_and(true_class, true_label).sum().float().item() - union = torch.logical_or(true_class, true_label).sum().float().item() - - iou = (intersect + 1e-10) / (union + 1e-10) - iou_per_class.append(iou) - - eval_iou_sum_score += np.nanmean(iou_per_class) + eval_loss = criterion(eval_output, eval_gt) eval_loss_sum += eval_loss - return eval_loss_sum, eval_iou_sum_score + eval_output = F.softmax(eval_output, dim=1) + eval_output = torch.argmax(eval_output, dim=1) + + eval_output = eval_output.detach().cpu().numpy() + eval_gt = eval_gt.detach().cpu().numpy() + + _iou_acc = fast_hist(eval_output.flatten(), eval_gt.flatten(), 19) + iou_acc += _iou_acc + + mean_eval_loss = eval_loss_sum/len(val_dataloader.data) + iu, acc, acc_cls = eval_metrics(iou_acc) + mean_iu = np.nanmean(iu) + return mean_iu, acc, acc_cls, mean_eval_loss + def main(): model_filename = args.model_name + '.py' command = 'mkdir ' + os.path.join(args.log_directory, args.model_name) @@ -112,22 +133,27 @@ def main(): os.system(command) if args.checkpoint_path == '': - aux_out_path = os.path.join(args.log_directory, args.model_name) + model_out_path = os.path.join(args.log_directory, args.model_name, model_filename) - command = 'cp denseASPP.py ' + aux_out_path + command = 'cp denseASPP.py ' + model_out_path os.system(command) + aux_out_path = os.path.join(args.log_directory, args.model_name) + command = 'cp msaspp_main.py ' + aux_out_path os.system(command) command = 'cp msaspp_dataloader.py ' + aux_out_path os.system(command) else: + args.checkpoint_path = os.path.join(args.checkpoint_path, args.model_name, 'model') loaded_model_dir = os.path.dirname(args.checkpoint_path) loaded_model_name = os.path.basename(loaded_model_dir) loaded_model_filename = loaded_model_name + '.py' - model_out_path = os.path.join(args.checkpoint_path, args.model_name, model_filename) + model_out_path = os.path.join(args.log_directory, args.model_name, model_filename) + command = 'cp ' + os.path.join(loaded_model_dir, loaded_model_filename) + ' ' + model_out_path + os.system(command) torch.cuda.empty_cache() args.distributed = args.world_size > 1 or args.multiprocessing_distributed @@ -168,6 +194,7 @@ def main_worker(ngpus_per_node, args): global_step = 0 optimizer = torch.optim.Adam(model.parameters(), weight_decay=args.weight_decay, lr=args.learning_rate) + model_just_loaded = False if args.checkpoint_path != '': if os.path.isfile(args.checkpoint_path): print("Loading checkpoint '{}'".format(args.checkpoint_path)) @@ -185,6 +212,7 @@ def main_worker(ngpus_per_node, args): else: print("No checkpoint found at '{}'".format(args.checkpoint_path)) + model_just_loaded = True if args.retrain: global_step = 0 @@ -236,13 +264,12 @@ def main_worker(ngpus_per_node, args): train_loss_list.append(loss) duration += time.time() - befor_op_time - if global_step and global_step % args.log_freq == 0: + if global_step and global_step % args.log_freq == 0 and not model_just_loaded: examples_per_sec = args.batch_size / duration * args.log_freq duration = 0 time_sofar = (time.time() - start_time) / 3600 training_time_left = (num_total_steps / global_step - 1.0) * time_sofar - # if not args.multiprocessing_distributed or (args.multiprocessing_distributed and args.rank % ngpus_per_node == 0): - # print("{}".format(args.model_name)) + print_string = 'GPU: {} | examples/s: {:4.2f} | Average train loss: {:.5f} | time elapsed: {:.2f}h | time left: {:.2f}h' print(print_string.format(args.gpu, examples_per_sec, sum(train_loss_list)/len(train_loss_list), time_sofar, training_time_left)) @@ -264,16 +291,15 @@ def main_worker(ngpus_per_node, args): 'optimizer': optimizer.state_dict()} torch.save(checkpoint, os.path.join(args.log_directory, args.model_name, 'model', 'model-{:07d}.pth'.format(global_step))) - if global_step and global_step % args.save_freq == 0: + if global_step and global_step % args.save_freq == 0 and not model_just_loaded: if not args.multiprocessing_distributed or (args.multiprocessing_distributed and args.rank % ngpus_per_node == 0): - eval_loss_sum, eval_iou_sum_score = evaluate_model(val_dataloader, model, criterion) - mIoU = eval_iou_sum_score / len(val_dataloader.data) - mIoU_list.append(mIoU) + mean_iu, acc, acc_cls, mean_eval_loss = evaluate_model(val_dataloader, model, criterion) + mIoU_list.append(mean_iu) idx_max = mIoU_list.index(max(mIoU_list)) if idx_max == 0 and idx_max == (global_step // args.save_freq - 1): - print_string = 'GPU: {} | Epoch: {}/{} | Average train loss: {:.5f} | Average validation loss: {:.5f} | evaluation mIoU: {:.5f}' - print(print_string.format(args.gpu, epoch, args.num_epochs, sum(train_loss_list)/len(train_loss_list), eval_loss_sum/len(val_dataloader.data), mIoU_list[idx_max])) + print_string = 'GPU: {} | Epoch: {}/{} | Average train loss: {:.5f} | Average validation loss: {:.5f} | evaluation mIoU: {:.5f} | evaluation acc: {:.5f} | evaluation acc cls: {:.5f}' + print(print_string.format(args.gpu, epoch, args.num_epochs, sum(train_loss_list)/len(train_loss_list), mean_eval_loss, mIoU_list[idx_max], acc, acc_cls)) checkpoint = {'global_step': global_step, 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), @@ -285,8 +311,8 @@ def main_worker(ngpus_per_node, args): pass elif idx_max !=0 and idx_max == (global_step // args.save_freq - 1): - print_string = 'GPU: {} | Epoch: {}/{} | Average train loss: {:.5f} | Average validation loss: {:.5f} | evaluation mIoU: {:.5f}' - print(print_string.format(args.gpu, epoch, args.num_epochs, sum(train_loss_list)/len(train_loss_list), eval_loss_sum/len(val_dataloader.data), max(mIoU_list))) + print_string = 'GPU: {} | Epoch: {}/{} | Average train loss: {:.5f} | Average validation loss: {:.5f} | evaluation mIoU: {:.5f} | evaluation acc: {:.5f} | evaluation acc cls: {:.5f}' + print(print_string.format(args.gpu, epoch, args.num_epochs, sum(train_loss_list)/len(train_loss_list), mean_eval_loss, max(mIoU_list), acc, acc_cls)) checkpoint = {'global_step': global_step, 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), @@ -294,6 +320,8 @@ def main_worker(ngpus_per_node, args): torch.save(checkpoint, os.path.join(args.log_directory, args.model_name, 'eval_model', 'model-{:07d}_mIoU-{:.3f}.pth'.format(global_step, max(mIoU_list)))) model.train() + + model_just_loaded = False global_step += 1 epoch += 1 diff --git a/msaspp_test.py b/msaspp_test.py deleted file mode 100644 index 1a26b3a..0000000 --- a/msaspp_test.py +++ /dev/null @@ -1,156 +0,0 @@ - -import torch.nn.functional as F -import matplotlib.pyplot as plt -import numpy as np -import argparse -import torch -import cv2 -import os - -import addToConfusionMatrix -from msaspp_dataloader import msasppDataLoader -from cityscapesLabels import labels, id2label -from denseASPP import DenseASPP -from cfgs import DenseASPP121 -from tqdm import tqdm - -parser = argparse.ArgumentParser(description="msaspp image segmentation inference code") -parser.add_argument('--model_name', type=str, help='model name to be trained', default='denseaspp-v2') -parser.add_argument('--data_path', type=str, help='test data path', default=os.getcwd()) -parser.add_argument('--input_height', type=int, help='input height', default=512) -parser.add_argument('--input_width', type=int, help='input width', default=512) -parser.add_argument('--batch_size', type=int, help='train batch size', default=8) -parser.add_argument('--checkpoint_path', type=str, help='path to a specific checkpoint to load', default=os.path.join(os.getcwd(), 'log')) -parser.add_argument('--num_checkpoint', type=str, help='model to be saved after training', default='model-0029000_mIoU-0.486.pth') -parser.add_argument('--num_seed', type=int, help='random seed number', default=1) -parser.add_argument('--num_threads', type=int, help='number of threads to use for data loading', default=5) -parser.add_argument('--gpu', type=int, help='GPU id to use', default=0) -args = parser.parse_args() - -def save_prediction(prediction): - command = 'mkdir ' + os.path.join(os.getcwd(), "prediction_image") - os.system(command) - - plt.figure(figsize=(40, 10)) - for i in range(19): - pred = prediction[:, i, :].permute(1, 2, 0) - pred = pred.detach().cpu().numpy() - plt.subplot(1, 19, i+1) - plt.imshow(pred) - plt.axis('off') - - plt.savefig(os.path.join(os.getcwd(), "prediction_image", "pred_imageset.png"), dpi=400, bbox_inches='tight') - print('Saving done') - -def cover_colormap(image): - COLOR_MAP = [(128, 64, 128), (244, 35, 232), (70, 70, 70), (102, 102, 156), (190, 153, 153), (153, 153, 153), - (250, 170, 30), (220, 220, 0), (107, 142, 35), (152, 251, 152), (70, 130, 180), (220, 20, 60), - (255, 0, 0), (0, 0, 142), (0, 0, 70), (0, 60, 100), (0, 80, 100), (0, 0, 230), (119, 11, 32)] - - image = image.data.cpu().numpy() - # image = image[0, :] - image = image.argmax(axis=1)[0] - row, col = image.shape - dst = np.zeros((row, col, 3), dtype=np.uint8) - for i in range(19): - dst[image == i] = COLOR_MAP[i] - dst = np.array(dst, dtype=np.uint8) - dst = cv2.cvtColor(dst, cv2.COLOR_RGB2BGR) - - return dst - -def get_iou(prediction, gt): - prediction = torch.argmax(prediction, dim=1) - prediction = prediction.contiguous().view(-1) - gt = gt.contiguous().view(-1) - iou_per_class = [] - for num_class in range(19): - true_class = (prediction == num_class) - true_label = (gt == num_class) - if true_label.long().sum().item() == 0: - iou_per_class.append(np.nan) - else: - intersect = torch.logical_and(true_class, true_label).sum().float().item() - union = torch.logical_or(true_class, true_label).sum().float().item() - - iou = (intersect + 1e-10) / (union + 1e-10) - iou_per_class.append(iou) - - return iou_per_class - -def generateMatrix(args): - - args.evalLabels = [] - for label in labels: - if (label.id < 0): - continue - # we append all found labels, regardless of being ignored - args.evalLabels.append(label.id) - maxId = max(args.evalLabels) - # We use longlong type to be sure that there are no overflows - return np.zeros(shape=(maxId+1, maxId+1),dtype=np.ulonglong) - -def getIouScoreForLabel(label, confMatrix, args): - if id2label[label].ignoreInEval: - return float('nan') - - # the number of true positive pixels for this label - # the entry on the diagonal of the confusion matrix - tp = np.longlong(confMatrix[label,label]) - - # the number of false negative pixels for this label - # the row sum of the matching row in the confusion matrix - # minus the diagonal entry - fn = np.longlong(confMatrix[label,:].sum()) - tp - - # the number of false positive pixels for this labels - # Only pixels that are not on a pixel with ground truth label that is ignored - # The column sum of the corresponding column in the confusion matrix - # without the ignored rows and without the actual label of interest - notIgnored = [l for l in args.evalLabels if not id2label[l].ignoreInEval and not l==label] - fp = np.longlong(confMatrix[notIgnored,label].sum()) - - # the denominator of the IOU score - denom = (tp + fp + fn) - if denom == 0: - return float('nan') - - # return IOU - return float(tp) / denom - -def test(args): - confMatrix = generateMatrix(args) - classScoreList = {} - for label in args.evalLabels: - labelName = id2label[label].name - classScoreList[labelName] = getIouScoreForLabel(label, confMatrix, args) - - dataloader_test = msasppDataLoader(args, mode='test') - - model = DenseASPP(args, model_cfg=DenseASPP121.Model_CFG) - torch.cuda.set_device(args.gpu) - model = torch.nn.DataParallel(model, device_ids=[args.gpu]) - - model_path = os.path.join(args.checkpoint_path, args.model_name, 'eval_model', args.num_checkpoint) - checkpoint = torch.load(model_path) - model.load_state_dict(checkpoint['model']) - model.eval() - model.cuda(args.gpu) - - with torch.no_grad(): - for data_name, sample in dataloader_test.data: - image = sample['image'].cuda(args.gpu) - gt = sample['gt'].cuda(args.gpu) - - prediction = model(image) - - prediction = F.softmax(prediction, dim=1) - confMatrix = addToConfusionMatrix.cEvaluatePair(predictionNp, groundTruthNp, confMatrix, args.evalLabels) - iou_per_class = get_iou(prediction, gt) - color_prediction = cover_colormap(prediction) - - plt.imshow(color_prediction) - plt.savefig('{}_color.png'.format(*data_name)) - a=1 -if __name__ == "__main__": - test(args) \ No newline at end of file diff --git a/transforms.py b/transforms.py new file mode 100644 index 0000000..474a2a7 --- /dev/null +++ b/transforms.py @@ -0,0 +1,76 @@ +import random + + +import numpy as np +from skimage.filters import gaussian +import torch +from PIL import Image, ImageFilter + + +class Compose(object): + def __init__(self, transform): + self.transform = transform + + + def __call__(self, mask): + for t in self.transform: + mask = t(mask) + return mask + + +class RandomVerticalFlip(object): + def __call__(self, img): + if random.random() < 0.5: + return img.transpose(Image.FLIP_TOP_BOTTOM) + return img + + + + +class DeNormalize(object): + def __init__(self, mean, std): + self.mean = mean + self.std = std + + + def __call__(self, tensor): + for t, m, s in zip(tensor, self.mean, self.std): + t.mul_(s).add_(m) + return tensor + + + + +class MaskToTensor(object): + def __call__(self, img): + return torch.from_numpy(np.array(img, dtype=np.int32)).long() + + + + +class FreeScale(object): + def __init__(self, size, interpolation=Image.BILINEAR): + self.size = tuple(reversed(size)) # size: (h, w) + self.interpolation = interpolation + + + def __call__(self, img): + return img.resize(self.size, self.interpolation) + + + + +class FlipChannels(object): + def __call__(self, img): + img = np.array(img)[:, :, ::-1] + return Image.fromarray(img.astype(np.uint8)) + + + + +class RandomGaussianBlur(object): + def __call__(self, img): + sigma = 0.15 + random.random() * 1.15 + blurred_img = gaussian(np.array(img), sigma=sigma, multichannel=True) + blurred_img *= 255 + return Image.fromarray(blurred_img.astype(np.uint8)) \ No newline at end of file