formatting, fix colorization rgb2bgr

Author: hangzhaomit

eval.py

@@ -8,8 +8,6 @@
 import torch.nn as nn
 from torch.autograd import Variable
 from scipy.io import loadmat
-from scipy.misc import imsave
-from scipy.ndimage import zoom
 # Our libs
 from dataset import ValDataset
 from models import ModelBuilder, SegmentationModule
@@ -84,7 +82,9 @@ def evaluate(segmentation_module, loader, args):
 
         # visualization
         if args.visualize:
-            visualize_result((batch_data['img_ori'], seg_label, batch_data['info']), preds, args)
+            visualize_result(
+                (batch_data['img_ori'], seg_label, batch_data['info']),
+                preds, args)
 
     iou = intersection_meter.sum / (union_meter.sum + 1e-10)
     for i, _iou in enumerate(iou):
@@ -136,7 +136,7 @@ def main(args):
     # Model related arguments
     parser.add_argument('--id', required=True,
                         help="a name for identifying the model to load")
-    parser.add_argument('--suffix', default='_best.pth',
+    parser.add_argument('--suffix', default='_epoch_13.pth',
                         help="which snapshot to load")
     parser.add_argument('--arch_encoder', default='resnet50_dilated8',
                         help="architecture of net_encoder")
@@ -170,7 +170,7 @@ def main(args):
     # Misc arguments
     parser.add_argument('--ckpt', default='./ckpt',
                         help='folder to output checkpoints')
-    parser.add_argument('--visualize', default=0,
+    parser.add_argument('--visualize', action='store_true',
                         help='output visualization?')
     parser.add_argument('--result', default='./result',
                         help='folder to output visualization results')
@@ -180,17 +180,15 @@ def main(args):
     args = parser.parse_args()
     print(args)
 
-    #torch.cuda.set_device(args.gpu_id)
-
-    # scales for evaluation
-    # args.scales = (1, )
-    # args.scales = (0.5, 0.75, 1, 1.25, 1.5)
+    # torch.cuda.set_device(args.gpu_id)
 
     # absolute paths of model weights
     args.weights_encoder = os.path.join(args.ckpt, args.id,
                                         'encoder' + args.suffix)
     args.weights_decoder = os.path.join(args.ckpt, args.id,
                                         'decoder' + args.suffix)
+    assert os.path.exists(args.weights_encoder) and \
+        os.path.exists(args.weights_encoder), 'checkpoint does not exitst!'
 
     args.result = os.path.join(args.result, args.id)
     if not os.path.isdir(args.result):

train.py

@@ -8,19 +8,16 @@
 import numpy as np
 import torch
 import torch.nn as nn
-from torch.autograd import Variable
-from scipy.io import loadmat
-from scipy.misc import imresize, imsave
 # Our libs
 from dataset import TrainDataset, ValDataset
 from models import ModelBuilder, SegmentationModule
 from utils import AverageMeter, colorEncode, accuracy
 from lib.nn import UserScatteredDataParallel, user_scattered_collate, patch_replication_callback
 import lib.utils.data as torchdata
 
-import matplotlib
-matplotlib.use('Agg')
-import matplotlib.pyplot as plt
+# import matplotlib
+# matplotlib.use('Agg')
+# import matplotlib.pyplot as plt
 
 
 # train one epoch
@@ -270,15 +267,14 @@ def main(args):
     args.id += '-' + str(args.arch_decoder)
     args.id += '-ngpus' + str(args.num_gpus)
     args.id += '-batchSize' + str(args.batch_size)
-    #args.id += '-imgSize' + str(args.imgSize)
     args.id += '-imgMaxSize' + str(args.imgMaxSize)
-    args.id += '-padding_constant' + str(args.padding_constant)
-    args.id += '-segm_downsampling_rate' + str(args.segm_downsampling_rate)
-    args.id += '-lr_encoder' + str(args.lr_encoder)
-    args.id += '-lr_decoder' + str(args.lr_decoder)
+    args.id += '-paddingConst' + str(args.padding_constant)
+    args.id += '-segmDownsampleRate' + str(args.segm_downsampling_rate)
+    args.id += '-LR_encoder' + str(args.lr_encoder)
+    args.id += '-LR_decoder' + str(args.lr_decoder)
     args.id += '-epoch' + str(args.num_epoch)
     args.id += '-decay' + str(args.weight_decay)
-    args.id += '-fix_bn' + str(args.fix_bn)
+    args.id += '-fixBN' + str(args.fix_bn)
     print('Model ID: {}'.format(args.id))
 
     args.ckpt = os.path.join(args.ckpt, args.id)

utils.py

@@ -1,4 +1,4 @@
-import torch
+# import torch
 import numpy as np
 
 
@@ -80,7 +80,7 @@ def unique(ar, return_index=False, return_inverse=False, return_counts=False):
     return ret
 
 
-def colorEncode(labelmap, colors):
+def colorEncode(labelmap, colors, mode='BGR'):
     labelmap = labelmap.astype('int')
     labelmap_rgb = np.zeros((labelmap.shape[0], labelmap.shape[1], 3),
                             dtype=np.uint8)
@@ -90,7 +90,11 @@ def colorEncode(labelmap, colors):
         labelmap_rgb += (labelmap == label)[:, :, np.newaxis] * \
             np.tile(colors[label],
                     (labelmap.shape[0], labelmap.shape[1], 1))
-    return labelmap_rgb
+
+    if mode == 'BGR':
+        return labelmap_rgb[:, :, ::-1]
+    else:
+        return labelmap_rgb
 
 
 def accuracy(preds, label):
@@ -100,22 +104,25 @@ def accuracy(preds, label):
     acc = float(acc_sum) / (valid_sum + 1e-10)
     return acc, valid_sum
 
+
 def intersectionAndUnion(imPred, imLab, numClass):
     imPred = np.asarray(imPred).copy()
     imLab = np.asarray(imLab).copy()
 
-    imPred += 1; imLab += 1
-    # Remove classes from unlabeled pixels in gt image. 
+    imPred += 1
+    imLab += 1
+    # Remove classes from unlabeled pixels in gt image.
     # We should not penalize detections in unlabeled portions of the image.
-    imPred = imPred * (imLab>0)
+    imPred = imPred * (imLab > 0)
 
     # Compute area intersection:
-    intersection = imPred * (imPred==imLab)
-    (area_intersection,_) = np.histogram(intersection, bins=numClass, range=(1, numClass))
+    intersection = imPred * (imPred == imLab)
+    (area_intersection, _) = np.histogram(
+        intersection, bins=numClass, range=(1, numClass))
 
     # Compute area union:
-    (area_pred,_) = np.histogram(imPred, bins=numClass, range=(1, numClass))
-    (area_lab,_) = np.histogram(imLab, bins=numClass, range=(1, numClass))
+    (area_pred, _) = np.histogram(imPred, bins=numClass, range=(1, numClass))
+    (area_lab, _) = np.histogram(imLab, bins=numClass, range=(1, numClass))
     area_union = area_pred + area_lab - area_intersection
-   
-    return (area_intersection, area_union) 
+
+    return (area_intersection, area_union)