Add files via upload

ASM/ASM_IOU.py

@@ -0,0 +1,74 @@
+import numpy as np
+import argparse
+import json
+from PIL import Image
+from os.path import join
+import csv
+
+
+def fast_hist(a, b, n):
+    k = (a >= 0) & (a < n)
+    return np.bincount(n * a[k].astype(int) + b[k], minlength=n ** 2).reshape(n, n)
+
+
+def per_class_iu(hist):
+    return np.diag(hist) / (hist.sum(1) + hist.sum(0) - np.diag(hist))
+
+
+def label_mapping(input, mapping):
+    output = np.copy(input)
+    for ind in range(len(mapping)):
+        output[input == mapping[ind][0]] = mapping[ind][1]
+    return np.array(output, dtype=np.int64)
+
+
+def compute_mIoU(gt_dir, pred_dir, devkit_dir=''):
+    """
+    Compute IoU given the predicted colorized images and
+    """
+    with open(join(devkit_dir, 'info.json'), 'r') as fp:
+      info = json.load(fp)
+    num_classes = np.int(info['classes'])
+    print('Num classes', num_classes)
+    name_classes = np.array(info['label'], dtype=np.str)
+    mapping = np.array(info['label2train'], dtype=np.int)
+    hist = np.zeros((num_classes, num_classes))
+
+    image_path_list = join(devkit_dir, 'val.txt')
+    label_path_list = join(devkit_dir, 'label.txt')
+    gt_imgs = open(label_path_list, 'r').read().splitlines()
+    gt_imgs = [join(gt_dir, x) for x in gt_imgs]
+    pred_imgs = open(image_path_list, 'r').read().splitlines()
+    pred_imgs = [join(pred_dir, x.split('/')[-1]) for x in pred_imgs]
+
+    for ind in range(len(gt_imgs)):
+        pred = np.array(Image.open(pred_imgs[ind]))
+        label = np.array(Image.open(gt_imgs[ind]))
+        label = label_mapping(label, mapping)
+        if len(label.flatten()) != len(pred.flatten()):
+            print('Skipping: len(gt) = {:d}, len(pred) = {:d}, {:s}, {:s}'.format(len(label.flatten()), len(pred.flatten()), gt_imgs[ind], pred_imgs[ind]))
+            continue
+        hist += fast_hist(label.flatten(), pred.flatten(), num_classes)
+        if ind > 0 and ind % 10 == 0:
+            print('{:d} / {:d}: {:0.2f}'.format(ind, len(gt_imgs), 100*np.mean(per_class_iu(hist))))
+
+    mIoUs = per_class_iu(hist)
+    for ind_class in range(num_classes):
+        print('===>' + name_classes[ind_class] + ':\t' + str(round(mIoUs[ind_class] * 100, 2)))
+    print('===> mIoU: ' + str(round(np.nanmean(mIoUs) * 100, 2)))
+    return str(round(np.nanmean(mIoUs) * 100, 2))
+
+
+def main(gt_dir, pred_dir, devkit_dir):
+   return compute_mIoU(gt_dir, pred_dir, devkit_dir)
+
+
+if __name__ == "__main__":
+    with open("mIoU_result/GTA2Cityscapes_mIoU.csv","a+",newline="") as datacsv:
+        csvwriter = csv.writer(datacsv,dialect = ("excel"))
+        for i in range(1, 50):
+            gt_dir = '/data02/yawei/Data/Cityscapes/gtFine/val/'
+            pred_dir = './result/GTA2Cityscapes_{0:d}'.format(i*2000)
+            devkit_dir = './dataset/cityscapes_list'
+            mIoU = main(gt_dir, pred_dir, devkit_dir)
+            csvwriter.writerow([mIoU])

ASM/ASM_evaluate.py

