├── .gitignore ├── LICENSE ├── README.md ├── data ├── __init__.py ├── base_data_loader.py ├── base_dataset.py ├── image_folder.py ├── nyuv2_dataset.py ├── scannetv2_dataset.py └── sunrgbd_dataset.py ├── datasets ├── calculate_class_weights.py ├── create_training_set.py ├── download_nyu2.sh ├── palette.txt ├── scannet │ ├── scannetv2_test.txt │ ├── scannetv2_train.txt │ ├── scannetv2_val.txt │ └── scannetv2_weigths.txt └── sunrgbd │ ├── class_weights │ └── splits.pkl ├── environment.yml ├── models ├── __init__.py ├── base_model.py ├── fusenet_model.py └── networks.py ├── options ├── __init__.py ├── base_options.py ├── test_options.py └── train_options.py ├── requirements.txt ├── test.py ├── train.py └── util ├── __init__.py ├── get_data.py ├── html.py ├── image_pool.py ├── util.py └── visualizer.py /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | checkpoints/ 3 | *.mat 4 | *.npy 5 | *.png 6 | */**/__pycache__ 7 | */*.pyc 8 | */**/*.pyc 9 | */**/**/*.pyc 10 | */**/**/**/*.pyc 11 | */**/**/**/**/*.pyc 12 | *~ 13 | .idea 14 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2017, Jun-Yan Zhu and Taesung Park 2 | All rights reserved. 3 | 4 | Redistribution and use in source and binary forms, with or without 5 | modification, are permitted provided that the following conditions are met: 6 | 7 | * Redistributions of source code must retain the above copyright notice, this 8 | list of conditions and the following disclaimer. 9 | 10 | * Redistributions in binary form must reproduce the above copyright notice, 11 | this list of conditions and the following disclaimer in the documentation 12 | and/or other materials provided with the distribution. 13 | 14 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 15 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 16 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 17 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 18 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 19 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 20 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 21 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 22 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 23 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 24 | 25 | 26 | --------------------------- LICENSE FOR pix2pix -------------------------------- 27 | BSD License 28 | 29 | For pix2pix software 30 | Copyright (c) 2016, Phillip Isola and Jun-Yan Zhu 31 | All rights reserved. 32 | 33 | Redistribution and use in source and binary forms, with or without 34 | modification, are permitted provided that the following conditions are met: 35 | 36 | * Redistributions of source code must retain the above copyright notice, this 37 | list of conditions and the following disclaimer. 38 | 39 | * Redistributions in binary form must reproduce the above copyright notice, 40 | this list of conditions and the following disclaimer in the documentation 41 | and/or other materials provided with the distribution. 42 | 43 | ----------------------------- LICENSE FOR DCGAN -------------------------------- 44 | BSD License 45 | 46 | For dcgan.torch software 47 | 48 | Copyright (c) 2015, Facebook, Inc. All rights reserved. 49 | 50 | Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: 51 | 52 | Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. 53 | 54 | Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. 55 | 56 | Neither the name Facebook nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. 57 | 58 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 59 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | # [FuseNet](https://github.com/tum-vision/fusenet) implementation in PyTorch 5 | 6 | This is the PyTorch implementation for FuseNet, developed based on [Pix2Pix](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix) code. 7 | 8 | ## Prerequisites 9 | - Linux 10 | - Python 3.7.0 11 | - CPU or NVIDIA GPU + CUDA CuDNN 12 | 13 | ## Getting Started 14 | ### Installation 15 | - Install PyTorch 0.4.1.post2 and dependencies from http://pytorch.org 16 | - Clone this repo: 17 | ```bash 18 | git clone https://github.com/MehmetAygun/fusenet-pytorch 19 | cd fusenet-pytorch 20 | pip install -r requirements.txt 21 | ``` 22 | ## Dataset preparation 23 | ### sunrgbd dataset 24 | - Download and untar the [preprocessed sunrgbd](https://vision.in.tum.de/webarchive/hazirbas/fusenet-pytorch/sun/sunrgbd.tar.gz) dataset under ```/datasets/sunrgbd``` 25 | 26 | ### nyuv2 dataset 27 | - Download the dataset and create the training set 28 | ```bash 29 | cd datasets 30 | sh download_nyuv2.sh 31 | python create_training_set.py 32 | ``` 33 | ### scannetv2 dataset 34 | - Download the ```scannet_frames_25k``` and ```scannet_frames_test``` under ```/datasets/scannet/tasks/``` 35 | 36 | ## FuseNet train/test 37 | 38 | ### visdom visualization 39 | - To view training errors and loss plots, set `--display_id 1`, run `python -m visdom.server` and click the URL http://localhost:8097 40 | - Checkpoints are saved under `./checkpoints/sunrgbd/` 41 | 42 | ### train & test on sunrgbd 43 | ```bash 44 | python train.py --dataroot datasets/sunrgbd --dataset sunrgbd --name sunrgbd 45 | 46 | python test.py --dataroot datasets/sunrgbd --dataset sunrgbd --name sunrgbd --epoch 400 47 | ``` 48 | 49 | ### train & test on nyuv2 50 | ```bash 51 | python train.py --dataroot datasets/nyuv2 --dataset nyuv2 --name nyuv2 52 | 53 | python test.py --dataroot datasets/nyuv2 --dataset nyuv2 --name nyuv2 --epoch 400 54 | ``` 55 | 56 | ### train & val & test on scannetv2 57 | ```bash 58 | python train.py --dataroot datasets/scannet/tasks/scannet_frames_25k --dataset scannetv2 \ 59 | --name scannetv2 60 | 61 | python test.py --dataroot datasets/scannet/tasks/scannet_frames_25k --dataset scannetv2 \ 62 | --name scannetv2 --epoch 380 --phase val 63 | 64 | python test.py --dataroot datasets/scannet/tasks/scannet_frames_test --dataset scannetv2 \ 65 | --name scannetv2 --epoch 380 --phase test 66 | ``` 67 | 68 | ## Results 69 | * We use the training scheme defined in FuseNet 70 | * Loss is weighted for SUNRGBD dataset 71 | * Learning rate is set to 0.01 for NYUv2 dataset 72 | * Results can be improved with a hyper-parameter search 73 | * Results on the scannetv2-test (w/o class-weighted loss) can be found [here](http://kaldir.vc.in.tum.de/scannet_benchmark/result_details?id=67) 74 | 75 | 76 | 77 | 79 | 81 | 83 | 85 | 87 |
Dataset FuseNet-SF5 (CAFFE) FuseNet-SF5 78 |
overall mean iou overall mean iou 80 |
sunrgbd 76.30 48.30 37.30 75.41 46.48 35.69 82 |
nyuv2 66.00 43.40 32.70 68.76 46.42 35.48 84 |
scannetv2-val -- -- -- 76.32 55.84 44.12 86 |
88 | scannetv2-cls_weighted-val -- -- -- 76.26 55.74 44.40 89 |
90 | 91 | | scannetv2-test | avg iou | bathtub | bed | bookshelf | cabinet | chair | counter | curtain | desk | door | floor | other furniture | picture | refrigerator | shower curtain | sink | sofa | table | toilet | wall | window | 92 | |-----------------|---------|---------|------|-----------|---------|-------|---------|---------|------|------|-------|----------------|---------|--------------|----------------|------|------|-------|--------|------|--------| 93 | | no-cls_weighted | 52.1 | 59.1 | 68.2 | 22.0 | 48.8 | 27.9 | 34.4 | 61.0 | 46.1 | 47.5 | 91.0 | 29.3 | 44.7 | 51.2 | 39.7 | 61.8 | 56.7 | 45.2 | 73.4 | 78.2 | 56.6 | 94 | | cls_weighted | 53.5 | 57.0 | 68.1 | 18.2 | 51.2 | 29.0 | 43.1 | 65.9 | 50.4 | 49.5 | 90.3 | 30.8 | 42.8 | 52.3 | 36.5 | 67.6 | 62.1 | 47.0 | 76.2 | 77.9 | 54.1 | 95 | 96 | 97 | 98 | ## Citation 99 | ``` 100 | @inproceedings{hazirbas16fusenet, 101 | Title = {{FuseNet}: Incorporating Depth into Semantic Segmentation via Fusion-Based CNN Architecture}, 102 | Author = {Hazirbas, Caner and Ma, Lingni and Domokos, Csaba and Cremers, Daniel}, 103 | Booktitle = {Asian Conference on Computer Vision ({ACCV})}, 104 | Year = {2016}, 105 | Doi = {10.1007/978-3-319-54181-5_14}, 106 | Url = {https://github.com/tum-vision/fusenet} 107 | } 108 | ``` 109 | ## Acknowledgments 110 | Code is inspired by [pytorch-CycleGAN-and-pix2pix]((https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix)). 111 | -------------------------------------------------------------------------------- /data/__init__.py: -------------------------------------------------------------------------------- 1 | import importlib 2 | import torch.utils.data 3 | from data.base_data_loader import BaseDataLoader 4 | from data.base_dataset import BaseDataset 5 | import numpy 6 | 7 | def find_dataset_using_name(dataset_name): 8 | # Given the option --dataset_mode [datasetname], 9 | # the file "data/datasetname_dataset.py" 10 | # will be imported. 11 | dataset_filename = "data." + dataset_name + "_dataset" 12 | datasetlib = importlib.import_module(dataset_filename) 13 | 14 | # In the file, the class called DatasetNameDataset() will 15 | # be instantiated. It has to be a subclass of BaseDataset, 16 | # and it is case-insensitive. 17 | dataset = None 18 | target_dataset_name = dataset_name.replace('_', '') + 'dataset' 19 | for name, cls in datasetlib.__dict__.items(): 20 | if name.lower() == target_dataset_name.lower() \ 21 | and issubclass(cls, BaseDataset): 22 | dataset = cls 23 | 24 | if dataset is None: 25 | print("In %s.py, there should be a subclass of BaseDataset with class name that matches %s in lowercase." % (dataset_filename, target_dataset_name)) 26 | exit(0) 27 | 28 | return dataset 29 | 30 | 31 | def get_option_setter(dataset_name): 32 | dataset_class = find_dataset_using_name(dataset_name) 33 | return dataset_class.modify_commandline_options 34 | 35 | 36 | def create_dataset(opt): 37 | dataset = find_dataset_using_name(opt.dataset_mode) 38 | instance = dataset() 39 | instance.initialize(opt) 40 | print("dataset [%s] was created" % (instance.name())) 41 | return instance 42 | 43 | 44 | def CreateDataLoader(opt): 45 | data_loader = CustomDatasetDataLoader() 46 | data_loader.initialize(opt) 47 | return data_loader 48 | 49 | 50 | # Wrapper class of Dataset class that performs 51 | # multi-threaded data loading 52 | class CustomDatasetDataLoader(BaseDataLoader): 53 | def name(self): 54 | return 'CustomDatasetDataLoader' 55 | 56 | def initialize(self, opt): 57 | BaseDataLoader.initialize(self, opt) 58 | self.dataset = create_dataset(opt) 59 | self.dataloader = torch.utils.data.DataLoader( 60 | self.dataset, 61 | batch_size=opt.batch_size, 62 | shuffle=not opt.serial_batches, 63 | num_workers=int(opt.num_threads), 64 | worker_init_fn=lambda worker_id: numpy.random.seed(opt.seed + worker_id)) 65 | 66 | def load_data(self): 67 | return self 68 | 69 | def __len__(self): 70 | return min(len(self.dataset), self.opt.max_dataset_size) 71 | 72 | def __iter__(self): 73 | for i, data in enumerate(self.dataloader): 74 | if i * self.opt.batch_size >= self.opt.max_dataset_size: 75 | break 76 | yield data 77 | -------------------------------------------------------------------------------- /data/base_data_loader.py: -------------------------------------------------------------------------------- 1 | class BaseDataLoader(): 2 | def __init__(self): 3 | pass 4 | 5 | def initialize(self, opt): 6 | self.opt = opt 7 | pass 8 | 9 | def load_data(): 10 | return None 11 | -------------------------------------------------------------------------------- /data/base_dataset.py: -------------------------------------------------------------------------------- 1 | import torch.utils.data as data 2 | from PIL import Image 3 | import torchvision.transforms as transforms 4 | 5 | 6 | class BaseDataset(data.Dataset): 7 | def __init__(self): 8 | super(BaseDataset, self).__init__() 9 | 10 | def name(self): 11 | return 'BaseDataset' 12 | 13 | @staticmethod 14 | def modify_commandline_options(parser, is_train): 15 | return parser 16 | 17 | def initialize(self, opt): 18 | pass 19 | 20 | def __len__(self): 21 | return 0 22 | 23 | 24 | def get_transform(opt): 25 | transform_list = [] 26 | if opt.resize_or_crop == 'resize_and_crop': 27 | osize = [opt.loadSize, opt.loadSize] 28 | transform_list.append(transforms.Resize(osize, Image.BICUBIC)) 29 | transform_list.append(transforms.RandomCrop(opt.fineSize)) 30 | elif opt.resize_or_crop == 'crop': 31 | transform_list.append(transforms.RandomCrop(opt.fineSize)) 32 | elif opt.resize_or_crop == 'scale_width': 33 | transform_list.append(transforms.Lambda( 34 | lambda img: __scale_width(img, opt.fineSize))) 35 | elif opt.resize_or_crop == 'scale_width_and_crop': 36 | transform_list.append(transforms.Lambda( 37 | lambda img: __scale_width(img, opt.loadSize))) 38 | transform_list.append(transforms.RandomCrop(opt.fineSize)) 39 | elif opt.resize_or_crop == 'none': 40 | transform_list.append(transforms.Lambda( 41 | lambda img: __adjust(img))) 42 | else: 43 | raise ValueError('--resize_or_crop %s is not a valid option.' % opt.resize_or_crop) 44 | 45 | if opt.isTrain and not opt.no_flip: 46 | transform_list.append(transforms.RandomHorizontalFlip()) 47 | 48 | transform_list += [transforms.ToTensor(), 49 | transforms.Normalize((0.5, 0.5, 0.5), 50 | (0.5, 0.5, 0.5))] 51 | return transforms.Compose(transform_list) 52 | 53 | 54 | # just modify the width and height to be multiple of 4 55 | def __adjust(img): 56 | ow, oh = img.size 57 | 58 | # the size needs to be a multiple of this number, 59 | # because going through generator network may change img size 60 | # and eventually cause size mismatch error 61 | mult = 4 62 | if ow % mult == 0 and oh % mult == 0: 63 | return img 64 | w = (ow - 1) // mult 65 | w = (w + 1) * mult 66 | h = (oh - 1) // mult 67 | h = (h + 1) * mult 68 | 69 | if ow != w or oh != h: 70 | __print_size_warning(ow, oh, w, h) 71 | 72 | return img.resize((w, h), Image.BICUBIC) 73 | 74 | 75 | def __scale_width(img, target_width): 76 | ow, oh = img.size 77 | 78 | # the size needs to be a multiple of this number, 79 | # because going through generator network may change img size 80 | # and eventually cause size mismatch error 81 | mult = 4 82 | assert target_width % mult == 0, "the target width needs to be multiple of %d." % mult 83 | if (ow == target_width and oh % mult == 0): 84 | return img 85 | w = target_width 86 | target_height = int(target_width * oh / ow) 87 | m = (target_height - 1) // mult 88 | h = (m + 1) * mult 89 | 90 | if target_height != h: 91 | __print_size_warning(target_width, target_height, w, h) 92 | 93 | return img.resize((w, h), Image.BICUBIC) 94 | 95 | 96 | def __print_size_warning(ow, oh, w, h): 97 | if not hasattr(__print_size_warning, 'has_printed'): 98 | print("The image size needs to be a multiple of 4. " 99 | "The loaded image size was (%d, %d), so it was adjusted to " 100 | "(%d, %d). This adjustment will be done to all images " 101 | "whose sizes are not multiples of 4" % (ow, oh, w, h)) 102 | __print_size_warning.has_printed = True 103 | -------------------------------------------------------------------------------- /data/image_folder.py: -------------------------------------------------------------------------------- 1 | ############################################################################### 2 | # Code from 3 | # https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py 4 | # Modified the original code so that it also loads images from the current 5 | # directory as well as the subdirectories 6 | ############################################################################### 7 | 8 | import torch.utils.data as data 9 | 10 | from PIL import Image 11 | import os 12 | import os.path 13 | 14 | IMG_EXTENSIONS = [ 15 | '.jpg', '.JPG', '.jpeg', '.JPEG', 16 | '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', 17 | ] 18 | 19 | 20 | def is_image_file(filename): 21 | return any(filename.endswith(extension) for extension in IMG_EXTENSIONS) 22 | 23 | 24 | def make_dataset(dir): 25 | images = [] 26 | assert os.path.isdir(dir), '%s is not a valid directory' % dir 27 | 28 | for root, _, fnames in sorted(os.walk(dir)): 29 | for fname in fnames: 30 | if is_image_file(fname): 31 | path = os.path.join(root, fname) 32 | images.append(path) 33 | 34 | return images 35 | 36 | 37 | def default_loader(path): 38 | return Image.open(path).convert('RGB') 39 | 40 | 41 | class ImageFolder(data.Dataset): 42 | 43 | def __init__(self, root, transform=None, return_paths=False, 44 | loader=default_loader): 45 | imgs = make_dataset(root) 46 | if len(imgs) == 0: 47 | raise(RuntimeError("Found 0 images in: " + root + "\n" 48 | "Supported image extensions are: " + 49 | ",".join(IMG_EXTENSIONS))) 50 | 51 | self.root = root 52 | self.imgs = imgs 53 | self.transform = transform 54 | self.return_paths = return_paths 55 | self.loader = loader 56 | 57 | def __getitem__(self, index): 58 | path = self.imgs[index] 59 | img = self.loader(path) 60 | if self.transform is not None: 61 | img = self.transform(img) 62 | if self.return_paths: 63 | return img, path 64 | else: 65 | return img 66 | 67 | def __len__(self): 68 | return len(self.imgs) 69 | -------------------------------------------------------------------------------- /data/nyuv2_dataset.py: -------------------------------------------------------------------------------- 1 | import os.path 2 | import random 3 | import torchvision.transforms as transforms 4 | import torch 5 | from data.base_dataset import BaseDataset 6 | from data.image_folder import make_dataset 7 | from PIL import Image 8 | import numpy as np 9 | import scipy.io as sio 10 | 11 | class nyuv2dataset(BaseDataset): 12 | @staticmethod 13 | def modify_commandline_options(parser, is_train): 14 | return parser 15 | 16 | def initialize(self, opt): 17 | self.opt = opt 18 | self.batch_size = opt.batch_size 19 | self.root = opt.dataroot # path for nyu2.npy 20 | self.nyu2 = np.load("{}/{}".format(self.root,"nyuv2.npy"),encoding = 'latin1').tolist() 21 | splits = sio.loadmat("{}/{}".format(self.root,"splits.mat")) 22 | self.indexes = [x[0] - 1 for x in splits["trainNdxs"]] if opt.phase == "train" else [x[0] -1 for x in splits["testNdxs"]] 23 | self.num_labels = 41 24 | self.ignore_label = 0 25 | self.class_weights = None 26 | 27 | def __getitem__(self, index): 28 | index = self.indexes[index] 29 | rgb_image = np.array(self.nyu2["rgb_images"][index],dtype=np.uint8) 30 | depth_image = self.nyu2["depth_images"][index] 31 | depth_image = np.expand_dims(depth_image,axis=2) 32 | mask = np.array(self.nyu2["masks"][index],dtype=np.uint8) 33 | 34 | rgb_image = transforms.ToTensor()(rgb_image) 35 | depth_image = transforms.ToTensor()(depth_image) 36 | 37 | mask = torch.from_numpy(mask) 38 | mask = mask.type(torch.LongTensor) 39 | 40 | return {'rgb_image': rgb_image, 'depth_image': depth_image, 'mask': mask, 'path': str(index)+".png"} 41 | 42 | def __len__(self): 43 | return len(self.indexes) 44 | 45 | def name(self): 46 | return 'nyuv2dataset' 47 | -------------------------------------------------------------------------------- /data/scannetv2_dataset.py: -------------------------------------------------------------------------------- 1 | import os.path 2 | import glob 3 | import random 4 | import torchvision.transforms as transforms 5 | import torch 6 | from data.base_dataset import BaseDataset 7 | from data.image_folder import make_dataset 8 | from PIL import Image 9 | import numpy as np 10 | import scipy.io as sio 11 | import cv2 12 | 13 | class Scannetv2Dataset(BaseDataset): 14 | @staticmethod 15 | def modify_commandline_options(parser, is_train): 16 | return parser 17 | 18 | def initialize(self, opt): 19 | self.opt = opt 20 | self.batch_size = opt.batch_size 21 | self.root = opt.dataroot 22 | self.num_labels = 41 23 | self.ignore_label = 0 24 | self.class_weights = None 25 | with open('./datasets/scannet/scannetv2_weigths.txt') as f: 26 | weights = f.readlines() 27 | self.class_weights = torch.from_numpy(np.array([float(x.strip()) for x in weights])) 28 | self.class_weights = self.class_weights.type(torch.FloatTensor) 29 | 30 | with open('./datasets/scannet/scannetv2_{}.txt'.format(opt.phase)) as f: 31 | scans = f.readlines() 32 | self.scans = [x.strip() for x in scans] 33 | 34 | self.rgb_frames = [] 35 | self.depth_frames = [] 36 | self.masks = [] 37 | 38 | self.total_frames = 0 39 | for scan in self.scans: 40 | rgb_frames = glob.glob("{}/{}/color/*.jpg".format(self.root, scan)) 41 | depth_frames = glob.glob("{}/{}/depth/*.png".format(self.root, scan)) 42 | masks = glob.glob("{}/{}/label/*.png".format(self.root, scan)) 43 | if len(rgb_frames) == len(depth_frames): 44 | rgb_frames.sort() 45 | depth_frames.sort() 46 | masks.sort() 47 | self.total_frames += len(rgb_frames) 48 | self.rgb_frames.extend(rgb_frames) 49 | self.depth_frames.extend(depth_frames) 50 | self.masks.extend(masks) 51 | 52 | def __getitem__(self, index): 53 | 54 | size = (320,240) 55 | rgb_image = np.array(Image.open(self.rgb_frames[index])) 56 | rgb_image = cv2.resize(rgb_image, size, interpolation=cv2.INTER_LINEAR) 57 | depth_image = np.array(Image.open(self.depth_frames[index])) 58 | depth_image = cv2.resize(depth_image, size,interpolation=cv2.INTER_NEAREST).astype(np.float) 59 | depth_image = (depth_image - depth_image.min()) / (depth_image.max() - depth_image.min()) * 255 60 | depth_image = depth_image.astype(np.uint8) 61 | mask_fullsize = [] 62 | if self.masks: 63 | mask = np.array(Image.open(self.masks[index])) 64 | mask = cv2.resize(mask, size, interpolation=cv2.INTER_NEAREST) 65 | if self.opt.phase == "val": 66 | mask_fullsize = np.array(Image.open(self.masks[index])) 67 | else: # test phase 68 | mask = np.zeros((240, 320), dtype=int) 69 | mask_fullsize = np.zeros((968, 1296), dtype=int) 70 | 71 | rgb_image = transforms.ToTensor()(rgb_image) 72 | rgb_image = rgb_image.type(torch.FloatTensor) 73 | depth_image = transforms.ToTensor()(depth_image[:, :, np.newaxis]) 74 | depth_image = depth_image.type(torch.FloatTensor) 75 | 76 | mask = torch.from_numpy(mask) 77 | mask = mask.type(torch.LongTensor) 78 | 79 | return {'rgb_image': rgb_image, 'depth_image': depth_image, 80 | 'mask': mask, 'mask_fullsize': mask_fullsize, 81 | 'path': self.depth_frames[index].split('/')[-1], 'scan':self.depth_frames[index].split('/')[-3]} 82 | 83 | def __len__(self): 84 | return self.total_frames 85 | 86 | def name(self): 87 | return 'Scannetv2' 88 | -------------------------------------------------------------------------------- /data/sunrgbd_dataset.py: -------------------------------------------------------------------------------- 1 | import os.path 2 | import random 3 | import torchvision.transforms as transforms 4 | import torch 5 | from data.base_dataset import BaseDataset 6 | from data.image_folder import make_dataset 7 | from PIL import Image 8 | import numpy as np 9 | import pickle 10 | 11 | class sunrgbddataset(BaseDataset): 12 | @staticmethod 13 | def modify_commandline_options(parser, is_train): 14 | return parser 15 | 16 | def initialize(self, opt): 17 | self.opt = opt 18 | self.batch_size = opt.batch_size 19 | self.root = opt.dataroot # path for the dataset 20 | splits = pickle.load(open(os.path.join(self.root, "splits.pkl"), "rb"), encoding="latin1") 21 | self.indexes = splits["trainval"] if opt.phase == "train" else splits["test"] 22 | self.num_labels = 38 23 | self.ignore_label = 0 24 | self.class_weights = torch.from_numpy(np.loadtxt(os.path.join(opt.dataroot, "class_weights"), 25 | delimiter=',').astype(np.float32) 26 | ) 27 | assert(opt.resize_or_crop == 'none') 28 | 29 | def __getitem__(self, index): 30 | index = self.indexes[index] 31 | rgb_image = np.array(Image.open(os.path.join(self.root, "images-224", str(index)+".png"))) 32 | depth_image = np.array(Image.open(os.path.join(self.root, "depth-inpaint-u8-224", str(index)+".png"))) 33 | mask = np.array(Image.open(os.path.join(self.root, "seglabel-224", str(index)+".png"))) 34 | 35 | rgb_image = transforms.ToTensor()(rgb_image) 36 | depth_image = transforms.ToTensor()(depth_image[:, :, np.newaxis]) 37 | 38 | mask = torch.from_numpy(mask) 39 | mask = mask.type(torch.LongTensor) 40 | 41 | return {'rgb_image': rgb_image, 'depth_image': depth_image, 'mask': mask, 'path': str(index)+".png"} 42 | 43 | def __len__(self): 44 | return len(self.indexes) 45 | 46 | def name(self): 47 | return 'sunrgbd' 48 | -------------------------------------------------------------------------------- /datasets/calculate_class_weights.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import h5py 3 | import scipy.io as sio 4 | import cv2 5 | import glob 6 | from PIL import Image 7 | 8 | def calc_scannetv2(data_root,n_class): 9 | masks = [] 10 | size = (320,240) 11 | with open('./datasets/scannet/scannetv2_{}.txt'.format('train')) as f: 12 | scans = f.readlines() 13 | scans = [x.strip() for x in scans] 14 | for scan in scans: 15 | ms = glob.glob("{}/{}/label/*.png".format(data_root, scan)) 16 | masks.extend(ms) 17 | mask_numpy = [] 18 | num_images = np.zeros((n_class)) 19 | for index in range(len(masks)): 20 | mask = np.array(Image.open(masks[index])) 21 | mask = cv2.resize(mask, size, interpolation=cv2.INTER_NEAREST) 22 | num_images[np.unique(mask)] += 1 23 | mask_numpy.append(mask) 24 | 25 | mask_numpy = np.array(mask_numpy) 26 | counts = np.array(np.unique(mask_numpy, return_counts=True)).T 27 | freqs = counts [:,1] / num_images 28 | weights = np.median(freqs) / freqs; 29 | np.savetxt('./datasets/scannet/scannetv2_weigths.txt',weights) 30 | 31 | def calc_weigths(dataset,data_root): 32 | if dataset == "scannetv2": 33 | n_class = 41 34 | calc_scannetv2(data_root,n_class) 35 | else: 36 | print ("Dataset {} is not implemented".format(dataset)) 37 | 38 | def main(): 39 | data_root = '/usr/data/cvpr_shared/common_datasets/scannet/tasks/scannet_frames_25k' 40 | calc_weigths("scannetv2",data_root) 41 | 42 | if __name__ == '__main__': 43 | main() 44 | -------------------------------------------------------------------------------- /datasets/create_training_set.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import h5py 3 | import scipy.io as sio 4 | import cv2 5 | import zipfile 6 | 7 | def create_nyu2(): 8 | 9 | print ("Loading mat file ...") 10 | 11 | f = h5py.File('nyuv2/nyu_depth_v2_labeled.mat') 12 | rgb_images = np.array([x for x in f["images"]]) 13 | depth_images = np.array([x for x in f["depths"]]) 14 | labels= np.array([x for x in f["labels"]]) 15 | 16 | mapping = sio.loadmat("nyuv2/nyuv2_40class_mapping.mat")["mapping"][0] 17 | 18 | print ("Mapping 894 class to 40 class...") 19 | 20 | for map_from, map_to in enumerate(mapping): 21 | # label 0 means unlabeled and will be ignored in training 22 | map_from = map_from + 1 23 | print("--- mapping from " + str(map_from) + " to " + str(map_to), end='\r') 24 | labels[labels==map_from] = map_to 25 | 26 | # for i,label in enumerate(labels): 27 | # for j,row in enumerate(label): 28 | # for k,value in enumerate(row): 29 | # labels[i][j][k] = mapping[value-1] - 1 30 | 31 | print ("Resizing and scaling images and labels...") 32 | 33 | rgbs=[] 34 | depths=[] 35 | masks=[] 36 | 37 | max_depth,min_depth = np.max(depth_images),np.min(depth_images) 38 | 39 | 40 | for rgb_image in rgb_images: 41 | img = np.transpose(rgb_image,(2,1,0)) 42 | img = cv2.resize(img,(320,240),interpolation=cv2.INTER_LINEAR) 43 | rgbs.append(img) 44 | 45 | for depth_image in depth_images: 46 | depth_image = np.transpose(depth_image,(1,0)) 47 | depth_image = cv2.resize(depth_image,(320,240), interpolation=cv2.INTER_NEAREST).astype(np.float) 48 | depth_image = (depth_image - depth_image.min()) / (depth_image.max() - depth_image.min()) * 255 49 | depths.append(depth_image.astype(np.uint8)) 50 | 51 | for label in labels: 52 | label = np.transpose(label,(1,0)) 53 | label = cv2.resize(label,(320,240),interpolation=cv2.INTER_NEAREST) 54 | masks.append(label) 55 | 56 | d = { "rgb_images":rgbs, "depth_images":depths, "masks":masks} 57 | 58 | print ("Saving images and labels...") 59 | 60 | np.save("nyuv2/nyuv2.npy", d) 61 | 62 | print ("Finished !") 63 | def create_dataset(dataset): 64 | 65 | if dataset == "nyu2": 66 | create_nyu2() 67 | else: 68 | print ("Dataset {} is not implemented".format(dataset)) 69 | 70 | def main(): 71 | create_dataset("nyu2") 72 | 73 | if __name__ == '__main__': 74 | main() 75 | -------------------------------------------------------------------------------- /datasets/download_nyu2.sh: -------------------------------------------------------------------------------- 1 | wget http://horatio.cs.nyu.edu/mit/silberman/nyu_depth_v2/nyu_depth_v2_labeled.mat -P nyuv2 2 | wget http://horatio.cs.nyu.edu/mit/silberman/indoor_seg_sup/splits.mat -P nyuv2 3 | wget https://github.com/tum-vision/fusenet/blob/master/fusenet/data/nyuv2_40class_mapping.mat -P nyuv2 4 | -------------------------------------------------------------------------------- /datasets/palette.txt: -------------------------------------------------------------------------------- 1 | 0 0 0 2 | 128 0 0 3 | 0 128 0 4 | 128 128 0 5 | 0 0 128 6 | 128 0 128 7 | 0 128 128 8 | 128 128 128 9 | 64 0 0 10 | 192 0 0 11 | 64 128 0 12 | 192 128 0 13 | 64 0 128 14 | 192 0 128 15 | 64 128 128 16 | 192 128 128 17 | 0 64 0 18 | 128 64 0 19 | 0 192 0 20 | 128 192 0 21 | 0 64 128 22 | 128 64 128 23 | 0 192 128 24 | 128 192 128 25 | 64 64 0 26 | 192 64 0 27 | 64 192 0 28 | 192 192 0 29 | 64 64 128 30 | 192 64 128 31 | 64 192 128 32 | 192 192 128 33 | 0 0 64 34 | 128 0 64 35 | 0 128 64 36 | 128 128 64 37 | 0 0 192 38 | 128 0 192 39 | 0 128 192 40 | 128 128 192 41 | 64 0 64 42 | 192 0 64 43 | 64 128 64 44 | 192 128 64 45 | 64 0 192 46 | 192 0 192 47 | 64 128 192 48 | 192 128 192 49 | 0 64 64 50 | 128 64 64 51 | 0 192 64 52 | 128 192 64 53 | 0 64 192 54 | 128 64 192 55 | 0 192 192 56 | 128 192 192 57 | 64 64 64 58 | 192 64 64 59 | 64 192 64 60 | 192 192 64 61 | 64 64 192 62 | 192 64 192 63 | 64 192 192 64 | 192 192 192 65 | 32 0 0 66 | 160 0 0 67 | 32 128 0 68 | 160 128 0 69 | 32 0 128 70 | 160 0 128 71 | 32 128 128 72 | 160 128 128 73 | 96 0 0 74 | 224 0 0 75 | 96 128 0 76 | 224 128 0 77 | 96 0 128 78 | 224 0 128 79 | 96 128 128 80 | 224 128 128 81 | 32 64 0 82 | 160 64 0 83 | 32 192 0 84 | 160 192 0 85 | 32 64 128 86 | 160 64 128 87 | 32 192 128 88 | 160 192 128 89 | 96 64 0 90 | 224 64 0 91 | 96 192 0 92 | 224 192 0 93 | 96 64 128 94 | 224 64 128 95 | 96 192 128 96 | 224 192 128 97 | 32 0 64 98 | 160 0 64 99 | 32 128 64 100 | 160 128 64 101 | 32 0 192 102 | 160 0 192 103 | 32 128 192 104 | 160 128 192 105 | 96 0 64 106 | 224 0 64 107 | 96 128 64 108 | 224 128 64 109 | 96 0 192 110 | 224 0 192 111 | 96 128 192 112 | 224 128 192 113 | 32 64 64 114 | 160 64 64 115 | 32 192 64 116 | 160 192 64 117 | 32 64 192 118 | 160 64 192 119 | 32 192 192 120 | 160 192 192 121 | 96 64 64 122 | 224 64 64 123 | 96 192 64 124 | 224 192 64 125 | 96 64 192 126 | 224 64 192 127 | 96 192 192 128 | 224 192 192 129 | 0 32 0 130 | 128 32 0 131 | 0 160 0 132 | 128 160 0 133 | 0 32 128 134 | 128 32 128 135 | 0 160 128 136 | 128 160 128 137 | 64 32 0 138 | 192 32 0 139 | 64 160 0 140 | 192 160 0 141 | 64 32 128 142 | 192 32 128 143 | 64 160 128 144 | 192 160 128 145 | 0 96 0 146 | 128 96 0 147 | 0 224 0 148 | 128 224 0 149 | 0 96 128 150 | 128 96 128 151 | 0 224 128 152 | 128 224 128 153 | 64 96 0 154 | 192 96 0 155 | 64 224 0 156 | 192 224 0 157 | 64 96 128 158 | 192 96 128 159 | 64 224 128 160 | 192 224 128 161 | 0 32 64 162 | 128 32 64 163 | 0 160 64 164 | 128 160 64 165 | 0 32 192 166 | 128 32 192 167 | 0 160 192 168 | 128 160 192 169 | 64 32 64 170 | 192 32 64 171 | 64 160 64 172 | 192 160 64 173 | 64 32 192 174 | 192 32 192 175 | 64 160 192 176 | 192 160 192 177 | 0 96 64 178 | 128 96 64 179 | 0 224 64 180 | 128 224 64 181 | 0 96 192 182 | 128 96 192 183 | 0 224 192 184 | 128 224 192 185 | 64 96 64 186 | 192 96 64 187 | 64 224 64 188 | 192 224 64 189 | 64 96 192 190 | 192 96 192 191 | 64 224 192 192 | 192 224 192 193 | 32 32 0 194 | 160 32 0 195 | 32 160 0 196 | 160 160 0 197 | 32 32 128 198 | 160 32 128 199 | 32 160 128 200 | 160 160 128 201 | 96 32 0 202 | 224 32 0 203 | 96 160 0 204 | 224 160 0 205 | 96 32 128 206 | 224 32 128 207 | 96 160 128 208 | 224 160 128 209 | 32 96 0 210 | 160 96 0 211 | 32 224 0 212 | 160 224 0 213 | 32 96 128 214 | 160 96 128 215 | 32 224 128 216 | 160 224 128 217 | 96 96 0 218 | 224 96 0 219 | 96 224 0 220 | 224 224 0 221 | 96 96 128 222 | 224 96 128 223 | 96 224 128 224 | 224 224 128 225 | 32 32 64 226 | 160 32 64 227 | 32 160 64 228 | 160 160 64 229 | 32 32 192 230 | 160 32 192 231 | 32 160 192 232 | 160 160 192 233 | 96 32 64 234 | 224 32 64 235 | 96 160 64 236 | 224 160 64 237 | 96 32 192 238 | 224 32 192 239 | 96 160 192 240 | 224 160 192 241 | 32 96 64 242 | 160 96 64 243 | 32 224 64 244 | 160 224 64 245 | 32 96 192 246 | 160 96 192 247 | 32 224 192 248 | 160 224 192 249 | 96 96 64 250 | 224 96 64 251 | 96 224 64 252 | 224 224 64 253 | 96 96 192 254 | 224 96 192 255 | 96 224 192 256 | 224 224 192 257 | -------------------------------------------------------------------------------- /datasets/scannet/scannetv2_test.txt: -------------------------------------------------------------------------------- 1 | scene0707_00 2 | scene0708_00 3 | scene0709_00 4 | scene0710_00 5 | scene0711_00 6 | scene0712_00 7 | scene0713_00 8 | scene0714_00 9 | scene0715_00 10 | scene0716_00 11 | scene0717_00 12 | scene0718_00 13 | scene0719_00 14 | scene0720_00 15 | scene0721_00 16 | scene0722_00 17 | scene0723_00 18 | scene0724_00 19 | scene0725_00 20 | scene0726_00 21 | scene0727_00 22 | scene0728_00 23 | scene0729_00 24 | scene0730_00 25 | scene0731_00 26 | scene0732_00 27 | scene0733_00 28 | scene0734_00 29 | scene0735_00 30 | scene0736_00 31 | scene0737_00 32 | scene0738_00 33 | scene0739_00 34 | scene0740_00 35 | scene0741_00 36 | scene0742_00 37 | scene0743_00 38 | scene0744_00 39 | scene0745_00 40 | scene0746_00 41 | scene0747_00 42 | scene0748_00 43 | scene0749_00 44 | scene0750_00 45 | scene0751_00 46 | scene0752_00 47 | scene0753_00 48 | scene0754_00 49 | scene0755_00 50 | scene0756_00 51 | scene0757_00 52 | scene0758_00 53 | scene0759_00 54 | scene0760_00 55 | scene0761_00 56 | scene0762_00 57 | scene0763_00 58 | scene0764_00 59 | scene0765_00 60 | scene0766_00 61 | scene0767_00 62 | scene0768_00 63 | scene0769_00 64 | scene0770_00 65 | scene0771_00 66 | scene0772_00 67 | scene0773_00 68 | scene0774_00 69 | scene0775_00 70 | scene0776_00 71 | scene0777_00 72 | scene0778_00 73 | scene0779_00 74 | scene0780_00 75 | scene0781_00 76 | scene0782_00 77 | scene0783_00 78 | scene0784_00 79 | scene0785_00 80 | scene0786_00 81 | scene0787_00 82 | scene0788_00 83 | scene0789_00 84 | scene0790_00 85 | scene0791_00 86 | scene0792_00 87 | scene0793_00 88 | scene0794_00 89 | scene0795_00 90 | scene0796_00 91 | scene0797_00 92 | scene0798_00 93 | scene0799_00 94 | scene0800_00 95 | scene0801_00 96 | scene0802_00 97 | scene0803_00 98 | scene0804_00 99 | scene0805_00 100 | scene0806_00 101 | -------------------------------------------------------------------------------- /datasets/scannet/scannetv2_train.txt: -------------------------------------------------------------------------------- 1 | scene0191_00 2 | scene0191_01 3 | scene0191_02 4 | scene0119_00 5 | scene0230_00 6 | scene0528_00 7 | scene0528_01 8 | scene0705_00 9 | scene0705_01 10 | scene0705_02 11 | scene0415_00 12 | scene0415_01 13 | scene0415_02 14 | scene0007_00 15 | scene0141_00 16 | scene0141_01 17 | scene0141_02 18 | scene0515_00 19 | scene0515_01 20 | scene0515_02 21 | scene0447_00 22 | scene0447_01 23 | scene0447_02 24 | scene0531_00 25 | scene0503_00 26 | scene0285_00 27 | scene0069_00 28 | scene0584_00 29 | scene0584_01 30 | scene0584_02 31 | scene0581_00 32 | scene0581_01 33 | scene0581_02 34 | scene0620_00 35 | scene0620_01 36 | scene0263_00 37 | scene0263_01 38 | scene0481_00 39 | scene0481_01 40 | scene0020_00 41 | scene0020_01 42 | scene0291_00 43 | scene0291_01 44 | scene0291_02 45 | scene0469_00 46 | scene0469_01 47 | scene0469_02 48 | scene0659_00 49 | scene0659_01 50 | scene0024_00 51 | scene0024_01 52 | scene0024_02 53 | scene0564_00 54 | scene0117_00 55 | scene0027_00 56 | scene0027_01 57 | scene0027_02 58 | scene0028_00 59 | scene0330_00 60 | scene0418_00 61 | scene0418_01 62 | scene0418_02 63 | scene0233_00 64 | scene0233_01 65 | scene0673_00 66 | scene0673_01 67 | scene0673_02 68 | scene0673_03 69 | scene0673_04 70 | scene0673_05 71 | scene0585_00 72 | scene0585_01 73 | scene0362_00 74 | scene0362_01 75 | scene0362_02 76 | scene0362_03 77 | scene0035_00 78 | scene0035_01 79 | scene0358_00 80 | scene0358_01 81 | scene0358_02 82 | scene0037_00 83 | scene0194_00 84 | scene0321_00 85 | scene0293_00 86 | scene0293_01 87 | scene0623_00 88 | scene0623_01 89 | scene0592_00 90 | scene0592_01 91 | scene0569_00 92 | scene0569_01 93 | scene0413_00 94 | scene0313_00 95 | scene0313_01 96 | scene0313_02 97 | scene0480_00 98 | scene0480_01 99 | scene0401_00 100 | scene0517_00 101 | scene0517_01 102 | scene0517_02 103 | scene0032_00 104 | scene0032_01 105 | scene0613_00 106 | scene0613_01 107 | scene0613_02 108 | scene0306_00 109 | scene0306_01 110 | scene0052_00 111 | scene0052_01 112 | scene0052_02 113 | scene0053_00 114 | scene0444_00 115 | scene0444_01 116 | scene0055_00 117 | scene0055_01 118 | scene0055_02 119 | scene0560_00 120 | scene0589_00 121 | scene0589_01 122 | scene0589_02 123 | scene0610_00 124 | scene0610_01 125 | scene0610_02 126 | scene0364_00 127 | scene0364_01 128 | scene0383_00 129 | scene0383_01 130 | scene0383_02 131 | scene0006_00 132 | scene0006_01 133 | scene0006_02 134 | scene0275_00 135 | scene0451_00 136 | scene0451_01 137 | scene0451_02 138 | scene0451_03 139 | scene0451_04 140 | scene0451_05 141 | scene0135_00 142 | scene0065_00 143 | scene0065_01 144 | scene0065_02 145 | scene0104_00 146 | scene0674_00 147 | scene0674_01 148 | scene0448_00 149 | scene0448_01 150 | scene0448_02 151 | scene0502_00 152 | scene0502_01 153 | scene0502_02 154 | scene0440_00 155 | scene0440_01 156 | scene0440_02 157 | scene0071_00 158 | scene0072_00 159 | scene0072_01 160 | scene0072_02 161 | scene0509_00 162 | scene0509_01 163 | scene0509_02 164 | scene0649_00 165 | scene0649_01 166 | scene0602_00 167 | scene0694_00 168 | scene0694_01 169 | scene0101_00 170 | scene0101_01 171 | scene0101_02 172 | scene0101_03 173 | scene0101_04 174 | scene0101_05 175 | scene0218_00 176 | scene0218_01 177 | scene0579_00 178 | scene0579_01 179 | scene0579_02 180 | scene0039_00 181 | scene0039_01 182 | scene0493_00 183 | scene0493_01 184 | scene0242_00 185 | scene0242_01 186 | scene0242_02 187 | scene0083_00 188 | scene0083_01 189 | scene0127_00 190 | scene0127_01 191 | scene0662_00 192 | scene0662_01 193 | scene0662_02 194 | scene0018_00 195 | scene0087_00 196 | scene0087_01 197 | scene0087_02 198 | scene0332_00 199 | scene0332_01 200 | scene0332_02 201 | scene0628_00 202 | scene0628_01 203 | scene0628_02 204 | scene0134_00 205 | scene0134_01 206 | scene0134_02 207 | scene0238_00 208 | scene0238_01 209 | scene0092_00 210 | scene0092_01 211 | scene0092_02 212 | scene0092_03 213 | scene0092_04 214 | scene0022_00 215 | scene0022_01 216 | scene0467_00 217 | scene0392_00 218 | scene0392_01 219 | scene0392_02 220 | scene0424_00 221 | scene0424_01 222 | scene0424_02 223 | scene0646_00 224 | scene0646_01 225 | scene0646_02 226 | scene0098_00 227 | scene0098_01 228 | scene0044_00 229 | scene0044_01 230 | scene0044_02 231 | scene0510_00 232 | scene0510_01 233 | scene0510_02 234 | scene0571_00 235 | scene0571_01 236 | scene0166_00 237 | scene0166_01 238 | scene0166_02 239 | scene0563_00 240 | scene0172_00 241 | scene0172_01 242 | scene0388_00 243 | scene0388_01 244 | scene0215_00 245 | scene0215_01 246 | scene0252_00 247 | scene0287_00 248 | scene0668_00 249 | scene0572_00 250 | scene0572_01 251 | scene0572_02 252 | scene0026_00 253 | scene0224_00 254 | scene0113_00 255 | scene0113_01 256 | scene0551_00 257 | scene0381_00 258 | scene0381_01 259 | scene0381_02 260 | scene0371_00 261 | scene0371_01 262 | scene0460_00 263 | scene0118_00 264 | scene0118_01 265 | scene0118_02 266 | scene0417_00 267 | scene0008_00 268 | scene0634_00 269 | scene0521_00 270 | scene0123_00 271 | scene0123_01 272 | scene0123_02 273 | scene0045_00 274 | scene0045_01 275 | scene0511_00 276 | scene0511_01 277 | scene0114_00 278 | scene0114_01 279 | scene0114_02 280 | scene0070_00 281 | scene0029_00 282 | scene0029_01 283 | scene0029_02 284 | scene0129_00 285 | scene0103_00 286 | scene0103_01 287 | scene0002_00 288 | scene0002_01 289 | scene0132_00 290 | scene0132_01 291 | scene0132_02 292 | scene0124_00 293 | scene0124_01 294 | scene0143_00 295 | scene0143_01 296 | scene0143_02 297 | scene0604_00 298 | scene0604_01 299 | scene0604_02 300 | scene0507_00 301 | scene0105_00 302 | scene0105_01 303 | scene0105_02 304 | scene0428_00 305 | scene0428_01 306 | scene0311_00 307 | scene0140_00 308 | scene0140_01 309 | scene0182_00 310 | scene0182_01 311 | scene0182_02 312 | scene0142_00 313 | scene0142_01 314 | scene0399_00 315 | scene0399_01 316 | scene0012_00 317 | scene0012_01 318 | scene0012_02 319 | scene0060_00 320 | scene0060_01 321 | scene0370_00 322 | scene0370_01 323 | scene0370_02 324 | scene0310_00 325 | scene0310_01 326 | scene0310_02 327 | scene0661_00 328 | scene0650_00 329 | scene0152_00 330 | scene0152_01 331 | scene0152_02 332 | scene0158_00 333 | scene0158_01 334 | scene0158_02 335 | scene0482_00 336 | scene0482_01 337 | scene0600_00 338 | scene0600_01 339 | scene0600_02 340 | scene0393_00 341 | scene0393_01 342 | scene0393_02 343 | scene0562_00 344 | scene0174_00 345 | scene0174_01 346 | scene0157_00 347 | scene0157_01 348 | scene0161_00 349 | scene0161_01 350 | scene0161_02 351 | scene0159_00 352 | scene0254_00 353 | scene0254_01 354 | scene0115_00 355 | scene0115_01 356 | scene0115_02 357 | scene0162_00 358 | scene0163_00 359 | scene0163_01 360 | scene0523_00 361 | scene0523_01 362 | scene0523_02 363 | scene0459_00 364 | scene0459_01 365 | scene0175_00 366 | scene0085_00 367 | scene0085_01 368 | scene0279_00 369 | scene0279_01 370 | scene0279_02 371 | scene0201_00 372 | scene0201_01 373 | scene0201_02 374 | scene0283_00 375 | scene0456_00 376 | scene0456_01 377 | scene0429_00 378 | scene0043_00 379 | scene0043_01 380 | scene0419_00 381 | scene0419_01 382 | scene0419_02 383 | scene0368_00 384 | scene0368_01 385 | scene0348_00 386 | scene0348_01 387 | scene0348_02 388 | scene0442_00 389 | scene0178_00 390 | scene0380_00 391 | scene0380_01 392 | scene0380_02 393 | scene0165_00 394 | scene0165_01 395 | scene0165_02 396 | scene0181_00 397 | scene0181_01 398 | scene0181_02 399 | scene0181_03 400 | scene0333_00 401 | scene0614_00 402 | scene0614_01 403 | scene0614_02 404 | scene0404_00 405 | scene0404_01 406 | scene0404_02 407 | scene0185_00 408 | scene0126_00 409 | scene0126_01 410 | scene0126_02 411 | scene0519_00 412 | scene0236_00 413 | scene0236_01 414 | scene0189_00 415 | scene0075_00 416 | scene0267_00 417 | scene0192_00 418 | scene0192_01 419 | scene0192_02 420 | scene0281_00 421 | scene0420_00 422 | scene0420_01 423 | scene0420_02 424 | scene0195_00 425 | scene0195_01 426 | scene0195_02 427 | scene0597_00 428 | scene0597_01 429 | scene0597_02 430 | scene0041_00 431 | scene0041_01 432 | scene0111_00 433 | scene0111_01 434 | scene0111_02 435 | scene0666_00 436 | scene0666_01 437 | scene0666_02 438 | scene0200_00 439 | scene0200_01 440 | scene0200_02 441 | scene0536_00 442 | scene0536_01 443 | scene0536_02 444 | scene0390_00 445 | scene0280_00 446 | scene0280_01 447 | scene0280_02 448 | scene0344_00 449 | scene0344_01 450 | scene0205_00 451 | scene0205_01 452 | scene0205_02 453 | scene0484_00 454 | scene0484_01 455 | scene0009_00 456 | scene0009_01 457 | scene0009_02 458 | scene0302_00 459 | scene0302_01 460 | scene0209_00 461 | scene0209_01 462 | scene0209_02 463 | scene0210_00 464 | scene0210_01 465 | scene0395_00 466 | scene0395_01 467 | scene0395_02 468 | scene0683_00 469 | scene0601_00 470 | scene0601_01 471 | scene0214_00 472 | scene0214_01 473 | scene0214_02 474 | scene0477_00 475 | scene0477_01 476 | scene0439_00 477 | scene0439_01 478 | scene0468_00 479 | scene0468_01 480 | scene0468_02 481 | scene0546_00 482 | scene0466_00 483 | scene0466_01 484 | scene0220_00 485 | scene0220_01 486 | scene0220_02 487 | scene0122_00 488 | scene0122_01 489 | scene0130_00 490 | scene0110_00 491 | scene0110_01 492 | scene0110_02 493 | scene0327_00 494 | scene0156_00 495 | scene0266_00 496 | scene0266_01 497 | scene0001_00 498 | scene0001_01 499 | scene0228_00 500 | scene0199_00 501 | scene0219_00 502 | scene0464_00 503 | scene0232_00 504 | scene0232_01 505 | scene0232_02 506 | scene0299_00 507 | scene0299_01 508 | scene0530_00 509 | scene0363_00 510 | scene0453_00 511 | scene0453_01 512 | scene0570_00 513 | scene0570_01 514 | scene0570_02 515 | scene0183_00 516 | scene0239_00 517 | scene0239_01 518 | scene0239_02 519 | scene0373_00 520 | scene0373_01 521 | scene0241_00 522 | scene0241_01 523 | scene0241_02 524 | scene0188_00 525 | scene0622_00 526 | scene0622_01 527 | scene0244_00 528 | scene0244_01 529 | scene0691_00 530 | scene0691_01 531 | scene0206_00 532 | scene0206_01 533 | scene0206_02 534 | scene0247_00 535 | scene0247_01 536 | scene0061_00 537 | scene0061_01 538 | scene0082_00 539 | scene0250_00 540 | scene0250_01 541 | scene0250_02 542 | scene0501_00 543 | scene0501_01 544 | scene0501_02 545 | scene0320_00 546 | scene0320_01 547 | scene0320_02 548 | scene0320_03 549 | scene0631_00 550 | scene0631_01 551 | scene0631_02 552 | scene0255_00 553 | scene0255_01 554 | scene0255_02 555 | scene0047_00 556 | scene0265_00 557 | scene0265_01 558 | scene0265_02 559 | scene0004_00 560 | scene0336_00 561 | scene0336_01 562 | scene0058_00 563 | scene0058_01 564 | scene0260_00 565 | scene0260_01 566 | scene0260_02 567 | scene0243_00 568 | scene0603_00 569 | scene0603_01 570 | scene0093_00 571 | scene0093_01 572 | scene0093_02 573 | scene0109_00 574 | scene0109_01 575 | scene0434_00 576 | scene0434_01 577 | scene0434_02 578 | scene0290_00 579 | scene0627_00 580 | scene0627_01 581 | scene0470_00 582 | scene0470_01 583 | scene0137_00 584 | scene0137_01 585 | scene0137_02 586 | scene0270_00 587 | scene0270_01 588 | scene0270_02 589 | scene0271_00 590 | scene0271_01 591 | scene0504_00 592 | scene0274_00 593 | scene0274_01 594 | scene0274_02 595 | scene0036_00 596 | scene0036_01 597 | scene0276_00 598 | scene0276_01 599 | scene0272_00 600 | scene0272_01 601 | scene0499_00 602 | scene0698_00 603 | scene0698_01 604 | scene0051_00 605 | scene0051_01 606 | scene0051_02 607 | scene0051_03 608 | scene0108_00 609 | scene0245_00 610 | scene0369_00 611 | scene0369_01 612 | scene0369_02 613 | scene0284_00 614 | scene0289_00 615 | scene0289_01 616 | scene0286_00 617 | scene0286_01 618 | scene0286_02 619 | scene0286_03 620 | scene0031_00 621 | scene0031_01 622 | scene0031_02 623 | scene0545_00 624 | scene0545_01 625 | scene0545_02 626 | scene0557_00 627 | scene0557_01 628 | scene0557_02 629 | scene0533_00 630 | scene0533_01 631 | scene0116_00 632 | scene0116_01 633 | scene0116_02 634 | scene0611_00 635 | scene0611_01 636 | scene0688_00 637 | scene0294_00 638 | scene0294_01 639 | scene0294_02 640 | scene0295_00 641 | scene0295_01 642 | scene0296_00 643 | scene0296_01 644 | scene0596_00 645 | scene0596_01 646 | scene0596_02 647 | scene0532_00 648 | scene0532_01 649 | scene0637_00 650 | scene0638_00 651 | scene0121_00 652 | scene0121_01 653 | scene0121_02 654 | scene0040_00 655 | scene0040_01 656 | scene0197_00 657 | scene0197_01 658 | scene0197_02 659 | scene0410_00 660 | scene0410_01 661 | scene0305_00 662 | scene0305_01 663 | scene0615_00 664 | scene0615_01 665 | scene0703_00 666 | scene0703_01 667 | scene0555_00 668 | scene0297_00 669 | scene0297_01 670 | scene0297_02 671 | scene0582_00 672 | scene0582_01 673 | scene0582_02 674 | scene0023_00 675 | scene0094_00 676 | scene0013_00 677 | scene0013_01 678 | scene0013_02 679 | scene0136_00 680 | scene0136_01 681 | scene0136_02 682 | scene0407_00 683 | scene0407_01 684 | scene0062_00 685 | scene0062_01 686 | scene0062_02 687 | scene0386_00 688 | scene0318_00 689 | scene0554_00 690 | scene0554_01 691 | scene0497_00 692 | scene0213_00 693 | scene0258_00 694 | scene0323_00 695 | scene0323_01 696 | scene0324_00 697 | scene0324_01 698 | scene0016_00 699 | scene0016_01 700 | scene0016_02 701 | scene0681_00 702 | scene0398_00 703 | scene0398_01 704 | scene0227_00 705 | scene0090_00 706 | scene0066_00 707 | scene0262_00 708 | scene0262_01 709 | scene0155_00 710 | scene0155_01 711 | scene0155_02 712 | scene0352_00 713 | scene0352_01 714 | scene0352_02 715 | scene0038_00 716 | scene0038_01 717 | scene0038_02 718 | scene0335_00 719 | scene0335_01 720 | scene0335_02 721 | scene0261_00 722 | scene0261_01 723 | scene0261_02 724 | scene0261_03 725 | scene0640_00 726 | scene0640_01 727 | scene0640_02 728 | scene0080_00 729 | scene0080_01 730 | scene0080_02 731 | scene0403_00 732 | scene0403_01 733 | scene0282_00 734 | scene0282_01 735 | scene0282_02 736 | scene0682_00 737 | scene0173_00 738 | scene0173_01 739 | scene0173_02 740 | scene0522_00 741 | scene0687_00 742 | scene0345_00 743 | scene0345_01 744 | scene0612_00 745 | scene0612_01 746 | scene0411_00 747 | scene0411_01 748 | scene0411_02 749 | scene0625_00 750 | scene0625_01 751 | scene0211_00 752 | scene0211_01 753 | scene0211_02 754 | scene0211_03 755 | scene0676_00 756 | scene0676_01 757 | scene0179_00 758 | scene0498_00 759 | scene0498_01 760 | scene0498_02 761 | scene0547_00 762 | scene0547_01 763 | scene0547_02 764 | scene0269_00 765 | scene0269_01 766 | scene0269_02 767 | scene0366_00 768 | scene0680_00 769 | scene0680_01 770 | scene0588_00 771 | scene0588_01 772 | scene0588_02 773 | scene0588_03 774 | scene0346_00 775 | scene0346_01 776 | scene0359_00 777 | scene0359_01 778 | scene0014_00 779 | scene0120_00 780 | scene0120_01 781 | scene0212_00 782 | scene0212_01 783 | scene0212_02 784 | scene0176_00 785 | scene0049_00 786 | scene0259_00 787 | scene0259_01 788 | scene0586_00 789 | scene0586_01 790 | scene0586_02 791 | scene0309_00 792 | scene0309_01 793 | scene0125_00 794 | scene0455_00 795 | scene0177_00 796 | scene0177_01 797 | scene0177_02 798 | scene0326_00 799 | scene0372_00 800 | scene0171_00 801 | scene0171_01 802 | scene0374_00 803 | scene0654_00 804 | scene0654_01 805 | scene0445_00 806 | scene0445_01 807 | scene0475_00 808 | scene0475_01 809 | scene0475_02 810 | scene0349_00 811 | scene0349_01 812 | scene0234_00 813 | scene0669_00 814 | scene0669_01 815 | scene0375_00 816 | scene0375_01 817 | scene0375_02 818 | scene0387_00 819 | scene0387_01 820 | scene0387_02 821 | scene0312_00 822 | scene0312_01 823 | scene0312_02 824 | scene0384_00 825 | scene0385_00 826 | scene0385_01 827 | scene0385_02 828 | scene0000_00 829 | scene0000_01 830 | scene0000_02 831 | scene0376_00 832 | scene0376_01 833 | scene0376_02 834 | scene0301_00 835 | scene0301_01 836 | scene0301_02 837 | scene0322_00 838 | scene0542_00 839 | scene0079_00 840 | scene0079_01 841 | scene0099_00 842 | scene0099_01 843 | scene0476_00 844 | scene0476_01 845 | scene0476_02 846 | scene0394_00 847 | scene0394_01 848 | scene0147_00 849 | scene0147_01 850 | scene0067_00 851 | scene0067_01 852 | scene0067_02 853 | scene0397_00 854 | scene0397_01 855 | scene0337_00 856 | scene0337_01 857 | scene0337_02 858 | scene0431_00 859 | scene0223_00 860 | scene0223_01 861 | scene0223_02 862 | scene0010_00 863 | scene0010_01 864 | scene0402_00 865 | scene0268_00 866 | scene0268_01 867 | scene0268_02 868 | scene0679_00 869 | scene0679_01 870 | scene0405_00 871 | scene0128_00 872 | scene0408_00 873 | scene0408_01 874 | scene0190_00 875 | scene0107_00 876 | scene0076_00 877 | scene0167_00 878 | scene0361_00 879 | scene0361_01 880 | scene0361_02 881 | scene0216_00 882 | scene0202_00 883 | scene0303_00 884 | scene0303_01 885 | scene0303_02 886 | scene0446_00 887 | scene0446_01 888 | scene0089_00 889 | scene0089_01 890 | scene0089_02 891 | scene0360_00 892 | scene0150_00 893 | scene0150_01 894 | scene0150_02 895 | scene0421_00 896 | scene0421_01 897 | scene0421_02 898 | scene0454_00 899 | scene0626_00 900 | scene0626_01 901 | scene0626_02 902 | scene0186_00 903 | scene0186_01 904 | scene0538_00 905 | scene0479_00 906 | scene0479_01 907 | scene0479_02 908 | scene0656_00 909 | scene0656_01 910 | scene0656_02 911 | scene0656_03 912 | scene0525_00 913 | scene0525_01 914 | scene0525_02 915 | scene0308_00 916 | scene0396_00 917 | scene0396_01 918 | scene0396_02 919 | scene0624_00 920 | scene0292_00 921 | scene0292_01 922 | scene0632_00 923 | scene0253_00 924 | scene0021_00 925 | scene0325_00 926 | scene0325_01 927 | scene0437_00 928 | scene0437_01 929 | scene0438_00 930 | scene0590_00 931 | scene0590_01 932 | scene0400_00 933 | scene0400_01 934 | scene0541_00 935 | scene0541_01 936 | scene0541_02 937 | scene0677_00 938 | scene0677_01 939 | scene0677_02 940 | scene0443_00 941 | scene0315_00 942 | scene0288_00 943 | scene0288_01 944 | scene0288_02 945 | scene0422_00 946 | scene0672_00 947 | scene0672_01 948 | scene0184_00 949 | scene0449_00 950 | scene0449_01 951 | scene0449_02 952 | scene0048_00 953 | scene0048_01 954 | scene0138_00 955 | scene0452_00 956 | scene0452_01 957 | scene0452_02 958 | scene0667_00 959 | scene0667_01 960 | scene0667_02 961 | scene0463_00 962 | scene0463_01 963 | scene0078_00 964 | scene0078_01 965 | scene0078_02 966 | scene0636_00 967 | scene0457_00 968 | scene0457_01 969 | scene0457_02 970 | scene0465_00 971 | scene0465_01 972 | scene0577_00 973 | scene0151_00 974 | scene0151_01 975 | scene0339_00 976 | scene0573_00 977 | scene0573_01 978 | scene0154_00 979 | scene0096_00 980 | scene0096_01 981 | scene0096_02 982 | scene0235_00 983 | scene0168_00 984 | scene0168_01 985 | scene0168_02 986 | scene0594_00 987 | scene0587_00 988 | scene0587_01 989 | scene0587_02 990 | scene0587_03 991 | scene0229_00 992 | scene0229_01 993 | scene0229_02 994 | scene0512_00 995 | scene0106_00 996 | scene0106_01 997 | scene0106_02 998 | scene0472_00 999 | scene0472_01 1000 | scene0472_02 1001 | scene0489_00 1002 | scene0489_01 1003 | scene0489_02 1004 | scene0425_00 1005 | scene0425_01 1006 | scene0641_00 1007 | scene0526_00 1008 | scene0526_01 1009 | scene0317_00 1010 | scene0317_01 1011 | scene0544_00 1012 | scene0017_00 1013 | scene0017_01 1014 | scene0017_02 1015 | scene0042_00 1016 | scene0042_01 1017 | scene0042_02 1018 | scene0576_00 1019 | scene0576_01 1020 | scene0576_02 1021 | scene0347_00 1022 | scene0347_01 1023 | scene0347_02 1024 | scene0436_00 1025 | scene0226_00 1026 | scene0226_01 1027 | scene0485_00 1028 | scene0486_00 1029 | scene0487_00 1030 | scene0487_01 1031 | scene0619_00 1032 | scene0097_00 1033 | scene0367_00 1034 | scene0367_01 1035 | scene0491_00 1036 | scene0492_00 1037 | scene0492_01 1038 | scene0005_00 1039 | scene0005_01 1040 | scene0543_00 1041 | scene0543_01 1042 | scene0543_02 1043 | scene0657_00 1044 | scene0341_00 1045 | scene0341_01 1046 | scene0534_00 1047 | scene0534_01 1048 | scene0319_00 1049 | scene0273_00 1050 | scene0273_01 1051 | scene0225_00 1052 | scene0198_00 1053 | scene0003_00 1054 | scene0003_01 1055 | scene0003_02 1056 | scene0409_00 1057 | scene0409_01 1058 | scene0331_00 1059 | scene0331_01 1060 | scene0505_00 1061 | scene0505_01 1062 | scene0505_02 1063 | scene0505_03 1064 | scene0505_04 1065 | scene0506_00 1066 | scene0057_00 1067 | scene0057_01 1068 | scene0074_00 1069 | scene0074_01 1070 | scene0074_02 1071 | scene0091_00 1072 | scene0112_00 1073 | scene0112_01 1074 | scene0112_02 1075 | scene0240_00 1076 | scene0102_00 1077 | scene0102_01 1078 | scene0513_00 1079 | scene0514_00 1080 | scene0514_01 1081 | scene0537_00 1082 | scene0516_00 1083 | scene0516_01 1084 | scene0495_00 1085 | scene0617_00 1086 | scene0133_00 1087 | scene0520_00 1088 | scene0520_01 1089 | scene0635_00 1090 | scene0635_01 1091 | scene0054_00 1092 | scene0473_00 1093 | scene0473_01 1094 | scene0524_00 1095 | scene0524_01 1096 | scene0379_00 1097 | scene0471_00 1098 | scene0471_01 1099 | scene0471_02 1100 | scene0566_00 1101 | scene0248_00 1102 | scene0248_01 1103 | scene0248_02 1104 | scene0529_00 1105 | scene0529_01 1106 | scene0529_02 1107 | scene0391_00 1108 | scene0264_00 1109 | scene0264_01 1110 | scene0264_02 1111 | scene0675_00 1112 | scene0675_01 1113 | scene0350_00 1114 | scene0350_01 1115 | scene0350_02 1116 | scene0450_00 1117 | scene0068_00 1118 | scene0068_01 1119 | scene0237_00 1120 | scene0237_01 1121 | scene0365_00 1122 | scene0365_01 1123 | scene0365_02 1124 | scene0605_00 1125 | scene0605_01 1126 | scene0539_00 1127 | scene0539_01 1128 | scene0539_02 1129 | scene0540_00 1130 | scene0540_01 1131 | scene0540_02 1132 | scene0170_00 1133 | scene0170_01 1134 | scene0170_02 1135 | scene0433_00 1136 | scene0340_00 1137 | scene0340_01 1138 | scene0340_02 1139 | scene0160_00 1140 | scene0160_01 1141 | scene0160_02 1142 | scene0160_03 1143 | scene0160_04 1144 | scene0059_00 1145 | scene0059_01 1146 | scene0059_02 1147 | scene0056_00 1148 | scene0056_01 1149 | scene0478_00 1150 | scene0478_01 1151 | scene0548_00 1152 | scene0548_01 1153 | scene0548_02 1154 | scene0204_00 1155 | scene0204_01 1156 | scene0204_02 1157 | scene0033_00 1158 | scene0145_00 1159 | scene0483_00 1160 | scene0508_00 1161 | scene0508_01 1162 | scene0508_02 1163 | scene0180_00 1164 | scene0148_00 1165 | scene0556_00 1166 | scene0556_01 1167 | scene0416_00 1168 | scene0416_01 1169 | scene0416_02 1170 | scene0416_03 1171 | scene0416_04 1172 | scene0073_00 1173 | scene0073_01 1174 | scene0073_02 1175 | scene0073_03 1176 | scene0034_00 1177 | scene0034_01 1178 | scene0034_02 1179 | scene0639_00 1180 | scene0561_00 1181 | scene0561_01 1182 | scene0298_00 1183 | scene0692_00 1184 | scene0692_01 1185 | scene0692_02 1186 | scene0692_03 1187 | scene0692_04 1188 | scene0642_00 1189 | scene0642_01 1190 | scene0642_02 1191 | scene0642_03 1192 | scene0630_00 1193 | scene0630_01 1194 | scene0630_02 1195 | scene0630_03 1196 | scene0630_04 1197 | scene0630_05 1198 | scene0630_06 1199 | scene0706_00 1200 | scene0567_00 1201 | scene0567_01 1202 | -------------------------------------------------------------------------------- /datasets/scannet/scannetv2_val.txt: -------------------------------------------------------------------------------- 1 | scene0568_00 2 | scene0568_01 3 | scene0568_02 4 | scene0304_00 5 | scene0488_00 6 | scene0488_01 7 | scene0412_00 8 | scene0412_01 9 | scene0217_00 10 | scene0019_00 11 | scene0019_01 12 | scene0414_00 13 | scene0575_00 14 | scene0575_01 15 | scene0575_02 16 | scene0426_00 17 | scene0426_01 18 | scene0426_02 19 | scene0426_03 20 | scene0549_00 21 | scene0549_01 22 | scene0578_00 23 | scene0578_01 24 | scene0578_02 25 | scene0665_00 26 | scene0665_01 27 | scene0050_00 28 | scene0050_01 29 | scene0050_02 30 | scene0257_00 31 | scene0025_00 32 | scene0025_01 33 | scene0025_02 34 | scene0583_00 35 | scene0583_01 36 | scene0583_02 37 | scene0701_00 38 | scene0701_01 39 | scene0701_02 40 | scene0580_00 41 | scene0580_01 42 | scene0565_00 43 | scene0169_00 44 | scene0169_01 45 | scene0655_00 46 | scene0655_01 47 | scene0655_02 48 | scene0063_00 49 | scene0221_00 50 | scene0221_01 51 | scene0591_00 52 | scene0591_01 53 | scene0591_02 54 | scene0678_00 55 | scene0678_01 56 | scene0678_02 57 | scene0462_00 58 | scene0427_00 59 | scene0595_00 60 | scene0193_00 61 | scene0193_01 62 | scene0164_00 63 | scene0164_01 64 | scene0164_02 65 | scene0164_03 66 | scene0598_00 67 | scene0598_01 68 | scene0598_02 69 | scene0599_00 70 | scene0599_01 71 | scene0599_02 72 | scene0328_00 73 | scene0300_00 74 | scene0300_01 75 | scene0354_00 76 | scene0458_00 77 | scene0458_01 78 | scene0423_00 79 | scene0423_01 80 | scene0423_02 81 | scene0307_00 82 | scene0307_01 83 | scene0307_02 84 | scene0606_00 85 | scene0606_01 86 | scene0606_02 87 | scene0432_00 88 | scene0432_01 89 | scene0608_00 90 | scene0608_01 91 | scene0608_02 92 | scene0651_00 93 | scene0651_01 94 | scene0651_02 95 | scene0430_00 96 | scene0430_01 97 | scene0689_00 98 | scene0357_00 99 | scene0357_01 100 | scene0574_00 101 | scene0574_01 102 | scene0574_02 103 | scene0329_00 104 | scene0329_01 105 | scene0329_02 106 | scene0153_00 107 | scene0153_01 108 | scene0616_00 109 | scene0616_01 110 | scene0671_00 111 | scene0671_01 112 | scene0618_00 113 | scene0382_00 114 | scene0382_01 115 | scene0490_00 116 | scene0621_00 117 | scene0607_00 118 | scene0607_01 119 | scene0149_00 120 | scene0695_00 121 | scene0695_01 122 | scene0695_02 123 | scene0695_03 124 | scene0389_00 125 | scene0377_00 126 | scene0377_01 127 | scene0377_02 128 | scene0342_00 129 | scene0139_00 130 | scene0629_00 131 | scene0629_01 132 | scene0629_02 133 | scene0496_00 134 | scene0633_00 135 | scene0633_01 136 | scene0518_00 137 | scene0652_00 138 | scene0406_00 139 | scene0406_01 140 | scene0406_02 141 | scene0144_00 142 | scene0144_01 143 | scene0494_00 144 | scene0278_00 145 | scene0278_01 146 | scene0316_00 147 | scene0609_00 148 | scene0609_01 149 | scene0609_02 150 | scene0609_03 151 | scene0084_00 152 | scene0084_01 153 | scene0084_02 154 | scene0696_00 155 | scene0696_01 156 | scene0696_02 157 | scene0351_00 158 | scene0351_01 159 | scene0643_00 160 | scene0644_00 161 | scene0645_00 162 | scene0645_01 163 | scene0645_02 164 | scene0081_00 165 | scene0081_01 166 | scene0081_02 167 | scene0647_00 168 | scene0647_01 169 | scene0535_00 170 | scene0353_00 171 | scene0353_01 172 | scene0353_02 173 | scene0559_00 174 | scene0559_01 175 | scene0559_02 176 | scene0593_00 177 | scene0593_01 178 | scene0246_00 179 | scene0653_00 180 | scene0653_01 181 | scene0064_00 182 | scene0064_01 183 | scene0356_00 184 | scene0356_01 185 | scene0356_02 186 | scene0030_00 187 | scene0030_01 188 | scene0030_02 189 | scene0222_00 190 | scene0222_01 191 | scene0338_00 192 | scene0338_01 193 | scene0338_02 194 | scene0378_00 195 | scene0378_01 196 | scene0378_02 197 | scene0660_00 198 | scene0553_00 199 | scene0553_01 200 | scene0553_02 201 | scene0527_00 202 | scene0663_00 203 | scene0663_01 204 | scene0663_02 205 | scene0664_00 206 | scene0664_01 207 | scene0664_02 208 | scene0334_00 209 | scene0334_01 210 | scene0334_02 211 | scene0046_00 212 | scene0046_01 213 | scene0046_02 214 | scene0203_00 215 | scene0203_01 216 | scene0203_02 217 | scene0088_00 218 | scene0088_01 219 | scene0088_02 220 | scene0088_03 221 | scene0086_00 222 | scene0086_01 223 | scene0086_02 224 | scene0670_00 225 | scene0670_01 226 | scene0256_00 227 | scene0256_01 228 | scene0256_02 229 | scene0249_00 230 | scene0441_00 231 | scene0658_00 232 | scene0704_00 233 | scene0704_01 234 | scene0187_00 235 | scene0187_01 236 | scene0131_00 237 | scene0131_01 238 | scene0131_02 239 | scene0207_00 240 | scene0207_01 241 | scene0207_02 242 | scene0461_00 243 | scene0011_00 244 | scene0011_01 245 | scene0343_00 246 | scene0251_00 247 | scene0077_00 248 | scene0077_01 249 | scene0684_00 250 | scene0684_01 251 | scene0550_00 252 | scene0686_00 253 | scene0686_01 254 | scene0686_02 255 | scene0208_00 256 | scene0500_00 257 | scene0500_01 258 | scene0552_00 259 | scene0552_01 260 | scene0648_00 261 | scene0648_01 262 | scene0435_00 263 | scene0435_01 264 | scene0435_02 265 | scene0435_03 266 | scene0690_00 267 | scene0690_01 268 | scene0693_00 269 | scene0693_01 270 | scene0693_02 271 | scene0700_00 272 | scene0700_01 273 | scene0700_02 274 | scene0699_00 275 | scene0231_00 276 | scene0231_01 277 | scene0231_02 278 | scene0697_00 279 | scene0697_01 280 | scene0697_02 281 | scene0697_03 282 | scene0474_00 283 | scene0474_01 284 | scene0474_02 285 | scene0474_03 286 | scene0474_04 287 | scene0474_05 288 | scene0355_00 289 | scene0355_01 290 | scene0146_00 291 | scene0146_01 292 | scene0146_02 293 | scene0196_00 294 | scene0702_00 295 | scene0702_01 296 | scene0702_02 297 | scene0314_00 298 | scene0277_00 299 | scene0277_01 300 | scene0277_02 301 | scene0095_00 302 | scene0095_01 303 | scene0015_00 304 | scene0100_00 305 | scene0100_01 306 | scene0100_02 307 | scene0558_00 308 | scene0558_01 309 | scene0558_02 310 | scene0685_00 311 | scene0685_01 312 | scene0685_02 313 | -------------------------------------------------------------------------------- /datasets/scannet/scannetv2_weigths.txt: -------------------------------------------------------------------------------- 1 | 9.620879743473732537e-01 2 | 4.190072537849464096e-01 3 | 4.728749009943037374e-01 4 | 5.703169262148852781e-01 5 | 3.837060948436303298e-01 6 | 7.708427560138058077e-01 7 | 4.715935110518766904e-01 8 | 7.060852007691431620e-01 9 | 7.029967467438303919e-01 10 | 1.363647406054545597e+00 11 | 5.501577870523420843e-01 12 | 1.697620280852317087e+00 13 | 1.000000000000000000e+00 14 | 1.124334838541653614e+00 15 | 7.600876939475833538e-01 16 | 7.561849619423417490e-01 17 | 6.659148023979147846e-01 18 | 7.083930519465178044e-01 19 | 1.714883368336828084e+00 20 | 4.040201071782831832e+00 21 | 1.407282274103964115e+00 22 | 1.021857512741301210e+00 23 | 9.723976304568668905e-01 24 | 1.031531143846733167e+00 25 | 6.608539193406804957e-01 26 | 1.507448203607812065e+00 27 | 3.006425611099400097e+00 28 | 1.659691622772110797e+00 29 | 6.915041294677738426e-01 30 | 1.588262366884082732e+00 31 | 1.067425213835515363e+00 32 | 1.195193101628283960e+00 33 | 1.864666616576375446e+00 34 | 7.903203182740298205e-01 35 | 1.717013497455768922e+00 36 | 5.557825196540050250e+00 37 | 5.930215676867716112e-01 38 | 3.102311202028769532e+00 39 | 8.882455326032755449e-01 40 | 1.075068274166168569e+00 41 | 1.235495932675423125e+00 42 | -------------------------------------------------------------------------------- /datasets/sunrgbd/class_weights: -------------------------------------------------------------------------------- 1 | 0.351185, 0.382592, 0.447844, 0.634237, 0.375678, 0.582794, 0.476692, 0.776787, 0.980661, 1.020118, 0.623396, 2.596563, 0.974491, 0.920240, 0.663878, 1.173357, 0.861062, 0.919955, 2.147320, 1.190958, 1.164314, 1.854754, 1.446550, 2.853078, 0.768276, 1.651721, 4.456313, 2.205633, 1.116695, 2.781543, 0.705917, 4.005834, 2.208329, 0.967071, 1.475710, 5.347752, 0.745528, 4.017548 2 | -------------------------------------------------------------------------------- /environment.yml: -------------------------------------------------------------------------------- 1 | name: pytorch-CycleGAN-and-pix2pix 2 | channels: 3 | - peterjc123 4 | - defaults 5 | dependencies: 6 | - python=3.5.5 7 | - pytorch=0.4 8 | - scipy 9 | - pip: 10 | - dominate==2.3.1 11 | - git+https://github.com/pytorch/vision.git 12 | - Pillow==5.0.0 13 | - numpy==1.14.1 14 | - visdom==0.1.7 15 | -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- 1 | import importlib 2 | from models.base_model import BaseModel 3 | 4 | 5 | def find_model_using_name(model_name): 6 | # Given the option --model [modelname], 7 | # the file "models/modelname_model.py" 8 | # will be imported. 9 | model_filename = "models." + model_name + "_model" 10 | modellib = importlib.import_module(model_filename) 11 | 12 | # In the file, the class called ModelNameModel() will 13 | # be instantiated. It has to be a subclass of BaseModel, 14 | # and it is case-insensitive. 15 | model = None 16 | target_model_name = model_name.replace('_', '') + 'model' 17 | for name, cls in modellib.__dict__.items(): 18 | if name.lower() == target_model_name.lower() \ 19 | and issubclass(cls, BaseModel): 20 | model = cls 21 | 22 | if model is None: 23 | print("In %s.py, there should be a subclass of BaseModel with class name that matches %s in lowercase." % (model_filename, target_model_name)) 24 | exit(0) 25 | 26 | return model 27 | 28 | 29 | def get_option_setter(model_name): 30 | model_class = find_model_using_name(model_name) 31 | return model_class.modify_commandline_options 32 | 33 | 34 | def create_model(opt, dataset): 35 | model = find_model_using_name(opt.model) 36 | instance = model() 37 | instance.initialize(opt, dataset) 38 | print("model [%s] was created" % (instance.name())) 39 | return instance 40 | -------------------------------------------------------------------------------- /models/base_model.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | from collections import OrderedDict 4 | from . import networks 5 | 6 | 7 | class BaseModel(): 8 | 9 | # modify parser to add command line options, 10 | # and also change the default values if needed 11 | @staticmethod 12 | def modify_commandline_options(parser, is_train): 13 | return parser 14 | 15 | def name(self): 16 | return 'BaseModel' 17 | 18 | def initialize(self, opt): 19 | self.opt = opt 20 | self.gpu_ids = opt.gpu_ids 21 | self.isTrain = opt.isTrain 22 | self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu') 23 | self.save_dir = os.path.join(opt.checkpoints_dir, opt.name) 24 | if opt.resize_or_crop != 'scale_width': 25 | torch.backends.cudnn.benchmark = True 26 | self.loss_names = [] 27 | self.model_names = [] 28 | self.visual_names = [] 29 | self.image_paths = [] 30 | 31 | def set_input(self, input): 32 | self.input = input 33 | 34 | def forward(self): 35 | pass 36 | 37 | # load and print networks; create schedulers 38 | def setup(self, opt, parser=None): 39 | if self.isTrain: 40 | self.schedulers = [networks.get_scheduler(optimizer, opt) for optimizer in self.optimizers] 41 | 42 | if not self.isTrain or opt.continue_train: 43 | self.load_networks(opt.epoch) 44 | self.print_networks(opt.verbose) 45 | 46 | # make models eval mode during test time 47 | def eval(self): 48 | for name in self.model_names: 49 | if isinstance(name, str): 50 | net = getattr(self, 'net' + name) 51 | net.eval() 52 | def train(self): 53 | for name in self.model_names: 54 | if isinstance(name, str): 55 | net = getattr(self, 'net' + name) 56 | net.train() 57 | 58 | 59 | # used in test time, wrapping `forward` in no_grad() so we don't save 60 | # intermediate steps for backprop 61 | def test(self): 62 | with torch.no_grad(): 63 | self.forward() 64 | 65 | # get image paths 66 | def get_image_paths(self): 67 | return self.image_paths 68 | 69 | def optimize_parameters(self): 70 | pass 71 | 72 | # update learning rate (called once every epoch) 73 | def update_learning_rate(self): 74 | for scheduler in self.schedulers: 75 | scheduler.step() 76 | lr = self.optimizers[0].param_groups[0]['lr'] 77 | print('learning rate = %.7f' % lr) 78 | 79 | # return visualization images. train.py will display these images, and save the images to a html 80 | def get_current_visuals(self): 81 | visual_ret = OrderedDict() 82 | for name in self.visual_names: 83 | if isinstance(name, str): 84 | visual_ret[name] = getattr(self, name) 85 | return visual_ret 86 | 87 | # return traning losses/errors. train.py will print out these errors as debugging information 88 | def get_current_losses(self): 89 | errors_ret = OrderedDict() 90 | for name in self.loss_names: 91 | if isinstance(name, str): 92 | # float(...) works for both scalar tensor and float number 93 | errors_ret[name] = float(getattr(self, 'loss_' + name)) 94 | return errors_ret 95 | 96 | # save models to the disk 97 | def save_networks(self, epoch): 98 | for name in self.model_names: 99 | if isinstance(name, str): 100 | save_filename = '%s_net_%s.pth' % (epoch, name) 101 | save_path = os.path.join(self.save_dir, save_filename) 102 | net = getattr(self, 'net' + name) 103 | 104 | if len(self.gpu_ids) > 0 and torch.cuda.is_available(): 105 | torch.save(net.module.cpu().state_dict(), save_path) 106 | net.cuda(self.gpu_ids[0]) 107 | else: 108 | torch.save(net.cpu().state_dict(), save_path) 109 | 110 | def __patch_instance_norm_state_dict(self, state_dict, module, keys, i=0): 111 | key = keys[i] 112 | if i + 1 == len(keys): # at the end, pointing to a parameter/buffer 113 | if module.__class__.__name__.startswith('InstanceNorm') and \ 114 | (key == 'running_mean' or key == 'running_var'): 115 | if getattr(module, key) is None: 116 | state_dict.pop('.'.join(keys)) 117 | if module.__class__.__name__.startswith('InstanceNorm') and \ 118 | (key == 'num_batches_tracked'): 119 | state_dict.pop('.'.join(keys)) 120 | else: 121 | self.__patch_instance_norm_state_dict(state_dict, getattr(module, key), keys, i + 1) 122 | 123 | # load models from the disk 124 | def load_networks(self, epoch): 125 | for name in self.model_names: 126 | if isinstance(name, str): 127 | load_filename = '%s_net_%s.pth' % (epoch, name) 128 | load_path = os.path.join(self.save_dir, load_filename) 129 | net = getattr(self, 'net' + name) 130 | if isinstance(net, torch.nn.DataParallel): 131 | net = net.module 132 | print('loading the model from %s' % load_path) 133 | # if you are using PyTorch newer than 0.4 (e.g., built from 134 | # GitHub source), you can remove str() on self.device 135 | state_dict = torch.load(load_path, map_location=str(self.device)) 136 | if hasattr(state_dict, '_metadata'): 137 | del state_dict._metadata 138 | 139 | # patch InstanceNorm checkpoints prior to 0.4 140 | for key in list(state_dict.keys()): # need to copy keys here because we mutate in loop 141 | self.__patch_instance_norm_state_dict(state_dict, net, key.split('.')) 142 | net.load_state_dict(state_dict) 143 | 144 | # print network information 145 | def print_networks(self, verbose): 146 | print('---------- Networks initialized -------------') 147 | for name in self.model_names: 148 | if isinstance(name, str): 149 | net = getattr(self, 'net' + name) 150 | num_params = 0 151 | for param in net.parameters(): 152 | num_params += param.numel() 153 | if verbose: 154 | print(net) 155 | print('[Network %s] Total number of parameters : %.3f M' % (name, num_params / 1e6)) 156 | print('-----------------------------------------------') 157 | 158 | # set requies_grad=Fasle to avoid computation 159 | def set_requires_grad(self, nets, requires_grad=False): 160 | if not isinstance(nets, list): 161 | nets = [nets] 162 | for net in nets: 163 | if net is not None: 164 | for param in net.parameters(): 165 | param.requires_grad = requires_grad 166 | -------------------------------------------------------------------------------- /models/fusenet_model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from util.image_pool import ImagePool 3 | from .base_model import BaseModel 4 | from . import networks 5 | import numpy as np 6 | import os 7 | 8 | class FuseNetModel(BaseModel): 9 | def name(self): 10 | return 'FuseNetModel' 11 | 12 | @staticmethod 13 | def modify_commandline_options(parser, is_train=True): 14 | 15 | # changing the default values 16 | if is_train: 17 | parser.add_argument('--lambda_L1', type=float, default=100.0, help='weight for L1 loss') 18 | return parser 19 | 20 | def initialize(self, opt, dataset): 21 | BaseModel.initialize(self, opt) 22 | self.isTrain = opt.isTrain 23 | # specify the training losses you want to print out. The program will call base_model.get_current_losses 24 | self.loss_names = ['segmentation'] 25 | # specify the images you want to save/display. The program will call base_model.get_current_visuals 26 | self.visual_names = ['rgb_image','depth_image','mask','output'] 27 | # specify the models you want to save to the disk. The program will call base_model.save_networks and base_model.load_networks 28 | 29 | self.model_names = ['FuseNet'] 30 | 31 | # load/define networks 32 | self.netFuseNet = networks.define_FuseNet(dataset.num_labels, rgb_enc=True, depth_enc=True, rgb_dec=True, depth_dec=False, norm=opt.norm,use_dropout= not opt.no_dropout, init_type=opt.init_type, init_gain= opt.init_gain, gpu_ids= self.gpu_ids) 33 | # define loss functions 34 | self.criterionSegmentation = networks.SegmantationLoss(ignore_label=dataset.ignore_label, class_weights=dataset.class_weights).to(self.device) 35 | 36 | if self.isTrain: 37 | #self.criterionL1 = torch.nn.L1Loss() 38 | 39 | # initialize optimizers 40 | self.optimizers = [] 41 | 42 | self.optimizer_FuseNet = torch.optim.SGD(self.netFuseNet.parameters(), lr=opt.lr, momentum=opt.momentum, weight_decay=opt.weight_decay) 43 | #self.optimizer_FuseNet = torch.optim.Adam(self.netFuseNet.parameters(), lr=opt.lr, weight_decay=opt.weight_decay) 44 | self.optimizers.append(self.optimizer_FuseNet) 45 | 46 | def set_input(self, input): 47 | self.rgb_image = input['rgb_image'].to(self.device) 48 | self.depth_image = input['depth_image'].to(self.device) 49 | self.mask = input['mask'].to(self.device) 50 | self.image_paths = input['path'] 51 | 52 | def forward(self): 53 | self.output = self.netFuseNet(self.rgb_image,self.depth_image) 54 | 55 | def get_loss(self): 56 | self.loss_segmentation = self.criterionSegmentation(self.output, self.mask) 57 | 58 | def backward(self): 59 | 60 | self.loss_segmentation.backward() 61 | 62 | def optimize_parameters(self): 63 | self.forward() 64 | # update Fusenet 65 | self.set_requires_grad(self.netFuseNet, True) 66 | self.optimizer_FuseNet.zero_grad() 67 | self.get_loss() 68 | self.backward() 69 | self.optimizer_FuseNet.step() 70 | -------------------------------------------------------------------------------- /models/networks.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.nn import init 4 | import torchvision 5 | import functools 6 | from torch.optim import lr_scheduler 7 | import torch.nn.functional as F 8 | from copy import deepcopy 9 | 10 | 11 | ############################################################################### 12 | # Helper Functions 13 | ############################################################################### 14 | 15 | 16 | def get_norm_layer(norm_type='instance'): 17 | if norm_type == 'batch': 18 | norm_layer = functools.partial(nn.BatchNorm2d, affine=True) 19 | elif norm_type == 'instance': 20 | norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=True) 21 | elif norm_type == 'none': 22 | norm_layer = None 23 | else: 24 | raise NotImplementedError('normalization layer [%s] is not found' % norm_type) 25 | return norm_layer 26 | 27 | 28 | def get_scheduler(optimizer, opt): 29 | if opt.lr_policy == 'lambda': 30 | lambda_rule = lambda epoch: opt.lr_gamma ** ((epoch+1) // opt.lr_decay_epochs) 31 | # def lambda_rule(epoch): 32 | # lr_l = 1.0 - max(0, epoch + 1 + opt.epoch_count - opt.niter) / float(opt.niter_decay + 1) 33 | # return lr_l 34 | scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule) 35 | elif opt.lr_policy == 'step': 36 | scheduler = lr_scheduler.StepLR(optimizer,step_size=opt.lr_decay_iters, gamma=0.1) 37 | else: 38 | return NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy) 39 | return scheduler 40 | 41 | 42 | def init_weights(net, init_type='normal', gain=0.02): 43 | net = net 44 | def init_func(m): 45 | classname = m.__class__.__name__ 46 | if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1): 47 | if init_type == 'normal': 48 | init.normal_(m.weight.data, 0.0, gain) 49 | elif init_type == 'xavier': 50 | init.xavier_normal_(m.weight.data, gain=gain) 51 | elif init_type == 'kaiming': 52 | init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') 53 | elif init_type == 'orthogonal': 54 | init.orthogonal_(m.weight.data, gain=gain) 55 | elif init_type == 'pretrained': 56 | pass 57 | else: 58 | raise NotImplementedError('initialization method [%s] is not implemented' % init_type) 59 | if hasattr(m, 'bias') and m.bias is not None and init_type != 'pretrained': 60 | init.constant_(m.bias.data, 0.0) 61 | elif classname.find('BatchNorm2d') != -1: 62 | init.normal_(m.weight.data, 1.0, gain) 63 | init.constant_(m.bias.data, 0.0) 64 | print('initialize network with %s' % init_type) 65 | net.apply(init_func) 66 | 67 | 68 | def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]): 69 | if len(gpu_ids) > 0: 70 | assert(torch.cuda.is_available()) 71 | net.to(gpu_ids[0]) 72 | net = torch.nn.DataParallel(net, gpu_ids) 73 | 74 | for root_child in net.children(): 75 | for children in root_child.children(): 76 | if children in root_child.need_initialization: 77 | init_weights(children, init_type, gain=init_gain) 78 | else: 79 | init_weights(children,"pretrained",gain=init_gain) #for batchnorms 80 | return net 81 | 82 | def define_FuseNet(num_labels, rgb_enc=True, depth_enc=True, rgb_dec=True, depth_dec=False, norm='batch', use_dropout=True, init_type='xavier', init_gain=0.02, gpu_ids=[]): 83 | net = None 84 | norm_layer = get_norm_layer(norm_type=norm) 85 | 86 | net = FusenetGenerator(num_labels, rgb_enc=True, depth_enc=True, rgb_dec=True, depth_dec=False ) 87 | return init_net(net, init_type, init_gain, gpu_ids) 88 | 89 | 90 | ############################################################################## 91 | # Classes 92 | ############################################################################## 93 | 94 | def VGG16_initializator(): 95 | layer_names =["conv1_1","conv1_2","conv2_1","conv2_2","conv3_1","conv3_2","conv3_3","conv4_1","conv4_2","conv4_3","conv5_1","conv5_2","conv5_3"] 96 | layers = list(torchvision.models.vgg16(pretrained=True).features.children()) 97 | layers = [x for x in layers if isinstance(x, nn.Conv2d)] 98 | layer_dic = dict(zip(layer_names,layers)) 99 | return layer_dic 100 | 101 | def make_layers_from_names(names,model_dic,bn_dim,existing_layer=None): 102 | 103 | layers = [] 104 | if existing_layer is not None: 105 | layers = [existing_layer,nn.BatchNorm2d(bn_dim,momentum = 0.1),nn.ReLU(inplace=True)] 106 | for name in names: 107 | layers += [deepcopy(model_dic[name]), nn.BatchNorm2d(bn_dim,momentum = 0.1), nn.ReLU(inplace=True)] 108 | 109 | return nn.Sequential(*layers) 110 | 111 | def make_layers_from_size(sizes): 112 | layers = [] 113 | for size in sizes: 114 | layers += [nn.Conv2d(size[0], size[1], kernel_size=3, padding=1), nn.BatchNorm2d(size[1],momentum = 0.1), nn.ReLU(inplace=True)] 115 | return nn.Sequential(*layers) 116 | 117 | class MyUpsample2(nn.Module): 118 | def forward(self, x): 119 | return x[:, :, :, None, :, None].expand(-1, -1, -1, 2, -1, 2).reshape(x.size(0), x.size(1), x.size(2)*2, x.size(3)*2) 120 | 121 | class FusenetGenerator(nn.Module): 122 | def __init__(self, num_labels, rgb_enc=True, depth_enc=True, rgb_dec=True, depth_dec=False): 123 | super(FusenetGenerator, self).__init__() 124 | batchNorm_momentum = 0.1#TODO:make param 125 | 126 | self.need_initialization = [] #modules that need initialization 127 | model_dic = VGG16_initializator() 128 | 129 | if rgb_enc : 130 | 131 | ##### RGB ENCODER #### 132 | self.CBR1_RGB_ENC = make_layers_from_names(["conv1_1","conv1_2"], model_dic, 64) 133 | self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2, return_indices=True) 134 | 135 | self.CBR2_RGB_ENC = make_layers_from_names(["conv2_1","conv2_2"], model_dic, 128) 136 | self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2, return_indices=True) 137 | 138 | self.CBR3_RGB_ENC = make_layers_from_names(["conv3_1","conv3_2","conv3_3"], model_dic, 256) 139 | self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2, return_indices=True) 140 | self.dropout3 = nn.Dropout(p=0.4) 141 | 142 | self.CBR4_RGB_ENC = make_layers_from_names(["conv4_1","conv4_2","conv4_3"], model_dic, 512) 143 | self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2, return_indices=True) 144 | self.dropout4 = nn.Dropout(p=0.4) 145 | 146 | self.CBR5_RGB_ENC = make_layers_from_names(["conv5_1","conv5_2","conv5_3"], model_dic, 512) 147 | self.dropout5 = nn.Dropout(p=0.4) 148 | 149 | self.pool5 = nn.MaxPool2d(kernel_size=2, stride=2, return_indices=True) 150 | 151 | if depth_enc : 152 | feats_depth = list(torchvision.models.vgg16(pretrained=True).features.children()) 153 | avg = torch.mean(feats_depth[0].weight.data, dim=1) 154 | avg = avg.unsqueeze(1) 155 | 156 | conv11d = nn.Conv2d(1, 64, kernel_size=3,padding=1) 157 | conv11d.weight.data = avg 158 | 159 | self.CBR1_DEPTH_ENC = make_layers_from_names(["conv1_2"], model_dic, 64, conv11d) 160 | self.pool1_d = nn.MaxPool2d(kernel_size=2, stride=2, return_indices=True) 161 | 162 | self.CBR2_DEPTH_ENC = make_layers_from_names(["conv2_1","conv2_2"], model_dic, 128) 163 | self.pool2_d = nn.MaxPool2d(kernel_size=2, stride=2, return_indices=True) 164 | 165 | self.CBR3_DEPTH_ENC = make_layers_from_names(["conv3_1","conv3_2","conv3_3"], model_dic, 256) 166 | self.pool3_d = nn.MaxPool2d(kernel_size=2, stride=2, return_indices=True) 167 | self.dropout3_d = nn.Dropout(p=0.4) 168 | 169 | self.CBR4_DEPTH_ENC = make_layers_from_names(["conv4_1","conv4_2","conv4_3"], model_dic, 512) 170 | self.pool4_d = nn.MaxPool2d(kernel_size=2, stride=2, return_indices=True) 171 | self.dropout4_d = nn.Dropout(p=0.4) 172 | 173 | self.CBR5_DEPTH_ENC = make_layers_from_names(["conv5_1","conv5_2","conv5_3"], model_dic, 512) 174 | 175 | if rgb_dec : 176 | #### RGB DECODER #### 177 | 178 | self.unpool5 = nn.MaxUnpool2d(kernel_size=2, stride=2) 179 | self.CBR5_RGB_DEC = make_layers_from_size([[512,512],[512,512],[512,512]]) 180 | self.dropout5_dec = nn.Dropout(p=0.4) 181 | 182 | self.need_initialization.append(self.CBR5_RGB_DEC) 183 | 184 | self.unpool4 = nn.MaxUnpool2d(kernel_size=2, stride=2) 185 | self.CBR4_RGB_DEC = make_layers_from_size([[512,512],[512,512],[512,256]]) 186 | self.dropout4_dec = nn.Dropout(p=0.4) 187 | 188 | self.need_initialization.append(self.CBR4_RGB_DEC) 189 | 190 | self.unpool3 = nn.MaxUnpool2d(kernel_size=2, stride=2) 191 | self.CBR3_RGB_DEC = make_layers_from_size([[256,256],[256,256],[256,128]]) 192 | self.dropout3_dec = nn.Dropout(p=0.4) 193 | 194 | self.need_initialization.append(self.CBR3_RGB_DEC) 195 | 196 | self.unpool2 = nn.MaxUnpool2d(kernel_size=2, stride=2) 197 | self.CBR2_RGB_DEC = make_layers_from_size([[128,128],[128,64]]) 198 | 199 | self.need_initialization.append(self.CBR2_RGB_DEC) 200 | 201 | self.unpool1 = nn.MaxUnpool2d(kernel_size=2, stride=2) 202 | self.CBR1_RGB_DEC = nn.Sequential ( 203 | nn.Conv2d(64, 64, kernel_size=3, padding=1), 204 | nn.BatchNorm2d(64, momentum= batchNorm_momentum), 205 | nn.ReLU(), 206 | nn.Conv2d(64, num_labels, kernel_size=3, padding=1), 207 | ) 208 | 209 | self.need_initialization.append(self.CBR1_RGB_DEC) 210 | 211 | def forward(self, rgb_inputs, depth_inputs): 212 | 213 | ######## DEPTH ENCODER ######## 214 | # Stage 1 215 | #x = self.conv11d(depth_inputs) 216 | x_1 = self.CBR1_DEPTH_ENC(depth_inputs) 217 | x, id1_d = self.pool1_d(x_1) 218 | 219 | # Stage 2 220 | x_2 = self.CBR2_DEPTH_ENC(x) 221 | x, id2_d = self.pool2_d(x_2) 222 | 223 | # Stage 3 224 | x_3 = self.CBR3_DEPTH_ENC(x) 225 | x, id3_d = self.pool4_d(x_3) 226 | x = self.dropout3_d(x) 227 | 228 | # Stage 4 229 | x_4 = self.CBR4_DEPTH_ENC(x) 230 | x, id4_d = self.pool4_d(x_4) 231 | x = self.dropout4_d(x) 232 | 233 | # Stage 5 234 | x_5 = self.CBR5_DEPTH_ENC(x) 235 | 236 | ######## RGB ENCODER ######## 237 | 238 | # Stage 1 239 | y = self.CBR1_RGB_ENC(rgb_inputs) 240 | y = torch.add(y,x_1) 241 | y = torch.div(y,2) 242 | y, id1 = self.pool1(y) 243 | 244 | # Stage 2 245 | y = self.CBR2_RGB_ENC(y) 246 | y = torch.add(y,x_2) 247 | y = torch.div(y,2) 248 | y, id2 = self.pool2(y) 249 | 250 | # Stage 3 251 | y = self.CBR3_RGB_ENC(y) 252 | y = torch.add(y,x_3) 253 | y = torch.div(y,2) 254 | y, id3 = self.pool3(y) 255 | y = self.dropout3(y) 256 | 257 | # Stage 4 258 | y = self.CBR4_RGB_ENC(y) 259 | y = torch.add(y,x_4) 260 | y = torch.div(y,2) 261 | y, id4 = self.pool4(y) 262 | y = self.dropout4(y) 263 | 264 | # Stage 5 265 | y = self.CBR5_RGB_ENC(y) 266 | y = torch.add(y,x_5) 267 | y = torch.div(y,2) 268 | y_size = y.size() 269 | 270 | y, id5 = self.pool5(y) 271 | y = self.dropout5(y) 272 | 273 | ######## DECODER ######## 274 | 275 | # Stage 5 dec 276 | y = self.unpool5(y, id5,output_size=y_size) 277 | y = self.CBR5_RGB_DEC(y) 278 | y = self.dropout5_dec(y) 279 | 280 | # Stage 4 dec 281 | y = self.unpool4(y, id4) 282 | y = self.CBR4_RGB_DEC(y) 283 | y = self.dropout4_dec(y) 284 | 285 | # Stage 3 dec 286 | y = self.unpool3(y, id3) 287 | y = self.CBR3_RGB_DEC(y) 288 | y = self.dropout3_dec(y) 289 | 290 | # Stage 2 dec 291 | y = self.unpool2(y, id2) 292 | y = self.CBR2_RGB_DEC(y) 293 | 294 | # Stage 1 dec 295 | y = self.unpool1(y, id1) 296 | y = self.CBR1_RGB_DEC(y) 297 | 298 | return y 299 | 300 | class FusenetGeneratorTest(nn.Module): 301 | def __init__(self, num_labels, rgb_enc=True, depth_enc=True, rgb_dec=True, depth_dec=False): 302 | super(FusenetGeneratorTest, self).__init__() 303 | batchNorm_momentum = 0.1#TODO:make param 304 | self.c = nn.Conv2d(3, num_labels, kernel_size=3, padding=1) 305 | self.need_initialization = [] 306 | 307 | def forward(self, rgb_inputs, depth_inputs): 308 | 309 | ######## DEPTH ENCODER ######## 310 | # Stage 1 311 | #x = self.conv11d(depth_inputs) 312 | y = self.c(rgb_inputs) 313 | return y 314 | 315 | class SegmantationLoss(nn.Module): 316 | def __init__(self, ignore_label=None, class_weights=None): 317 | super(SegmantationLoss, self).__init__() 318 | self.loss = nn.CrossEntropyLoss(ignore_index=ignore_label, weight=class_weights) 319 | def __call__(self, output, target, pixel_average=True): 320 | if pixel_average: 321 | return self.loss(output, target) #/ target.data.sum() 322 | else: 323 | return self.loss(output, target) 324 | -------------------------------------------------------------------------------- /options/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MehmetAygun/fusenet-pytorch/26d28b6cf40a6be28ad7e2190a0abe1f0f765b09/options/__init__.py -------------------------------------------------------------------------------- /options/base_options.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | from util import util 4 | import torch 5 | import models 6 | import data 7 | 8 | 9 | class BaseOptions(): 10 | def __init__(self): 11 | self.initialized = False 12 | 13 | def initialize(self, parser): 14 | parser.add_argument('--dataroot', required=True, help='path to images (should have subfolders trainA, trainB, valA, valB, etc)') 15 | parser.add_argument('--batch_size', type=int, default=4, help='input batch size') 16 | parser.add_argument('--loadSize', type=int, default=224, help='scale images to this size') 17 | parser.add_argument('--fineSize', type=int, default=224, help='then crop to this size') 18 | parser.add_argument('--display_winsize', type=int, default=224, help='display window size for both visdom and HTML') 19 | parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0 0,1,2, 0,2. use -1 for CPU') 20 | parser.add_argument('--name', type=str, default='experiment_name', help='name of the experiment. It decides where to store samples and models') 21 | parser.add_argument('--dataset_mode', type=str, default='unaligned', help='chooses how datasets are loaded. [unaligned | aligned | single]') 22 | parser.add_argument('--model', type=str, default='fusenet', 23 | help='chooses which model to use.') 24 | #parser.add_argument('--direction', type=str, default='AtoB', help='AtoB or BtoA') 25 | parser.add_argument('--epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model') 26 | parser.add_argument('--num_threads', default=8, type=int, help='# threads for loading data') 27 | parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here') 28 | parser.add_argument('--norm', type=str, default='instance', help='instance normalization or batch normalization') 29 | parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly') 30 | parser.add_argument('--no_dropout', action='store_true', help='no dropout for the generator') 31 | parser.add_argument('--max_dataset_size', type=int, default=float("inf"), help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.') 32 | parser.add_argument('--resize_or_crop', type=str, default='none', help='scaling and cropping of images at load time [resize_and_crop|crop|scale_width|scale_width_and_crop|none]') 33 | parser.add_argument('--no_flip', action='store_true', default=True, help='if specified, do not flip the images for data augmentation') 34 | parser.add_argument('--init_type', type=str, default='kaiming', help='network initialization [normal|xavier|kaiming|orthogonal]') 35 | parser.add_argument('--init_gain', type=float, default=0.02, help='scaling factor for normal, xavier and orthogonal.') 36 | parser.add_argument('--verbose', action='store_true', help='if specified, print more debugging information') 37 | parser.add_argument('--suffix', default='', type=str, help='customized suffix: opt.name = opt.name + suffix: e.g., {model}_{netG}_size{loadSize}') 38 | parser.add_argument('--seed', type=int, default=0, help='seed for random generators') 39 | self.initialized = True 40 | return parser 41 | 42 | def gather_options(self): 43 | # initialize parser with basic options 44 | if not self.initialized: 45 | parser = argparse.ArgumentParser( 46 | formatter_class=argparse.ArgumentDefaultsHelpFormatter) 47 | parser = self.initialize(parser) 48 | 49 | # get the basic options 50 | opt, _ = parser.parse_known_args() 51 | 52 | # modify model-related parser options 53 | model_name = opt.model 54 | model_option_setter = models.get_option_setter(model_name) 55 | parser = model_option_setter(parser, self.isTrain) 56 | opt, _ = parser.parse_known_args() # parse again with the new defaults 57 | 58 | # modify dataset-related parser options 59 | dataset_name = opt.dataset_mode 60 | dataset_option_setter = data.get_option_setter(dataset_name) 61 | parser = dataset_option_setter(parser, self.isTrain) 62 | 63 | self.parser = parser 64 | 65 | return parser.parse_args() 66 | 67 | def print_options(self, opt): 68 | message = '' 69 | message += '----------------- Options ---------------\n' 70 | for k, v in sorted(vars(opt).items()): 71 | comment = '' 72 | default = self.parser.get_default(k) 73 | if v != default: 74 | comment = '\t[default: %s]' % str(default) 75 | message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment) 76 | message += '----------------- End -------------------' 77 | print(message) 78 | 79 | # save to the disk 80 | expr_dir = os.path.join(opt.checkpoints_dir, opt.name) 81 | util.mkdirs(expr_dir) 82 | file_name = os.path.join(expr_dir, 'opt.txt') 83 | with open(file_name, 'wt') as opt_file: 84 | opt_file.write(message) 85 | opt_file.write('\n') 86 | 87 | def parse(self): 88 | 89 | opt = self.gather_options() 90 | opt.isTrain = self.isTrain # train or test 91 | 92 | # process opt.suffix 93 | if opt.suffix: 94 | suffix = ('_' + opt.suffix.format(**vars(opt))) if opt.suffix != '' else '' 95 | opt.name = opt.name + suffix 96 | 97 | self.print_options(opt) 98 | 99 | # set gpu ids 100 | str_ids = opt.gpu_ids.split(',') 101 | opt.gpu_ids = [] 102 | for str_id in str_ids: 103 | id = int(str_id) 104 | if id >= 0: 105 | opt.gpu_ids.append(id) 106 | if len(opt.gpu_ids) > 0: 107 | torch.cuda.set_device(opt.gpu_ids[0]) 108 | 109 | self.opt = opt 110 | return self.opt 111 | -------------------------------------------------------------------------------- /options/test_options.py: -------------------------------------------------------------------------------- 1 | from .base_options import BaseOptions 2 | 3 | 4 | class TestOptions(BaseOptions): 5 | def initialize(self, parser): 6 | parser = BaseOptions.initialize(self, parser) 7 | parser.add_argument('--ntest', type=int, default=float("inf"), help='# of test examples.') 8 | parser.add_argument('--results_dir', type=str, default='./checkpoints/', help='saves results here.') 9 | parser.add_argument('--aspect_ratio', type=float, default=1.0, help='aspect ratio of result images') 10 | parser.add_argument('--phase', type=str, default='test', help='train, val, test, etc') 11 | parser.set_defaults(model='fusenet') 12 | # To avoid cropping, the loadSize should be the same as fineSize 13 | parser.set_defaults(loadSize=parser.get_default('fineSize')) 14 | self.isTrain = False 15 | return parser 16 | -------------------------------------------------------------------------------- /options/train_options.py: -------------------------------------------------------------------------------- 1 | from .base_options import BaseOptions 2 | 3 | 4 | class TrainOptions(BaseOptions): 5 | def initialize(self, parser): 6 | parser = BaseOptions.initialize(self, parser) 7 | parser.add_argument('--display_freq', type=int, default=100, help='frequency of showing training results on screen') 8 | parser.add_argument('--display_ncols', type=int, default=4, help='if positive, display all images in a single visdom web panel with certain number of images per row.') 9 | parser.add_argument('--display_id', type=int, default=1, help='window id of the web display') 10 | parser.add_argument('--display_server', type=str, default="http://localhost", help='visdom server of the web display') 11 | parser.add_argument('--display_env', type=str, default='main', help='visdom display environment name (default is "main")') 12 | parser.add_argument('--display_port', type=int, default=8097, help='visdom port of the web display') 13 | parser.add_argument('--update_html_freq', type=int, default=100, help='frequency of saving training results to html') 14 | parser.add_argument('--print_freq', type=int, default=100, help='frequency of showing training results on console') 15 | parser.add_argument('--save_epoch_freq', type=int, default=10, help='frequency of saving checkpoints at the end of epochs') 16 | parser.add_argument('--continue_train', action='store_true', help='continue training: load the latest model') 17 | parser.add_argument('--epoch_count', type=int, default=1, help='the starting epoch count, we save the model by , +, ...') 18 | parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc') 19 | parser.add_argument('--niter', type=int, default=400, help='# of iter at starting learning rate') 20 | parser.add_argument('--beta1', type=float, default=0.5, help='momentum term of adam') 21 | parser.add_argument('--no_html', action='store_true', help='do not save intermediate training results to [opt.checkpoints_dir]/[opt.name]/web/') 22 | parser.add_argument('--lr', type=float, default=0.001, help='initial learning rate for optimizer') 23 | parser.add_argument('--momentum', type=float, default=0.9, help='momentum factor for SGD') 24 | parser.add_argument('--weight_decay', type=float, default=0.0005, help='momentum factor for optimizer') 25 | parser.add_argument('--lr_policy', type=str, default='lambda', help='learning rate policy: lambda|step|plateau|cosine') 26 | parser.add_argument('--lr_decay_iters', type=int, default=5000000, help='multiply by a gamma every lr_decay_iters iterations') 27 | parser.add_argument('--lr_decay_epochs', type=int, default=25, help='multiply by a gamma every lr_decay_epoch epochs') 28 | parser.add_argument('--lr_gamma', type=float, default=0.9, help='gamma factor for lr_scheduler') 29 | self.isTrain = True 30 | return parser 31 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | backcall==0.1.0 2 | certifi==2018.10.15 3 | chardet==3.0.4 4 | decorator==4.3.0 5 | dominate==2.3.4 6 | h5py==2.8.0 7 | idna==2.7 8 | ipython==7.16.3 9 | ipython-genutils==0.2.0 10 | jedi==0.13.1 11 | numpy==1.21.0 12 | opencv-python==4.2.0.32 13 | parso==0.3.1 14 | pexpect==4.6.0 15 | pickleshare==0.7.5 16 | Pillow==9.0.0 17 | prompt-toolkit==2.0.6 18 | ptyprocess==0.6.0 19 | Pygments==2.7.4 20 | pyzmq==17.1.2 21 | requests==2.20.0 22 | scipy==1.1.0 23 | simplegeneric==0.8.1 24 | six==1.11.0 25 | torch==0.4.1.post2 26 | torchfile==0.1.0 27 | torchvision==0.2.1 28 | tornado==5.1.1 29 | traitlets==4.3.2 30 | urllib3==1.26.5 31 | visdom==0.1.8.5 32 | wcwidth==0.1.7 33 | websocket-client==0.53.0 34 | -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | import os 2 | from options.test_options import TestOptions 3 | from data import CreateDataLoader 4 | from models import create_model 5 | from util.visualizer import Visualizer 6 | from util.visualizer import save_images,save_scannet_prediction 7 | from util.util import confusion_matrix, getScores 8 | from util import html 9 | import torch 10 | import numpy as np 11 | import cv2 12 | 13 | if __name__ == '__main__': 14 | opt = TestOptions().parse() 15 | # hard-code some parameters for test 16 | # opt.num_threads = 1 # test code only supports num_threads = 1 17 | opt.batch_size = 1 # test code only supports batch_size = 1 18 | opt.serial_batches = True # no shuffle 19 | opt.no_flip = True # no flip 20 | opt.display_id = -1 # no visdom display 21 | 22 | if opt.dataset_mode == "scannetv2" and opt.phase == "test": 23 | save_dir = os.path.join(opt.results_dir, opt.name, opt.phase + '_' + opt.epoch, 'prediction') 24 | if not os.path.exists(save_dir): 25 | os.makedirs(save_dir) 26 | 27 | data_loader = CreateDataLoader(opt) 28 | dataset = data_loader.load_data() 29 | model = create_model(opt, dataset.dataset) 30 | model.setup(opt) 31 | model.eval() 32 | visualizer = Visualizer(opt) 33 | 34 | # create a website 35 | web_dir = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.epoch)) 36 | webpage = html.HTML(web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (opt.name, opt.phase, opt.epoch)) 37 | 38 | test_loss_iter = [] 39 | epoch_iter = 0 40 | conf_mat = np.zeros((dataset.dataset.num_labels, dataset.dataset.num_labels), dtype=np.float) 41 | with torch.no_grad(): 42 | for i, data in enumerate(dataset): 43 | model.set_input(data) 44 | model.forward() 45 | model.get_loss() 46 | epoch_iter += opt.batch_size 47 | gt = model.mask.cpu().int().numpy() 48 | _, pred = torch.max(model.output.data.cpu(), 1) 49 | pred = pred.float().detach().int().numpy() 50 | if dataset.dataset.name() == 'Scannetv2': 51 | gt = data["mask_fullsize"].cpu().int().numpy()[0] 52 | pred = cv2.resize(pred[0], (gt.shape[1], gt.shape[0]), interpolation=cv2.INTER_NEAREST) 53 | if opt.phase == "test": 54 | save_scannet_prediction(pred, data['scan'][0], data['path'][0], save_dir) 55 | save_images(webpage, model.get_current_visuals(), model.get_image_paths()) 56 | conf_mat += confusion_matrix(gt, pred, dataset.dataset.num_labels, ignore_label=dataset.dataset.ignore_label) 57 | test_loss_iter.append(model.loss_segmentation.cpu().numpy()) 58 | print('Epoch {0:}, iters: {1:}/{2:}, loss: {3:.3f} '.format(opt.epoch, 59 | epoch_iter, 60 | len(dataset) * opt.batch_size, 61 | test_loss_iter[-1]), end='\r') 62 | 63 | avg_test_loss = np.mean(test_loss_iter) 64 | print ('Epoch {0:} test loss: {1:.3f} '.format(opt.epoch, avg_test_loss)) 65 | glob,mean,iou = getScores(conf_mat) 66 | print ('Epoch {0:} glob acc : {1:.2f}, mean acc : {2:.2f}, IoU : {3:.2f}'.format(opt.epoch, glob, mean, iou)) 67 | print('Confusim matrix is saved to ' + visualizer.conf_mat_name) 68 | visualizer.save_confusion_matrix(conf_mat, opt.epoch) 69 | 70 | # save the website 71 | webpage.save() 72 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | #SBATCH --job-name=fusenet 3 | #SBATCH --nodes=1 4 | #SBATCH --cpus=4 5 | #SBATCH --gres=gpu:1 6 | #SBATCH --time="UNLIMITED" 7 | 8 | import time 9 | from options.train_options import TrainOptions 10 | from data import CreateDataLoader 11 | from models import create_model 12 | from util.visualizer import Visualizer 13 | # from util.visualize_mask import * 14 | from util.util import confusion_matrix, getScores 15 | import numpy as np 16 | import random 17 | import torch 18 | import cv2 19 | 20 | if __name__ == '__main__': 21 | train_opt = TrainOptions().parse() 22 | 23 | np.random.seed(train_opt.seed) 24 | random.seed(train_opt.seed) 25 | torch.manual_seed(train_opt.seed) 26 | torch.backends.cudnn.deterministic = True 27 | torch.backends.cudnn.benchmark = False 28 | torch.cuda.manual_seed(train_opt.seed) 29 | 30 | train_data_loader = CreateDataLoader(train_opt) 31 | train_dataset = train_data_loader.load_data() 32 | train_dataset_size = len(train_data_loader) 33 | print('#training images = %d' % train_dataset_size) 34 | 35 | test_opt = TrainOptions().parse() 36 | test_opt.phase = 'val' 37 | test_opt.batch_size = 1 38 | test_opt.num_threads = 1 39 | test_opt.serial_batches = True 40 | test_opt.no_flip = True 41 | test_opt.display_id = -1 42 | test_data_loader = CreateDataLoader(test_opt) 43 | test_dataset = test_data_loader.load_data() 44 | test_dataset_size = len(test_data_loader) 45 | print('#test images = %d' % test_dataset_size) 46 | 47 | model = create_model(train_opt, train_dataset.dataset) 48 | model.setup(train_opt) 49 | visualizer = Visualizer(train_opt) 50 | total_steps = 0 51 | for epoch in range(train_opt.epoch_count, train_opt.niter + 1): 52 | model.train() 53 | epoch_start_time = time.time() 54 | iter_data_time = time.time() 55 | epoch_iter = 0 56 | for i, data in enumerate(train_dataset): 57 | iter_start_time = time.time() 58 | if total_steps % train_opt.print_freq == 0: 59 | t_data = iter_start_time - iter_data_time 60 | visualizer.reset() 61 | total_steps += train_opt.batch_size 62 | epoch_iter += train_opt.batch_size 63 | model.set_input(data) 64 | model.optimize_parameters() 65 | 66 | if total_steps % train_opt.display_freq == 0: 67 | save_result = total_steps % train_opt.update_html_freq == 0 68 | visualizer.display_current_results(model.get_current_visuals(), epoch, save_result) 69 | 70 | if total_steps % train_opt.print_freq == 0: 71 | losses = model.get_current_losses() 72 | t = (time.time() - iter_start_time) / train_opt.batch_size 73 | visualizer.print_current_losses(epoch, epoch_iter, losses, t, t_data) 74 | if train_opt.display_id > 0: 75 | visualizer.plot_current_losses(epoch, 76 | float(epoch_iter) / train_dataset_size, train_opt, losses) 77 | 78 | iter_data_time = time.time() 79 | 80 | print('End of epoch %d / %d \t Time Taken: %d sec' % (epoch, train_opt.niter, time.time() - epoch_start_time)) 81 | model.update_learning_rate() 82 | if epoch > 100 and epoch % train_opt.save_epoch_freq == 0: 83 | print('saving the model at the end of epoch %d, iters %d' % (epoch, total_steps)) 84 | model.save_networks('latest') 85 | model.save_networks(epoch) 86 | 87 | model.eval() 88 | test_loss_iter = [] 89 | gts = None 90 | preds = None 91 | epoch_iter = 0 92 | conf_mat = np.zeros((test_dataset.dataset.num_labels, test_dataset.dataset.num_labels), dtype=np.float) 93 | with torch.no_grad(): 94 | for i, data in enumerate(test_dataset): 95 | model.set_input(data) 96 | model.forward() 97 | model.get_loss() 98 | epoch_iter += test_opt.batch_size 99 | gt = model.mask.cpu().int().numpy() 100 | _, pred = torch.max(model.output.data.cpu(), 1) 101 | pred = pred.float().detach().int().numpy() 102 | if test_dataset.dataset.name() == 'Scannetv2': 103 | gt = data["mask_fullsize"].cpu().int().numpy()[0] 104 | pred = cv2.resize(pred[0], (gt.shape[1], gt.shape[0]), interpolation=cv2.INTER_NEAREST) 105 | conf_mat += confusion_matrix(gt, pred, test_dataset.dataset.num_labels, ignore_label=test_dataset.dataset.ignore_label) 106 | # visualizer.display_current_results(model.get_current_visuals(), epoch, False) 107 | losses = model.get_current_losses() 108 | test_loss_iter.append(model.loss_segmentation) 109 | print('test epoch {0:}, iters: {1:}/{2:} '.format(epoch, epoch_iter, len(test_dataset) * test_opt.batch_size), end='\r') 110 | 111 | avg_test_loss = np.mean(test_loss_iter) 112 | glob,mean,iou = getScores(conf_mat) 113 | visualizer.print_current_scores(epoch, avg_test_loss, glob, mean, iou) 114 | visualizer.save_confusion_matrix(conf_mat, epoch) 115 | -------------------------------------------------------------------------------- /util/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MehmetAygun/fusenet-pytorch/26d28b6cf40a6be28ad7e2190a0abe1f0f765b09/util/__init__.py -------------------------------------------------------------------------------- /util/get_data.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import os 3 | import tarfile 4 | import requests 5 | from warnings import warn 6 | from zipfile import ZipFile 7 | from bs4 import BeautifulSoup 8 | from os.path import abspath, isdir, join, basename 9 | 10 | 11 | class GetData(object): 12 | """ 13 | 14 | Download CycleGAN or Pix2Pix Data. 15 | 16 | Args: 17 | technique : str 18 | One of: 'cyclegan' or 'pix2pix'. 19 | verbose : bool 20 | If True, print additional information. 21 | 22 | Examples: 23 | >>> from util.get_data import GetData 24 | >>> gd = GetData(technique='cyclegan') 25 | >>> new_data_path = gd.get(save_path='./datasets') # options will be displayed. 26 | 27 | """ 28 | 29 | def __init__(self, technique='cyclegan', verbose=True): 30 | url_dict = { 31 | 'pix2pix': 'https://people.eecs.berkeley.edu/~tinghuiz/projects/pix2pix/datasets', 32 | 'cyclegan': 'https://people.eecs.berkeley.edu/~taesung_park/CycleGAN/datasets' 33 | } 34 | self.url = url_dict.get(technique.lower()) 35 | self._verbose = verbose 36 | 37 | def _print(self, text): 38 | if self._verbose: 39 | print(text) 40 | 41 | @staticmethod 42 | def _get_options(r): 43 | soup = BeautifulSoup(r.text, 'lxml') 44 | options = [h.text for h in soup.find_all('a', href=True) 45 | if h.text.endswith(('.zip', 'tar.gz'))] 46 | return options 47 | 48 | def _present_options(self): 49 | r = requests.get(self.url) 50 | options = self._get_options(r) 51 | print('Options:\n') 52 | for i, o in enumerate(options): 53 | print("{0}: {1}".format(i, o)) 54 | choice = input("\nPlease enter the number of the " 55 | "dataset above you wish to download:") 56 | return options[int(choice)] 57 | 58 | def _download_data(self, dataset_url, save_path): 59 | if not isdir(save_path): 60 | os.makedirs(save_path) 61 | 62 | base = basename(dataset_url) 63 | temp_save_path = join(save_path, base) 64 | 65 | with open(temp_save_path, "wb") as f: 66 | r = requests.get(dataset_url) 67 | f.write(r.content) 68 | 69 | if base.endswith('.tar.gz'): 70 | obj = tarfile.open(temp_save_path) 71 | elif base.endswith('.zip'): 72 | obj = ZipFile(temp_save_path, 'r') 73 | else: 74 | raise ValueError("Unknown File Type: {0}.".format(base)) 75 | 76 | self._print("Unpacking Data...") 77 | obj.extractall(save_path) 78 | obj.close() 79 | os.remove(temp_save_path) 80 | 81 | def get(self, save_path, dataset=None): 82 | """ 83 | 84 | Download a dataset. 85 | 86 | Args: 87 | save_path : str 88 | A directory to save the data to. 89 | dataset : str, optional 90 | A specific dataset to download. 91 | Note: this must include the file extension. 92 | If None, options will be presented for you 93 | to choose from. 94 | 95 | Returns: 96 | save_path_full : str 97 | The absolute path to the downloaded data. 98 | 99 | """ 100 | if dataset is None: 101 | selected_dataset = self._present_options() 102 | else: 103 | selected_dataset = dataset 104 | 105 | save_path_full = join(save_path, selected_dataset.split('.')[0]) 106 | 107 | if isdir(save_path_full): 108 | warn("\n'{0}' already exists. Voiding Download.".format( 109 | save_path_full)) 110 | else: 111 | self._print('Downloading Data...') 112 | url = "{0}/{1}".format(self.url, selected_dataset) 113 | self._download_data(url, save_path=save_path) 114 | 115 | return abspath(save_path_full) 116 | -------------------------------------------------------------------------------- /util/html.py: -------------------------------------------------------------------------------- 1 | import dominate 2 | from dominate.tags import * 3 | import os 4 | 5 | 6 | class HTML: 7 | def __init__(self, web_dir, title, reflesh=0): 8 | self.title = title 9 | self.web_dir = web_dir 10 | self.img_dir = os.path.join(self.web_dir, 'images') 11 | if not os.path.exists(self.web_dir): 12 | os.makedirs(self.web_dir) 13 | if not os.path.exists(self.img_dir): 14 | os.makedirs(self.img_dir) 15 | # print(self.img_dir) 16 | 17 | self.doc = dominate.document(title=title) 18 | if reflesh > 0: 19 | with self.doc.head: 20 | meta(http_equiv="reflesh", content=str(reflesh)) 21 | 22 | def get_image_dir(self): 23 | return self.img_dir 24 | 25 | def add_header(self, str): 26 | with self.doc: 27 | h3(str) 28 | 29 | def add_table(self, border=1): 30 | self.t = table(border=border, style="table-layout: fixed;") 31 | self.doc.add(self.t) 32 | 33 | def add_images(self, ims, txts, links, width=400): 34 | self.add_table() 35 | with self.t: 36 | with tr(): 37 | for im, txt, link in zip(ims, txts, links): 38 | with td(style="word-wrap: break-word;", halign="center", valign="top"): 39 | with p(): 40 | with a(href=os.path.join('images', link)): 41 | img(style="width:%dpx" % width, src=os.path.join('images', im)) 42 | br() 43 | p(txt) 44 | 45 | def save(self): 46 | html_file = '%s/index.html' % self.web_dir 47 | f = open(html_file, 'wt') 48 | f.write(self.doc.render()) 49 | f.close() 50 | 51 | 52 | if __name__ == '__main__': 53 | html = HTML('web/', 'test_html') 54 | html.add_header('hello world') 55 | 56 | ims = [] 57 | txts = [] 58 | links = [] 59 | for n in range(4): 60 | ims.append('image_%d.png' % n) 61 | txts.append('text_%d' % n) 62 | links.append('image_%d.png' % n) 63 | html.add_images(ims, txts, links) 64 | html.save() 65 | -------------------------------------------------------------------------------- /util/image_pool.py: -------------------------------------------------------------------------------- 1 | import random 2 | import torch 3 | 4 | 5 | class ImagePool(): 6 | def __init__(self, pool_size): 7 | self.pool_size = pool_size 8 | if self.pool_size > 0: 9 | self.num_imgs = 0 10 | self.images = [] 11 | 12 | def query(self, images): 13 | if self.pool_size == 0: 14 | return images 15 | return_images = [] 16 | for image in images: 17 | image = torch.unsqueeze(image.data, 0) 18 | if self.num_imgs < self.pool_size: 19 | self.num_imgs = self.num_imgs + 1 20 | self.images.append(image) 21 | return_images.append(image) 22 | else: 23 | p = random.uniform(0, 1) 24 | if p > 0.5: 25 | random_id = random.randint(0, self.pool_size - 1) # randint is inclusive 26 | tmp = self.images[random_id].clone() 27 | self.images[random_id] = image 28 | return_images.append(tmp) 29 | else: 30 | return_images.append(image) 31 | return_images = torch.cat(return_images, 0) 32 | return return_images 33 | -------------------------------------------------------------------------------- /util/util.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import torch 3 | import numpy as np 4 | from PIL import Image 5 | import os 6 | 7 | 8 | # Converts a Tensor into an image array (numpy) 9 | # |imtype|: the desired type of the converted numpy array 10 | def tensor2im(input_image, imtype=np.uint8): 11 | if isinstance(input_image, torch.Tensor): 12 | image_tensor = input_image.data 13 | else: 14 | return input_image 15 | image_numpy = image_tensor[0].cpu().float().numpy() 16 | if image_numpy.shape[0] == 1: 17 | image_numpy = np.tile(image_numpy, (3, 1, 1)) 18 | image_numpy = (np.transpose(image_numpy, (1, 2, 0)))* 255.0 19 | return image_numpy.astype(imtype) 20 | 21 | 22 | def tensor2labelim(label_tensor, impalette, imtype=np.uint8): 23 | if len(label_tensor.shape) == 4: 24 | _, label_tensor = torch.max(label_tensor.data.cpu(), 1) 25 | 26 | label_numpy = label_tensor[0].cpu().float().detach().numpy() 27 | label_image = Image.fromarray(label_numpy.astype(np.uint8)) 28 | label_image = label_image.convert("P") 29 | label_image.putpalette(impalette) 30 | label_image = label_image.convert("RGB") 31 | return np.array(label_image).astype(imtype) 32 | 33 | def diagnose_network(net, name='network'): 34 | mean = 0.0 35 | count = 0 36 | for param in net.parameters(): 37 | if param.grad is not None: 38 | mean += torch.mean(torch.abs(param.grad.data)) 39 | count += 1 40 | if count > 0: 41 | mean = mean / count 42 | print(name) 43 | print(mean) 44 | 45 | 46 | def save_image(image_numpy, image_path): 47 | image_pil = Image.fromarray(image_numpy) 48 | image_pil.save(image_path) 49 | 50 | 51 | def print_numpy(x, val=True, shp=False): 52 | x = x.astype(np.float64) 53 | if shp: 54 | print('shape,', x.shape) 55 | if val: 56 | x = x.flatten() 57 | print('mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f' % ( 58 | np.mean(x), np.min(x), np.max(x), np.median(x), np.std(x))) 59 | 60 | 61 | def mkdirs(paths): 62 | if isinstance(paths, list) and not isinstance(paths, str): 63 | for path in paths: 64 | mkdir(path) 65 | else: 66 | mkdir(paths) 67 | 68 | 69 | def mkdir(path): 70 | if not os.path.exists(path): 71 | os.makedirs(path) 72 | 73 | def confusion_matrix(x , y, n, ignore_label=None, mask=None): 74 | if mask is None: 75 | mask = np.ones_like(x) == 1 76 | k = (x >= 0) & (y < n) & (x != ignore_label) & (mask.astype(np.bool)) 77 | return np.bincount(n * x[k].astype(int) + y[k], minlength=n**2).reshape(n, n) 78 | 79 | def getScores(conf_matrix): 80 | if conf_matrix.sum() == 0: 81 | return 0, 0, 0 82 | with np.errstate(divide='ignore',invalid='ignore'): 83 | overall = np.diag(conf_matrix).sum() / np.float(conf_matrix.sum()) 84 | perclass = np.diag(conf_matrix) / conf_matrix.sum(1).astype(np.float) 85 | IU = np.diag(conf_matrix) / (conf_matrix.sum(1) + conf_matrix.sum(0) - np.diag(conf_matrix)).astype(np.float) 86 | return overall * 100., np.nanmean(perclass) * 100., np.nanmean(IU) * 100. 87 | -------------------------------------------------------------------------------- /util/visualizer.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import os 3 | import sys 4 | import ntpath 5 | import time 6 | from . import util 7 | from . import html 8 | from scipy.misc import imresize 9 | import pickle 10 | import cv2 11 | 12 | if sys.version_info[0] == 2: 13 | VisdomExceptionBase = Exception 14 | else: 15 | VisdomExceptionBase = ConnectionError 16 | 17 | 18 | # save image to the disk 19 | def save_images(webpage, visuals, image_path, aspect_ratio=1.0, width=224): 20 | image_dir = webpage.get_image_dir() 21 | short_path = ntpath.basename(image_path[0]) 22 | name = os.path.splitext(short_path)[0] 23 | palet_file = 'datasets/palette.txt' 24 | impalette = list(np.genfromtxt(palet_file, dtype=np.uint8).reshape(3*256)) 25 | 26 | webpage.add_header(name) 27 | ims, txts, links = [], [], [] 28 | 29 | for label, im_data in visuals.items(): 30 | if label == 'mask' or label == 'output': 31 | im = util.tensor2labelim(im_data, impalette) 32 | else: 33 | im = util.tensor2im(im_data) 34 | image_name = '%s_%s.png' % (name, label) 35 | save_path = os.path.join(image_dir, image_name) 36 | h, w, _ = im.shape 37 | if aspect_ratio > 1.0: 38 | im = imresize(im, (h, int(w * aspect_ratio)), interp='bicubic') 39 | if aspect_ratio < 1.0: 40 | im = imresize(im, (int(h / aspect_ratio), w), interp='bicubic') 41 | util.save_image(im, save_path) 42 | 43 | ims.append(image_name) 44 | txts.append(label) 45 | links.append(image_name) 46 | webpage.add_images(ims, txts, links, width=width) 47 | 48 | def save_scannet_prediction(mask, scan, frame, save_dir): 49 | 50 | # size = (1296,968) 51 | palet_file = 'datasets/palette.txt' 52 | # impalette = list(np.genfromtxt(palet_file, dtype=np.uint8).reshape(3*256)) 53 | # im = util.tensor2labelim(mask, impalette) 54 | # im = cv2.resize(im,size,interpolation=cv2.INTER_NEAREST) 55 | save_path = os.path.join(save_dir, scan + '_' + frame) 56 | util.save_image(mask.astype(np.uint8), save_path) 57 | 58 | 59 | class Visualizer(): 60 | def __init__(self, opt): 61 | self.display_id = opt.display_id 62 | self.use_html = opt.isTrain and not opt.no_html 63 | self.win_size = opt.display_winsize 64 | self.name = opt.name 65 | self.opt = opt 66 | palet_file = 'datasets/palette.txt' 67 | self.impalette = list(np.genfromtxt(palet_file,dtype=np.uint8).reshape(3*256)) 68 | self.saved = False 69 | if self.display_id > 0: 70 | import visdom 71 | self.ncols = opt.display_ncols 72 | self.vis = visdom.Visdom(server=opt.display_server, port=opt.display_port, env=opt.display_env, raise_exceptions=True) 73 | 74 | if self.use_html: 75 | self.web_dir = os.path.join(opt.checkpoints_dir, opt.name, 'web') 76 | self.img_dir = os.path.join(self.web_dir, 'images') 77 | print('create web directory %s...' % self.web_dir) 78 | util.mkdirs([self.web_dir, self.img_dir]) 79 | self.log_name = os.path.join(opt.checkpoints_dir, opt.name, 'loss_log.txt') 80 | self.val_log_name = os.path.join(opt.checkpoints_dir, opt.name, 'accuracy_log.txt') 81 | if opt.phase == "train": 82 | with open(self.log_name, "w") as log_file: 83 | now = time.strftime("%c") 84 | log_file.write('================ Training Loss (%s) ================\n' % now) 85 | with open(self.val_log_name, "w") as log_file: 86 | now = time.strftime("%c") 87 | log_file.write('================ Validation Accuracy (%s) ================\n' % now) 88 | self.conf_mat_name = os.path.join(opt.checkpoints_dir, opt.name, opt.phase + '_conf_mat.pkl') 89 | with open(self.conf_mat_name, "wb") as conf_mat_file: 90 | conf_mat = {} 91 | pickle.dump(conf_mat, conf_mat_file) 92 | 93 | def reset(self): 94 | self.saved = False 95 | 96 | def throw_visdom_connection_error(self): 97 | print('\n\nCould not connect to Visdom server (https://github.com/facebookresearch/visdom) for displaying training progress.\nYou can suppress connection to Visdom using the option --display_id -1. To install visdom, run \n$ pip install visdom\n, and start the server by \n$ python -m visdom.server.\n\n') 98 | exit(1) 99 | 100 | # |visuals|: dictionary of images to display or save 101 | def display_current_results(self, visuals, epoch, save_result): 102 | if self.display_id > 0: # show images in the browser 103 | ncols = self.ncols 104 | if ncols > 0: 105 | ncols = min(ncols, len(visuals)) 106 | h, w = next(iter(visuals.values())).shape[:2] 107 | table_css = """""" % (w, h) 111 | title = self.name 112 | label_html = '' 113 | label_html_row = '' 114 | images = [] 115 | idx = 0 116 | for label, image in visuals.items(): 117 | if label == 'mask' or label == 'output': 118 | image_numpy = util.tensor2labelim(image,self.impalette) 119 | else: 120 | image_numpy = util.tensor2im(image) 121 | label_html_row += '%s' % label 122 | images.append(image_numpy.transpose([2, 0, 1])) 123 | idx += 1 124 | if idx % ncols == 0: 125 | label_html += '%s' % label_html_row 126 | label_html_row = '' 127 | white_image = np.ones_like(image_numpy.transpose([2, 0, 1])) * 255 128 | while idx % ncols != 0: 129 | images.append(white_image) 130 | label_html_row += '' 131 | idx += 1 132 | if label_html_row != '': 133 | label_html += '%s' % label_html_row 134 | # pane col = image row 135 | try: 136 | self.vis.images(images, nrow=ncols, win=self.display_id + 1, 137 | padding=2, opts=dict(title=title + ' images')) 138 | label_html = '%s
' % label_html 139 | self.vis.text(table_css + label_html, win=self.display_id + 2, 140 | opts=dict(title=title + ' labels')) 141 | except VisdomExceptionBase: 142 | self.throw_visdom_connection_error() 143 | 144 | else: 145 | idx = 1 146 | for label, image in visuals.items(): 147 | if label == 'mask' or label == 'output': 148 | image_numpy = util.tensor2labelim(image,self.impalette) 149 | else: 150 | image_numpy = util.tensor2im(image) 151 | self.vis.image(image_numpy.transpose([2, 0, 1]), opts=dict(title=label), 152 | win=self.display_id + idx) 153 | idx += 1 154 | 155 | if self.use_html and (save_result or not self.saved): # save images to a html file 156 | self.saved = True 157 | for label, image in visuals.items(): 158 | if label == 'mask' or label == 'output': 159 | image_numpy = util.tensor2labelim(image,self.impalette) 160 | else : 161 | image_numpy = util.tensor2im(image) 162 | img_path = os.path.join(self.img_dir, 'epoch%.3d_%s.png' % (epoch, label)) 163 | util.save_image(image_numpy, img_path) 164 | # update website 165 | webpage = html.HTML(self.web_dir, 'Experiment name = %s' % self.name, reflesh=1) 166 | for n in range(epoch, 0, -1): 167 | webpage.add_header('epoch [%d]' % n) 168 | ims, txts, links = [], [], [] 169 | 170 | for label, image in visuals.items(): 171 | if label == 'mask' or label == 'output': 172 | image_numpy = util.tensor2labelim(image,self.impalette) 173 | else: 174 | image_numpy = util.tensor2im(image) 175 | img_path = 'epoch%.3d_%s.png' % (n, label) 176 | ims.append(img_path) 177 | txts.append(label) 178 | links.append(img_path) 179 | webpage.add_images(ims, txts, links, width=self.win_size) 180 | webpage.save() 181 | 182 | # losses: dictionary of error labels and values 183 | def plot_current_losses(self, epoch, counter_ratio, opt, losses): 184 | if not hasattr(self, 'plot_data'): 185 | self.plot_data = {'X': [], 'Y': [], 'legend': list(losses.keys())} 186 | self.plot_data['X'].append(epoch + counter_ratio) 187 | self.plot_data['Y'].append([losses[k] for k in self.plot_data['legend']]) 188 | try: 189 | self.vis.line( 190 | X=np.stack([np.array(self.plot_data['X'])] * len(self.plot_data['legend']), 1), 191 | Y=np.array(self.plot_data['Y']), 192 | opts={ 193 | 'title': self.name + ' loss over time', 194 | 'legend': self.plot_data['legend'], 195 | 'xlabel': 'epoch', 196 | 'ylabel': 'loss'}, 197 | win=self.display_id) 198 | except VisdomExceptionBase: 199 | self.throw_visdom_connection_error() 200 | 201 | # losses: same format as |losses| of plot_current_losses 202 | def print_current_losses(self, epoch, i, losses, t, t_data): 203 | message = '(epoch: %d, iters: %d, time: %.3f, data: %.3f) ' % (epoch, i, t, t_data) 204 | for k, v in losses.items(): 205 | message += '%s: %.3f ' % (k, v) 206 | 207 | print(message) 208 | with open(self.log_name, "a") as log_file: 209 | log_file.write('%s\n' % message) 210 | 211 | def print_current_scores(self, epoch, loss, glob, mean, iou): 212 | message = '(epoch {0:} test loss: {1:.3f} '.format(epoch, loss) 213 | message += 'glob acc : {0:.2f}, mean acc : {1:.2f}, IoU : {2:.2f}'.format(glob, mean, iou) 214 | 215 | print(message) 216 | with open(self.val_log_name, "a") as log_file: 217 | log_file.write('%s\n' % message) 218 | 219 | def save_confusion_matrix(self, conf_mat, epoch): 220 | conf_mats = pickle.load(open(self.conf_mat_name, "rb")) 221 | conf_mats["epoch"+str(epoch)] = conf_mat 222 | conf_mats["epoch"+str(epoch)+"_scores"] = np.asarray(util.getScores(conf_mat)) 223 | pickle.dump(conf_mats, open(self.conf_mat_name, "wb")) 224 | --------------------------------------------------------------------------------