fixup legacy codebase

Author: leVirve

.gitignore

@@ -11,3 +11,5 @@ logs*/
 *.pyc
 *.log
 *.avi
+
+data/

README.md

@@ -13,84 +13,59 @@
 }
 ```
 
-**The code is under evaluation and update TBD. Deprecated information below.**
 Thanks [@shuuchen](https://github.com/shuuchen) for an all-in-one project, you may also refer to https://github.com/shuuchen/lsun-room-dsc!
 
 ## Prerequisite
 
 - Python 3.6+
+- PyTorch 1.0+
 - [OneGAN](https://github.com/leVirve/OneGAN) == `0.3.2`, clone and checkout to that tag.
 
   ```bash
   git clone https://github.com/leVirve/OneGA
   git checkout 0.3.2
   ```
 
+- `pip install -e requirements.txt`
+
 ## Dataset
 
-- Download dataset from http://lsun.cs.princeton.edu/2015.html#layout and put them in the following folders.
-- Unfortunately, the website hosted LSUN Layout challenge is down, you can only find them from the [web archive](https://web.archive.org/web/20190118150204/http://lsun.cs.princeton.edu/2016/), e.g. [layout.zip](https://web.archive.org/web/20170221111502/http://lsun.cs.princeton.edu/challenge/2015/roomlayout/data/layout.zip), [training.mat](https://web.archive.org/web/20180923231343/http://lsun.cs.princeton.edu/challenge/2015/roomlayout/data/training.mat), [validation.mat](https://web.archive.org/web/20180923231343/http://lsun.cs.princeton.edu/challenge/2015/roomlayout/data/validation.mat), [testing.mat](https://web.archive.org/web/20180923231343/http://lsun.cs.princeton.edu/challenge/2015/roomlayout/data/testing.mat).
-- However, the `image.zip` is too large and still unavailable to download from the web archive, fortunately, you can find them in https://github.com/liamw96/pytorch.room.layout#data. [@liamw96](https://github.com/liamw96) provides dataset image in the `lsun.zip` on Google drive!
+- Download dataset from http://lsun.cs.princeton.edu/2015.html#layout.
+  - Unfortunately, the website hosted LSUN Layout challenge is down.
+- Find them from the [web archive](https://web.archive.org/web/20190118150204/http://lsun.cs.princeton.edu/2016/)
+  - [layout.zip](https://web.archive.org/web/20170221111502/http://lsun.cs.princeton.edu/challenge/2015/roomlayout/data/layout.zip)
+  - [training.mat](https://web.archive.org/web/20180923231343/http://lsun.cs.princeton.edu/challenge/2015/roomlayout/data/training.mat), [validation.mat](https://web.archive.org/web/20180923231343/http://lsun.cs.princeton.edu/challenge/2015/roomlayout/data/validation.mat) and [testing.mat](https://web.archive.org/web/20180923231343/http://lsun.cs.princeton.edu/challenge/2015/roomlayout/data/testing.mat).
+- However, the `image.zip` is too large and still unavailable to download from the web archive.
+  - Fortunately, you can find them in https://github.com/liamw96/pytorch.room.layout#data. [@liamw96](https://github.com/liamw96) provides dataset image in the `lsun.zip` on Google drive!
 
 ## Usage
 
 - Dataset
 
-  - Edit the root of path at `config.yml` Line#8 `folder`.
-  - Put `LSUN Room Layout Dataset` into folder `${folder}` relative to this project.
+  - Put `LSUN Room Layout Dataset` into folder `./data/lsun_room`.
     - `images/`: RGB color image `*.jpg` of indoor room scene
     - `layout_seg/`: layout ground truth `*.mat` of indoor room scene
-    - `layout_seg_images/`: generated layout ground truth `*.png` of indoor room scene
   - Prepare layout images for trianing/validation.
+    - `layout_seg_images/`: generated layout ground truth `*.png` of indoor room scene
 
     ```bash
-    python script/re_label.py
+    python -m script.re_label
     ```
 
 - Training
 
-  - The trained model will be saved to folder `./exp/checkpoints/`
-  - You can modify `config.yml` to play with hyperparameters for training.
+  - The trained model will be saved to folder `./ckpts`
 
   Example
 
   ```bash
   python main.py --phase train --arch resnet --edge_factor 0.2 --l2_factor 0.2 --name baseline
   ```
 
-  Detials
-
-  ```bash
-  python main.py
-
-  Usage: main.py [OPTIONS]
-
-  Options:
-    --name TEXT
-    --dataset [lsun_room | others]
-    --dataset_root TEXT
-    --log_dir TEXT
-    --image_size <INTEGER INTEGER>
-    --epochs INTEGER
-    --batch_size INTEGER
-    --workers INTEGER
-    --l1_weight FLOAT
-    --resume PATH
-  ```
-
 - Demo
 
   ```bash
