feat(utils): saving to_multi_channel before del

data/gh/cam_rh.3dm.rhl

@@ -0,0 +1,3 @@
+Administrator
+LAPTOP-T062M1FB
+Sunday, April 25, 2021

deeprad/image2vec.py

@@ -17,7 +17,7 @@
 
 # deeprad modules
 from .utils import to_poly_sh, to_poly_np, extract_floorplan_ids, load_floorplan_data, \
-    load_json, make_dir_safely, DEEPRAD_MODELS_DIR
+    load_json, make_dir_safely, DEEPRAD_MODELS_DIR, DEEPRAD_GHOUT_DIR
 
 
 def viz_poly_sh_arr(poly_sh_arr, a=None, scale=1, iter=False, **kwargs):
@@ -404,7 +404,7 @@ def main(data_num):
     hdicts, _, labels, targ_id_dirs = load_floorplan_data(targ_ids, data_num)
 
     print('\nWriting {} polygon json files in {}.\n'.format(
-          data_num, DEEPRAD_DATA_DIR))
+          data_num, DEEPRAD_MODELS_DIR))
 
     for ii, (hdict, _rooms, targ_id_dir) in enumerate(zip(hdicts, labels, targ_id_dirs)):
 
@@ -571,7 +571,7 @@ def write_json(polygon_dict, dest_fpath):
         # Make image folders for train_test
         # -------------------------------------------------------------
         train_test_id_dir = os.path.join(
-            DEEPRAD_DATA_DIR, '..', 'train_test', model_id)
+            DEEPRAD_GHOUT_DIR, model_id)
         train_test_id_dir = os.path.abspath(train_test_id_dir)
         make_dir_safely(train_test_id_dir)
 

deeprad/make_traintest.py

@@ -1,6 +1,7 @@
 """Preprocess raw outputs from GHOUT_DIR for train/test dataset."""
 import os
 import numpy as np
+from typing import List
 
 import cv2
 
@@ -11,37 +12,84 @@
 DEEPRAD_TRAINTEST_DIR = utils.DEEPRAD_TRAINTEST_DIR
 
 
-def preprocess_img(img_fpath: str) -> np.ndarray:
+def preprocess_img(img_lst: List[np.ndarray], downsample: int = 1) -> np.ndarray:
     """Preprocess image for traintest."""
 
-    print(img_fpath)
+    ret_img = None
+    # TODO: contiguous ones and trim.
+    # concatenate all images acrosss x-axis
+
+    ret_img = np.concatenate(img_lst, axis=1)
+
+    # check dim
+    xdim = img_lst[0].shape[1]
+    assert ret_img.shape[1] == (xdim * len(img_lst))
+
+    # downsample
+    if downsample > 1:
+        ret_img = ret_img[::downsample, ::downsample]
+
+    return ret_img
+
+
+def in_channel_set(in_fpaths: List[str]) -> dict:
+    """Creates dictionary of unqiue channels from input directory."""
+
+    chset = set()
+    for fpath in in_fpaths:
+        fname = utils.fname_from_fpath(fpath)
+        chtype = fname.split('_')[0]  # remove orientation info
+        if chtype not in chset:
+            chset.add(chtype)
+
+    return chset
 
 
 def main():
     """Preproces main."""
 
     # load all directories in train/test
     data_num, floorplan_ids = \
-        utils.extract_floorplan_ids(1e6, DEEPRAD_TRAINTEST_DIR, verbose=False)
+        utils.extract_floorplan_ids(1e6, DEEPRAD_GHOUT_DIR, verbose=False)
+
+    chset = set()
 
     # for i, floorplan_id in enumerate(floorplan_ids):
     i = 0
     floorplan_id = floorplan_ids[0]
 
+    # Make filepaths and directores
     in_dir = os.path.join(DEEPRAD_GHOUT_DIR, floorplan_id, 'in_data')
     out_dir = os.path.join(DEEPRAD_GHOUT_DIR, floorplan_id, 'out_label')
