add evaluation

Author: wbw520

configs.py

@@ -7,8 +7,8 @@ def get_args_parser():
     # train settings
     parser.add_argument("--dataset", type=str, default="facade")
     parser.add_argument("--model_name", type=str, default="PSPNet")
-    parser.add_argument("--pre_model", type=str, default="ViT-B_16.npz")
-    parser.add_argument("--batch_size", type=int, default=2,
+    parser.add_argument("--pre_model", type=str, default="ViT-B_8.npz")
+    parser.add_argument("--batch_size", type=int, default=4,
                         help="Number of images sent to the network in one step.")
     parser.add_argument("--root", type=str, default="/home/wangbowen/DATA/",
                         help="Path to the directory containing the image list.")
@@ -27,12 +27,9 @@ def get_args_parser():
     parser.add_argument("--weight_decay", type=float, default=1e-4, help="weight decay.")
 
     # VIT settings
-    parser.add_argument("--patch_size", type=int, default=16, help="define the patch size.")
-    parser.add_argument("--encoder_embed_dim", type=int, default=768, help="dimension for encoder.")
+    parser.add_argument("--encoder", type=str, default="vit_base_patch8", help="name for encoder")
     parser.add_argument("--decoder_embed_dim", type=int, default=512, help="dimension for decoder.")
-    parser.add_argument("--encoder_depth", type=int, default=12, help="depth for encoder.")
     parser.add_argument("--decoder_depth", type=int, default=2, help="depth for decoder.")
-    parser.add_argument("--encoder_num_head", type=int, default=12, help="head number for encoder.")
     parser.add_argument("--decoder_num_head", type=int, default=8, help="head number for decoder.")
 
     # other settings
@@ -42,11 +39,10 @@ def get_args_parser():
     parser.add_argument("--use_ignore", type=bool, default=False)
 
     # # distributed training parameters
-    parser.add_argument('--num_workers', default=0, type=int)
-    parser.add_argument("--device", type=str, default='cuda:1',
-                        help="choose gpu device.")
+    parser.add_argument('--num_workers', default=4, type=int)
+    parser.add_argument("--device", type=str, default='cuda', help="choose gpu device.")
     parser.add_argument('--world_size', default=1, type=int, help='number of distributed processes')
-    parser.add_argument("--local_rank", type=int)
+    parser.add_argument("--local_rank", default=-1, type=int)
     parser.add_argument('--dist_on_itp', action='store_true')
     parser.add_argument('--dist_url', default='env://', help='url used to set up distributed training')
 

data/facade.py

@@ -69,7 +69,7 @@ def prepare_facade_data(args):
         record.append([roots + "images/" + item, roots + "binary_mask/" + item])
 
     train, other = train_test_split(record, train_size=0.8, random_state=1)
-    val, test = train_test_split(record, train_size=0.5, random_state=1)
+    val, test = train_test_split(other, train_size=0.5, random_state=1)
     return {"train": train, "val": val, "test": test}
 
 
@@ -87,7 +87,6 @@ def __init__(self, args, mode, joint_transform=None, standard_transform=None):
 
     def __getitem__(self, index):
         img_path, mask_path = self.imgs[index]
-        print(img_path)
         img, mask = Image.open(img_path).convert('RGB'), Image.open(mask_path)
 
         mask = np.array(mask)

data/facade_data_generation.py

@@ -34,8 +34,6 @@ def main():
 
         PT = PolygonTrans()
         mask = PT.polygon2mask((h, w), polygons, rectangles)
-        if np.sum(mask == 10):
-            print("--------------------")
 
         color_map = PT.id2trainId(mask)
 

data/get_data_set.py

@@ -3,7 +3,7 @@
 from data.loader_tools import get_joint_transformations, get_standard_transformations, get_joint_transformations_val
 
 
-def get_data(args):
+def get_data(args, evaluation_setting=None):
     joint_transformations = get_joint_transformations(args)
     joint_transformations_val = get_joint_transformations_val(args)
     standard_transformations = get_standard_transformations()
@@ -18,7 +18,13 @@ def get_data(args):
     elif args.dataset == "facade":
         train_set = facade.Facade(args, 'train', joint_transform=joint_transformations,
                                               standard_transform=standard_transformations)
-        val_set = facade.Facade(args, 'val', joint_transform=joint_transformations_val,
+
+        if evaluation_setting is not None:
+            current_set = "test"
+        else:
+            current_set = "val"
+
+        val_set = facade.Facade(args, 'test', joint_transform=joint_transformations_val,
                                             standard_transform=standard_transformations)
         ignore_index = facade.ignore_label
         args.num_classes = facade.num_classes

evaluation.py

@@ -0,0 +1,42 @@
+import argparse
+import torch.backends.cudnn as cudnn
+import torch
+from data.get_data_set import get_data
+import utils2.misc as misc
+from configs import get_args_parser
+from utils.engine import evaluation_none_training
+from torch.utils.data import DataLoader
+from model.get_model import model_generation
+import os
+
+
+def main():
+    # distribution
+    misc.init_distributed_mode(args)
+    device = torch.device(args.device)
+    cudnn.benchmark = True
+    train_set, test_set, ignore_index = get_data(args)
+    model = model_generation(args)
+    model.to(device)
+
+    if args.model_name == "Segmenter":
+        save_name = args.model_name + "_" + args.encoder
+    else:
+        save_name = args.model_name
+
+    checkpoint = torch.load(args.output_dir + args.dataset + "_" + save_name + ".pt", map_location="cuda:0")
+    model.load_state_dict(checkpoint, strict=True)
+    model.eval()
+    print("load trained model finished.")
+
+    sampler_val = torch.utils.data.SequentialSampler(test_set)
+    val_loader = DataLoader(test_set, batch_size=args.batch_size, sampler=sampler_val, num_workers=args.num_workers, shuffle=False)
+    evaluation_none_training(args, model, val_loader, device)
+
+
+if __name__ == '__main__':
+    os.makedirs('demo/', exist_ok=True)
+    parser = argparse.ArgumentParser('model training and evaluation script', parents=[get_args_parser()])
+    args = parser.parse_args()
+    img_path = "/home/wangbowen/DATA/Facade/translated_data/images/IMG_1287.png"
+    main()

inference.py

@@ -1,27 +1,20 @@
 import argparse
 import torch.backends.cudnn as cudnn
 import torch
+from data.get_data_set import get_data
 from PIL import Image
 from data.loader_tools import get_standard_transformations
 import utils2.misc as misc
+from utils.base_tools import get_name
 from configs import get_args_parser
 from model.get_model import model_generation
 from utils.engine import inference_sliding
-from data.cityscapes import ColorTransition
 from data.facade import PolygonTrans
 import matplotlib.pyplot as plt
 import numpy as np
 import os
 
 
-def get_name(root, mode_folder=True):
-    for root, dirs, file in os.walk(root):
-        if mode_folder:
-            return dirs
-        else:
-            return file
-
-
 def show_single(image, location=None, save=False):
     # show single image
     image = np.array(image, dtype=np.uint8)
@@ -40,9 +33,16 @@ def main():
     misc.init_distributed_mode(args)
     device = torch.device(args.device)
     cudnn.benchmark = True
+    train_set, val_set, ignore_index = get_data(args)
     model = model_generation(args)
     model.to(device)
-    checkpoint = torch.load(args.output_dir + args.dataset + "_" + args.model_name + ".pt", map_location="cuda:1")
+
+    if args.model_name == "Segmenter":
+        save_name = args.model_name + "_" + args.encoder
+    else:
+        save_name = args.model_name
+
+    checkpoint = torch.load(args.output_dir + args.dataset + "_" + save_name + ".pt", map_location="cuda:0")
     model.load_state_dict(checkpoint, strict=True)
     model.eval()
 
@@ -59,6 +59,5 @@ def main():
     os.makedirs('demo/', exist_ok=True)
     parser = argparse.ArgumentParser('model training and evaluation script', parents=[get_args_parser()])
     args = parser.parse_args()
-    args.num_classes = 10
-    img_path = "/home/wangbowen/DATA/Facade/zhao_translated_data/images/IMG_1282.jpg"
+    img_path = "/home/wangbowen/DATA/Facade/translated_data/images/IMG_1287.png"
     main()

model/segmenter.py

@@ -1,9 +1,10 @@
-from model.vit_model import vit_encoder
+import model.vit_model as vit
 from model.segmenter_decoder import sg_vit_mask_decoder
 import torch.nn as nn
 import torch
 import os
 import torch.nn.functional as F
+import re
 from timm.models.helpers import load_checkpoint
 
 
@@ -43,9 +44,25 @@ def get_attention_map_dec(self, im, layer_id):
         return self.decoder.get_attention_map(x, layer_id)
 
 
-def create_segmenter(args):
-    encoder = vit_encoder(img_size=args.crop_size[0], patch_size=args.patch_size, embed_dim=args.encoder_embed_dim, depth=args.encoder_depth, num_heads=args.encoder_num_head)
+def set_decoder_parameter(name):
+    if "small" in name:
+        encoder_dim = 384
+    elif "base" in name:
+        encoder_dim = 768
+    elif "large" in name:
+        encoder_dim = 1024
+    elif "huge" in name:
+        encoder_dim = 1280
+    else:
+        raise "type of encoder is not defined."
+
+    patch_size = re.findall("\d+", name)
 
+    return encoder_dim, int(patch_size[0])
+
+
+def create_segmenter(args):
+    encoder = vit.__dict__[args.encoder](img_size=args.crop_size[0])
     if "mae" not in args.pre_model:
         print("load pre-model trained by ImageNet")
         load_checkpoint(encoder, args.output_dir + args.pre_model)
@@ -64,7 +81,8 @@ def create_segmenter(args):
 
     print("load pre-trained weight from: ", args.pre_model)
 
-    decoder = sg_vit_mask_decoder(patch_size=args.patch_size, encoder_embed_dim=args.encoder_embed_dim,
+    encoder_embed_dim, patch_size = set_decoder_parameter(args.encoder)
+    decoder = sg_vit_mask_decoder(patch_size=patch_size, encoder_embed_dim=encoder_embed_dim,
                                   decoder_embed_dim=args.decoder_embed_dim, decoder_depth=args.decoder_depth, decoder_num_heads=args.decoder_num_head, n_cls=args.num_classes)
     model = Segmenter(encoder, decoder)
 

model/vit_model.py

@@ -78,6 +78,28 @@ def resize_pos_embed(posemb, grid_old_shape, grid_new_shape, num_extra_tokens):
     return posemb
 
 
-def vit_encoder(**kwargs):
-    model = VisionTransformer(mlp_ratio=4, qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), num_classes=0, **kwargs)
+def vit_base_patch8(**kwargs):
+    model = VisionTransformer(patch_size=8, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
+                              norm_layer=partial(nn.LayerNorm, eps=1e-6), num_classes=0,
+                              **kwargs)
     return model
+
+
+def vit_base_patch16(**kwargs):
+    model = VisionTransformer(patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
+                              norm_layer=partial(nn.LayerNorm, eps=1e-6), num_classes=0,
+                              **kwargs)
+    return model
+
+
+def vit_base_patch32(**kwargs):
+    model = VisionTransformer(patch_size=32, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
+                              norm_layer=partial(nn.LayerNorm, eps=1e-6), num_classes=0,
+                              **kwargs)
+    return model
+
+
+# set recommended archs
+vit_base_patch8 = vit_base_patch8
+vit_base_patch16 = vit_base_patch16
+vit_base_patch32 = vit_base_patch32

train.py

@@ -28,7 +28,7 @@ def main():
     model_without_ddp = model
 
     if args.distributed:
-        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
+        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=False)
         model_without_ddp = model.module
         sampler_train = DistributedSampler(train_set)
         sampler_val = DistributedSampler(val_set, shuffle=False)

utils/engine.py

@@ -86,7 +86,11 @@ def evaluation(args, best_record, epoch, model, model_without_ddp, val_loader, c
         best_record['acc_cls'] = acc_cls
         best_record['mean_iou'] = mean_iou
         if args.output_dir:
-            torch.save(model_without_ddp.state_dict(), args.output_dir + args.dataset + "_" + args.model_name + ".pt")
+            if args.model_name == "Segmenter":
+                save_name = args.model_name + "_" + args.encoder
+            else:
+                save_name = args.model_name
+            torch.save(model_without_ddp.state_dict(), args.output_dir + args.dataset + "_" + save_name + ".pt")
 
     print('-----------------------------------------------------------------------------------------------------------')
     print('[epoch %d], [val loss %.5f], [acc %.5f], [acc_cls %.5f], [mean_iou %.5f]' % (
@@ -99,6 +103,25 @@ def evaluation(args, best_record, epoch, model, model_without_ddp, val_loader, c
     print('-----------------------------------------------------------------------------------------------------------')
 
 
+@torch.no_grad()
+def evaluation_none_training(args, model, val_loader, device):
+    model.eval()
+    iou = IouCal(args)
+    for i_batch, data in enumerate(val_loader):
+        if i_batch % 5 == 0:
+            print(str(i_batch) + "/" + str(len(val_loader)))
+        inputs = data["images"].to(device, dtype=torch.float32)
+        mask = data["masks"].to(device, dtype=torch.int64)
+
+        pred, full_pred = inference_sliding(args, model, inputs)
+        iou.evaluate(pred, mask)
+
+    acc, acc_cls, mean_iou = iou.iou_demo()
+    print("acc:", acc)
+    print("acc_cls", acc_cls)
+    print("mean_iou", mean_iou)
+
+
 @torch.no_grad()
 def inference_sliding(args, model, image):
     image_size = image.size()