@@ -0,0 +1,132 @@
+import argparse
+import numpy as np
+import torch
+from torch.autograd import Variable
+from model.ASM_G import Res_Deeplab
+from dataset.cityscapes_dataset import cityscapesDataSet
+from torch.utils import data
+import os
+from PIL import Image
+import torch.nn as nn
+
+IMG_MEAN = np.array((104.00698793,116.66876762,122.67891434), dtype=np.float32)
+
+DATA_DIRECTORY = '/data02/yawei/Data/Cityscapes/'
+DATA_LIST_PATH = './dataset/cityscapes_list/val.txt'
+SAVE_PATH = './result/cityscapes'
+RESTORE_FROM = 'http://vllab.ucmerced.edu/ytsai/CVPR18/GTA2Cityscapes_multi-ed35151c.pth'
+IGNORE_LABEL = 255
+NUM_CLASSES = 19
+NUM_STEPS = 500 # Number of images in the validation set.
+SET = 'val'
+#128, 64, 128
+palette = [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]
+zero_pad = 256 * 3 - len(palette)
+for i in range(zero_pad):
+    palette.append(0)
+
+
+def colorize_mask(mask):
+    # mask: numpy array of the mask
+    new_mask = Image.fromarray(mask.astype(np.uint8)).convert('P')
+    new_mask.putpalette(palette)
+
+    return new_mask
+
+
+def create_map(input_size, mode):
+    if mode == 'h':
+        T_base = torch.arange(0, float(input_size[1]))
+        T_base = T_base.view(input_size[1], 1)
+        T = T_base
+        for i in range(input_size[0] - 1):
+            T = torch.cat((T, T_base), 1)
+        T = torch.div(T, float(input_size[1]))
+    if mode == 'w':
+        T_base = torch.arange(0, float(input_size[0]))
+        T_base = T_base.view(1, input_size[0])
+        T = T_base
+        for i in range(input_size[1] - 1):
+            T = torch.cat((T, T_base), 0)
+        T = torch.div(T, float(input_size[0]))
+    T = T.view(1, 1, T.size(0), T.size(1))
+    return T
+
+
+def get_arguments():
+    """Parse all the arguments provided from the CLI.
+    Returns:
+      A list of parsed arguments.
+    """
+    parser = argparse.ArgumentParser(description="DeepLab-ResNet Network")
+    parser.add_argument("--data-dir", type=str, default=DATA_DIRECTORY,
+                        help="Path to the directory containing the Cityscapes dataset.")
+    parser.add_argument("--data-list", type=str, default=DATA_LIST_PATH,
+                        help="Path to the file listing the images in the dataset.")
+    parser.add_argument("--ignore-label", type=int, default=IGNORE_LABEL,
+                        help="The index of the label to ignore during the training.")
+    parser.add_argument("--num-classes", type=int, default=NUM_CLASSES,
+                        help="Number of classes to predict (including background).")
+    parser.add_argument("--restore-from", type=str, default=RESTORE_FROM,
+                        help="Where restore model parameters from.")
+    parser.add_argument("--gpu", type=int, default=0,
+                        help="choose gpu device.")
+    parser.add_argument("--set", type=str, default=SET,
+                        help="choose evaluation set.")
+    parser.add_argument("--save", type=str, default=SAVE_PATH,
+                        help="Path to save result.")
+    return parser.parse_args()
+
+
+def main():
+    """Create the model and start the evaluation process."""
+
+    for i in range(1, 50):
+        model_path = './snapshots/GTA2Cityscapes/GTA5_{0:d}.pth'.format(i*2000)
+        save_path = './result/GTA2Cityscapes_{0:d}'.format(i*2000)
+        args = get_arguments()
+
+        gpu0 = args.gpu
+
+        if not os.path.exists(save_path):
+            os.makedirs(save_path)
+
+        model = Res_Deeplab(num_classes=args.num_classes)
+
+        saved_state_dict = torch.load(model_path)
+        model.load_state_dict(saved_state_dict)
+
+        model.eval()
+        model.cuda(gpu0)
+
+        testloader = data.DataLoader(cityscapesDataSet(args.data_dir, args.data_list, crop_size=(1024,512), mean=IMG_MEAN, scale=False, mirror=False, set=args.set),
+                                        batch_size=1, shuffle=False, pin_memory=True)
+
+        interp = nn.Upsample(size=(1024, 2048), mode='bilinear', align_corners=True)
+
+        with torch.no_grad():
+            for index, batch in enumerate(testloader):
+                if index % 100 == 0:
+                    print('%d processd' % index)
+                image, _, _, name = batch
+                #image = image.resize((1200,600),Image.BICUBIC)
+                output, _ = model(Variable(image).cuda(gpu0))
+
+                output = interp(output).cpu().data[0].numpy()
+
+                output = output.transpose(1,2,0)
+                output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
+
+                output_col = colorize_mask(output)
+                output = Image.fromarray(output)
+
+                name = name[0].split('/')[-1]
+                output.save('%s/%s' % (save_path, name))
+
+                output_col.save('%s/%s_color.png' % (save_path, name.split('.')[0]))
+
+        print(save_path)
+if __name__ == '__main__':
+    main()