-  python demo.py
-
-  Usage: demo.py [OPTIONS]
-
-  Options:
-    --device INTEGER
-    --video TEXT
-    --weight TEXT
-    --input_size <INTEGER INTEGER>.
-
+  python demo.py --weight {checkpoint_path} --video {test_video}
   ```
 
 - Toolkit

config.yml

@@ -8,9 +8,9 @@
 import torchvision.transforms as T
 from PIL import Image
 
-from lsun_room.label import Layout
-from lsun_room.loader import LsunRoomDataset as _BaseDataset
-from lsun_room.loader import get_meta
+from lib.lsun_room_api.lsun_room.label import Layout
+from lib.lsun_room_api.lsun_room.loader import LsunRoomDataset as _BaseDataset
+from lib.lsun_room_api.lsun_room.loader import get_meta
 
 
 class LsunRoomDataset(_BaseDataset):
@@ -31,7 +31,7 @@ def __init__(self, phase, args, **kwarges):
     def collect_meta(self, root, phase):
         ''' metadata fold1: original dataset '''
         meta = [
-            {'image_path': root / self.image_folder / f'{e["name"]}.jpg',
+            {'image_path': root / self.image_folder / f'{e["name"]}.png',
              'layout_path': root / self.layout_folder / f'{e["name"]}.png',
              'name': e['name'] + '.png', 'type': e['type']}
             for e in get_meta(dataset_root=root, phase=phase)

datasets/lsunroom.py

@@ -5,15 +5,16 @@
 import torchvision.transforms as T
 from PIL import Image
 
-from trainer.賣扣老師 import build_resnet101_FCN
+from trainer import core
 
 torch.backends.cudnn.benchmark = True
 
 
 class Predictor:
 
-    def __init__(self, input_size, weight=None):
-        self.model = self.build_model(weight)
+    def __init__(self, input_size, weight_path):
+        self.model = core.LayoutSeg.load_from_checkpoint(weight_path)
+        self.model.freeze()
         self.colorizer = onegan.extension.Colorizer(
             colors=[
                 [249, 69, 93], [255, 229, 170], [144, 206, 181],
@@ -24,22 +25,14 @@ def __init__(self, input_size, weight=None):
             T.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5])
         ])
 
-    def build_model(self, weight_path, joint_class=False):
-        model = build_resnet101_FCN(pretrained=False, nb_classes=37, stage_2=True, joint_class=joint_class)
-        weight = onegan.utils.export_checkpoint_weight(weight_path)
-        model.load_state_dict(weight)
-        model.eval()
-        return model.cuda()
-
     @onegan.utils.timeit
     def process(self, raw):
 
         def _batched_process(batched_img):
-            score, _ = self.model(onegan.utils.to_var(batched_img))
-            _, output = torch.max(score, 1)
+            _, outputs = self.model(batched_img)
 
             image = (batched_img / 2 + .5)
-            layout = self.colorizer.apply(output.data.cpu())
+            layout = self.colorizer.apply(outputs.data.cpu())
             return image * .6 + layout * .4
 
         img = Image.fromarray(raw)
@@ -55,8 +48,7 @@ def _batched_process(batched_img):
 @click.option('--weight', type=click.Path(exists=True))
 @click.option('--input_size', default=(320, 320), type=(int, int))
 def main(device, video, weight, input_size):
-
-    demo = Predictor(input_size, weight=weight)
+    demo = Predictor(input_size, weight_path=weight)
 
     reader = video if video else device
     cap = cv2.VideoCapture(reader)

demo.py

@@ -1 +1 @@
-__VERSION__ = '0.1.0'
+__VERSION__ = '0.1.1'

lib/lsun_room_api/lsun_room/__init__.py

@@ -4,7 +4,7 @@
 import scipy.io as sio
 
 import cv2
-from lsun_room.label import mapping_func
+from .label import mapping_func
 
 
 def load_mat(path):
@@ -47,16 +47,15 @@ def layout(self):
 
     def remap_layout(self, verbose=False):
         mapping = mapping_func(self.type)(self)
-
         old_layout = load_mat(self.layout_mat_path)
         layout = np.zeros_like(old_layout)
         for new_label, point in mapping:
             old_label = old_layout[point[1], point[0]]
-            layout[old_layout == old_label] = new_label
+            layout[old_layout == old_label] = new_label.value
         return layout
 
     def save_layout(self, visualization=False):
-        save_image(self.layout_path, self.layout)
+        save_image(self.layout_path, self.remap_layout())
 
     def __str__(self):
         return '<DataItem: %s>' % self.name

lib/lsun_room_api/lsun_room/item.py

@@ -5,7 +5,7 @@
 from onegan.io.loader import load_image, BaseDataset
 from onegan.io.transform import SegmentationPair
 
-from lsun_room.edge import gen_edge_map, gen_corner_map
+from .edge import gen_edge_map, gen_corner_map
 
 
 def get_meta(dataset_root, phase):

lib/lsun_room_api/lsun_room/loader.py

@@ -1,11 +1,11 @@
+import argparse
 import importlib
-import logging
+from pprint import pprint
 
-import torch
 import onegan
+import pytorch_lightning as pl
 
 from trainer import core
-from trainer.model import ResPlanarSeg
 
 
 def create_dataset(args):
@@ -24,19 +24,6 @@ def create_dataset(args):
             for phase in ['train', 'val'])
 
 
-def create_model(args):
-    return {
-        'resnet': lambda: ResPlanarSeg(num_classes=args.num_class, pretrained=True, base='resnet101')
-    }[args.arch]()
-
-
-def create_optim(args, model, optim='sgd'):
-    return {
-        'adam': lambda: torch.optim.Adam(model.parameters(), lr=args.lr, betas=(0.5, 0.999)),
-        'sgd': lambda: torch.optim.SGD(model.parameters(), lr=args.lr / 10, momentum=0.9)
-    }[optim]()
-
-
 def hyperparams_search(args):
     search_hyperparams = {
         'arch': ['vgg', 'mike'],
@@ -54,33 +41,35 @@ def hyperparams_search(args):
 
 
 def main(args):
-    log = logging.getLogger('room')
-    log.info(''.join([f'\n-- {k}: {v}' for k, v in args.items()]))
+    pprint(args)
 
     train_loader, val_loader = create_dataset(args)
-    model = create_model(args)
-
     if args.phase == 'train':
-        training_estimator = core.training_estimator(
-            torch.nn.DataParallel(model.cuda()),
-            create_optim(args, model, optim=args.optim), args)
-        training_estimator(train_loader, val_loader, epochs=args.epoch)
-
-    if args.phase in ['eval', 'eval_search']:
-        core_fn = core.evaluation_estimator if args.phase == 'eval' else core.weights_estimator
-        evaluate_estimator = core_fn(torch.nn.DataParallel(model.cuda()), args)
-        evaluate_estimator(val_loader)
+        model = core.LayoutSeg(args.num_class, args.lr, args)
+        trainer = pl.Trainer(
+            gpus=1,
+            limit_train_batches=.1,
+            limit_val_batches=.1,
+            max_epochs=args.epoch,
+            weights_save_path=f'ckpts/{args.name}',
+            logger=pl.loggers.TensorBoardLogger('ckpts/tb_logs', name=args.name)
+        )
+        trainer.fit(model, train_loader, val_loader)
 
 
 if __name__ == '__main__':
-    parser = onegan.option.Parser(description='Indoor room corner detection', config='./config.yml')
+    parser = argparse.ArgumentParser(description='Indoor room corner detection')
     parser.add_argument('--name', help='experiment name')
-    parser.add_argument('--folder', help='where\'s the dataset')
+    parser.add_argument('--folder', default='data/lsun_room', help='where is the dataset')
     parser.add_argument('--dataset', default='lsunroom', choices=['lsunroom', 'hedau', 'sunrgbd'])
     parser.add_argument('--phase', default='eval', choices=['train', 'eval', 'eval_search'])
+    parser.add_argument('--worker', default=4, type=int)
+    parser.add_argument('--epoch', default=30, type=int)
+    parser.add_argument('--batch_size', default=8, type=int)
+    parser.add_argument('--lr', default=0.0001, type=float)
 
     # data
-    parser.add_argument('--image_size', type=int)
+    parser.add_argument('--image_size', default=320, type=int)
     parser.add_argument('--use_edge', action='store_true')
     parser.add_argument('--use_corner', action='store_true')
     parser.add_argument('--datafold', type=int, default=1)
@@ -98,11 +87,8 @@ def main(args):
     parser.add_argument('--l1_factor', type=float, default=0.0)
     parser.add_argument('--l2_factor', type=float, default=0.0)
     parser.add_argument('--edge_factor', type=float, default=0.0)
+    parser.add_argument('--type_factor', type=float, default=0.0)
     parser.add_argument('--focal_gamma', type=float, default=0)
-    args = parser.parse()
+    args = parser.parse_args()
 
-    logging.basicConfig(level=logging.INFO,
-                        format='%(asctime)s %(name)-12s %(levelname)-8s %(message)s',
-                        datefmt='%m-%d %H:%M',
-                        handlers=[logging.StreamHandler(), ])
     main(args)

main.py

@@ -4,22 +4,21 @@
 
 import click
 
-from datasets.lsun_room.item import DataItems
+from lib.lsun_room_api.lsun_room.item import DataItems
 
 
 def worker(item):
-    item.remap_layout()
     item.save_layout()
 
 
 @click.command()
-@click.option('--dataset_root', default='../data/lsun_room/')
+@click.option('--dataset_root', default='data/lsun_room/')
 def main(dataset_root):
 
     for phase in ['train', 'val']:
         print('==> re-label for data in %s phase' % phase)
         s = time.time()
-        dataset = DataItems(root_dir=dataset_root, phase=phase)
+        dataset = DataItems(root=dataset_root, phase=phase)
         with Pool(cpu_count()) as pool:
             pool.map(worker, dataset.items)
         print('==> Done in %.4f sec.' % (time.time() - s))