modify some content

Author: wbw520

configs.py

@@ -5,18 +5,19 @@ def get_args_parser():
     parser = argparse.ArgumentParser(description="PSP-Net Network", add_help=False)
 
     # train settings
-    parser.add_argument("--dataset", type=str, default="cityscapes")
+    parser.add_argument("--dataset", type=str, default="facade")
     parser.add_argument("--model_name", type=str, default="Segmenter")
-    parser.add_argument("--pre_model", type=str, default="mae_pre_epoch99_crop640_patch8_ed768_depth12_head12.pt")
-    parser.add_argument("--batch_size", type=int, default=1,
+    parser.add_argument("--pre_model", type=str, default="ViT-B_16.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/cityscapes",
+    parser.add_argument("--root", type=str, default="/home/wangbowen/DATA/",
                         help="Path to the directory containing the image list.")
+    parser.add_argument("--setting_size", type=int, default=[1024, 2048],
+                        help="original size of data set image.")
     parser.add_argument("--crop_size", type=int, default=[640, 640],
                         help="crop size for training and inference slice.")
     parser.add_argument("--stride_rate", type=float, default=0.5, help="stride ratio.")
-    parser.add_argument("--num_epoch", type=int, default=200, help="Number of training steps.")
-    parser.add_argument("--num_classes", type=int, default=19, help="Number of class for dataset.")
+    parser.add_argument("--num_epoch", type=int, default=60, help="Number of training steps.")
     parser.add_argument('--accum_iter', default=1, type=int,
                         help='Accumulate gradient iterations (for increasing the effective batch size under memory constraints)')
 
@@ -26,7 +27,7 @@ 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=8, help="define the patch size.")
+    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("--decoder_embed_dim", type=int, default=512, help="dimension for decoder.")
     parser.add_argument("--encoder_depth", type=int, default=12, help="depth for encoder.")
@@ -38,10 +39,11 @@ def get_args_parser():
     parser.add_argument("--save_summary", type=str, default="save_model")
     parser.add_argument("--print_freq", type=str, default=5, help="print frequency.")
     parser.add_argument('--output_dir', default='save_model/', help='path where to save, empty for no saving')
+    parser.add_argument("--use_ignore", type=bool, default=False)
 
     # # distributed training parameters
     parser.add_argument('--num_workers', default=4, type=int)
-    parser.add_argument("--device", type=str, default='cuda',
+    parser.add_argument("--device", type=str, default='cuda:1',
                         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)

data/facade.py

@@ -2,19 +2,79 @@
 from PIL import Image
 import numpy as np
 from utils.base_tools import get_name
+import json
+from sklearn.model_selection import train_test_split
+import cv2
 
 
 ignore_label = 255
+num_classes = 10
+colors = {0: [0, 0, 0], 1: [70, 70, 70], 2: [250, 170, 30], 3: [70, 130, 180], 4: [0, 60, 100], 5: [153, 153, 153],
+          6: [107, 142, 35], 7: [255, 0, 0], 8: [0, 0, 142], 9: [220, 220, 0]}
+
+
+class PolygonTrans():
+    def __init__(self):
+        self.binary = {"building": 1, "window": 2, "sky": 3, "roof": 4, "door": 5, "tree": 6, "people": 7, "car": 8, "sign": 9}
+        self.overlap_order = ["sky", "building", "roof", "door", "window", "tree", "people", "car", "sign"]
+
+    def polygon2mask(self, img_size, polygons, rectangles):
+        mask = np.zeros(img_size, dtype=np.uint8)
+        for cat in self.overlap_order:
+            polygon = polygons[cat]
+            cv2.fillPoly(mask, polygon, color=self.binary[cat])
+            rectangle = rectangles[cat]
+            for ret in rectangle:
+                x1, y1 = ret[0]
+                x2, y2 = ret[1]
+                mask[y1:y2, x1:x2] = self.binary[cat]
+        return mask
+
+    # translate label_id to color img
+    def id2trainId(self, label):
+        w, h = label.shape
+        label_copy = np.zeros((w, h, 3), dtype=np.uint8)
+        for index, color in colors.items():
+            label_copy[label == index] = color
+        return label_copy.astype(np.uint8)
+
+
+def read_json(file_name):
+    record = {"building": [], "window": [], "sky": [], "roof": [], "door": [], "tree": [], "people": [], "car": [], "sign": []}
+    record_rectangle = {"building": [], "window": [], "sky": [], "roof": [], "door": [], "tree": [], "people": [], "car": [], "sign": []}
+    with open(file_name, "r") as load_polygon:
+        data = json.load(load_polygon)
+
+    data = data["shapes"]
+    for item in data:
+        label = item["label"]
+        points = item["points"]
+        shape = item["shape_type"]
+        if label not in record:
+            continue
+
+        if shape == "rectangle":
+            record_rectangle[label].append(np.array(points, dtype=np.int32))
+        else:
+            record[label].append(np.array(points, dtype=np.int32))
+    return record, record_rectangle
 
 
 def prepare_facade_data(args):
-    items = get_name(args.root + "/translated_data/images")
+    roots = args.root + "Facade/wang_translated_data/"
+    items = get_name(roots + "images", mode_folder=False)
+    record = []
+    for item in items:
+        record.append([roots + "images/" + item, roots + "binary_mask/" + item])
+
+    train, val = train_test_split(record, train_size=0.9, random_state=1)
+    return {"train": train, "val": val}
 
 
 class Facade(torch.utils.data.Dataset):
     def __init__(self, args, mode, joint_transform=None, standard_transform=None):
         self.args = args
-        self.imgs = ""
+        self.imgs = prepare_facade_data(args)[mode]
         if len(self.imgs) == 0:
             raise RuntimeError('Found 0 images, please check the data set')
 

data/facade_data_generation.py

@@ -0,0 +1,58 @@
+import cv2
+from utils.base_tools import get_name
+from data.facade import read_json, PolygonTrans
+import os
+from PIL import Image
+import shutil
+
+
+def main():
+    item_list = get_name(root + name, mode_folder=False)
+    image_list = []
+    for item in item_list:
+        name_ = item.split(".")[0]
+        if name_ not in image_list:
+            image_list.append(name_)
+
+    for img in image_list:
+        print(img)
+        json_root = root + name + "/" + img + ".json"
+        if not os.path.exists(json_root):
+            print("file not exist: ", json_root)
+            continue
+
+        polygons, rectangles = read_json(json_root)
+        if "IMG_E" in img:
+            suffix = ".JPG"
+        else:
+            suffix = ".jpg"
+
+        image = cv2.imread(root + name + "/" + img + suffix, cv2.IMREAD_COLOR)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        h, w, c = image.shape
+
+        PT = PolygonTrans()
+        mask = PT.polygon2mask((h, w), polygons, rectangles)
+        color_map = PT.id2trainId(mask)
+
+        image = Image.fromarray(image)
+        mask = Image.fromarray(mask)
+        color_map = Image.fromarray(color_map)
+
+        image.save(root_img + "/" + img + ".jpg")
+        mask.save(root_binary_mask + "/" + img + ".jpg")
+        color_map.save(root_color_mask + "/" + img + ".jpg")
+
+
+if __name__ == '__main__':
+    root = "/home/wangbowen/DATA/Facade/"
+    name = "zhao"
+    use_predict = ""
+    shutil.rmtree(root + name + "translated_data", ignore_errors=True)
+    root_img = root + name + "_translated_data/images"
+    root_color_mask = root + name + "_translated_data/color_mask"
+    root_binary_mask = root + name + "_translated_data/binary_mask"
+    os.makedirs(root_img, exist_ok=True)
+    os.makedirs(root_color_mask, exist_ok=True)
+    os.makedirs(root_binary_mask, exist_ok=True)
+    main()

data/get_data_set.py

@@ -0,0 +1,28 @@
+from data import cityscapes
+from data import facade
+from data.loader_tools import get_joint_transformations, get_standard_transformations, get_joint_transformations_val
+
+
+def get_data(args):
+    joint_transformations = get_joint_transformations(args)
+    joint_transformations_val = get_joint_transformations_val(args)
+    standard_transformations = get_standard_transformations()
+
+    if args.dataset == "cityscapes":
+        train_set = cityscapes.CityScapes(args, 'fine', 'train', joint_transform=joint_transformations,
+                                          standard_transform=standard_transformations)
+        val_set = cityscapes.CityScapes(args, 'fine', 'val', joint_transform=None,
+                                        standard_transform=standard_transformations)
+        ignore_index = cityscapes.ignore_label
+        args.num_classes = cityscapes.num_classes
+    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,
+                                            standard_transform=standard_transformations)
+        ignore_index = facade.ignore_label
+        args.num_classes = facade.num_classes
+    else:
+        raise "dataset name error !"
+
+    return train_set, val_set, ignore_index

data/loader_tools.py

@@ -4,14 +4,22 @@
 
 def get_joint_transformations(args):
     aug_list = [
-                joint_transforms.RandomSized(),
+                joint_transforms.Resize(args),
+                # joint_transforms.RandomSized(),
                 # joint_transforms.RandomRotate(10),
                 joint_transforms.RandomCrop(args.crop_size),
                 joint_transforms.RandomHorizontallyFlip(),
                 ]
     return joint_transforms.Compose(aug_list)
 
 
+def get_joint_transformations_val(args):
+    aug_list = [
+                joint_transforms.Resize(args),
+                ]
+    return joint_transforms.Compose(aug_list)
+
+
 def get_standard_transformations():
     mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
     return standard_transforms.Compose([

data/transforms.py

@@ -152,3 +152,14 @@ def __call__(self, img, mask):
 
         # return self.crop(*self.scale(img, mask))
         return img, mask
+
+
+class Resize(object):
+    def __init__(self, args):
+        self.h = args.setting_size[0]
+        self.w = args.setting_size[1]
+
+    def __call__(self, img, mask):
+        img, mask = img.resize((self.w, self.h), Image.BILINEAR), mask.resize((self.w, self.h), Image.NEAREST)
+
+        return img, mask

inference.py

@@ -1,6 +1,5 @@
 import argparse
 import torch.backends.cudnn as cudnn
-from data import cityscapes
 import torch
 from PIL import Image
 from data.loader_tools import get_standard_transformations
@@ -9,6 +8,7 @@
 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
@@ -30,8 +30,8 @@ def show_single(image, location=None, save=False):
     plt.axis('off')
     plt.subplots_adjust(top=1, bottom=0, left=0, right=1, hspace=0, wspace=0)
     plt.margins(0, 0)
-    if save:
-        plt.savefig("demo/" + img_name, bbox_inches='tight', pad_inches=0)
+    # if save:
+    #     plt.savefig("demo/" + img_name, bbox_inches='tight', pad_inches=0)
     plt.show()
 
 
@@ -42,29 +42,23 @@ def main():
     cudnn.benchmark = True
     model = model_generation(args)
     model.to(device)
-    checkpoint = torch.load(args.output_dir + "/48_epoch_PSPNet.pt", map_location="cuda:0")
+    checkpoint = torch.load(args.output_dir + args.dataset + "_" + args.model_name + ".pt", map_location="cuda:1")
     model.load_state_dict(checkpoint, strict=True)
     model.eval()
 
     standard_transformations = get_standard_transformations()
     img = Image.open(img_path).convert('RGB')
+    img = img.resize((args.setting_size[1], args.setting_size[0]), Image.BILINEAR)
     img = standard_transformations(img).to(device, dtype=torch.float32)
     pred, full_pred = inference_sliding(args, model, img.unsqueeze(0))
-    color_img = ColorTransition().recover(torch.squeeze(pred, dim=0))
+    color_img = PolygonTrans().id2trainId(torch.squeeze(pred, dim=0).cpu().detach().numpy())
     show_single(color_img, save=True)
 
 
 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()
-    if args.dataset == "cityscapes":
-        args.num_classes = cityscapes.num_classes
-    else:
-        args.num_classes = 1
-
-    root = "/home/wangbowen/streetview/"
-    imgs = get_name(root, mode_folder=False)
-    for img_name in imgs:
-        img_path = root + img_name
-        main()
+    args.num_classes = 10
+    img_path = "/home/wangbowen/DATA/Facade/zhao_translated_data/images/IMG_1282.jpg"
+    main()