-
-    # Preprocess labels
-    label_fpaths = [os.path.join(out_dir, out) for out in os.listdir(out_dir)]
-    label_imgs = [preprocess_img(fpath) for fpath in label_fpaths]
-    #[save_img(img) for img in label_imgs]
-
-    # Preprocess input images
-    # add as channels
-
-    # save new images
-
-    pass
+    in_img_fpath = os.path.join(
+        DEEPRAD_TRAINTEST_DIR, "in_{}.jpg".format(floorplan_id))
+    out_img_fpath = os.path.join(
+        DEEPRAD_TRAINTEST_DIR, "out_{}.jpg".format(floorplan_id))
+
+    # # Preprocess traintest_out
+    # label_fpaths = [os.path.join(out_dir, _out)
+    #                 for _out in os.listdir(out_dir)]
+    # label_imgs = [utils.load_img_gray(f) for f in label_fpaths]
+    # out_img = preprocess_img(label_imgs)
+    # assert utils.write_img(out_img, out_img_fpath)
+
+    # Preprocess traintest_in
+    in_fpaths = [os.path.join(in_dir, _in) for _in in os.listdir(in_dir)]
+    in_fpaths = [os.path.join(in_dir, _in) for _in in os.listdir(in_dir)]
+    if len(chset) == 0:
+        chset = in_channel_set(in_fpaths)
+
+    # Concat multiple images into one image depthwise (as chennels)
+    n_channels = len(chset)  # number of channels
+    ch_imgs = [0] * n_channels
+    for j, ch in enumerate(chset):
+        ch_fpaths = [in_fpath for in_fpath in in_fpaths if ch in in_fpath]
+        _ch_imgs = [utils.load_img_gray(ch_fpath) for ch_fpath in ch_fpaths]
+        ch_imgs[j] = preprocess_img(_ch_imgs)  # concats across xdim
+    print(floorplan_id)
+    # print(ch_imgs)
+    in_img = utils.to_multi_channel_img(ch_imgs)
+    assert utils.write_img(in_img, in_img_fpath)
 
 
 if __name__ == "__main__":

deeprad/utils.py

@@ -5,6 +5,9 @@
 from pathlib import Path
 import matplotlib.pyplot as plt
 import numpy as np
+import cv2
+from typing import List
+import shapely.geometry as geom
 
 # Path to all models in deep_rad
 DEEPRAD_MODELS_DIR = os.path.abspath(os.path.join(
@@ -45,6 +48,10 @@ def load_img_rgb(img_fpath: str) -> np.ndarray:
     return cv2.imread(img_fpath, cv2.COLOR_BGR2RGB)
 
 
+def write_img(img: np.ndarray, img_fpath: str) -> bool:
+    return cv2.imwrite(img_fpath, img)
+
+
 def extract_floorplan_ids(data_num, target_data_dir=None, verbose=True):
     """Safely extract root model directories for polygon extraction."""
 
@@ -105,18 +112,41 @@ def load_floorplan_data(targ_id_dirs, data_num):
 
         idx += 1
         hdict_arr[idx] = hdict
-        src_img_arr[idx] = koad_img_rgb(targ_src_fpath)
+        src_img_arr[idx] = load_img_rgb(targ_src_fpath)
         label_img_arr[idx] = load_img_gray(targ_label_fpath)
         targ_id_dir_arr[idx] = targ_id_dir
         null_lst.append(targ_id_dirs[i] + '\n')
 
     # Write to null list
     null_fpath = os.path.join(DEEPRAD_MODELS_DIR, '_null.txt')
     with open(null_fpath, 'w') as fp:
-        #[fp.writeline(null_) for null_ in null_lst]
+        # [fp.writeline(null_) for null_ in null_lst]
         fp.writelines(null_lst)
     return hdict_arr, src_img_arr, label_img_arr, targ_id_dir_arr
 
 
 def make_dir_safely(dest_dir):
     Path(dest_dir).mkdir(parents=True, exist_ok=True)
+
+
+def fname_from_fpath(fpath):
+    """Splits filename from extension and preceding directories."""
+    return Path(fpath).stem
+
+
+def to_multi_channel_img(imgs: List[np.ndarray]):
+    """Safely constructs a multichannel image from multiple grayscale images."""
+
+    n_channels = len(imgs)
+
+    assert np.all([np.array_equal([len(img.shape)], [2])
+                  for img in imgs]), 'img must be 2d grayscale.'
+
+    assert np.all([np.array_equal(img.shape, imgs[0].shape) for img in imgs]), \
+        'All images must be same shape. Got {}.'.format(
+            [img.shape for img in imgs])
+
+    cat_img = np.dstack(imgs)  # Concat depthwise channels
+    assert cat_img.shape[2] == n_channels
+
+    return cat_img