├── .gitignore ├── .gitmodules ├── README.md ├── environment.yml ├── fabricflownet ├── eval.py ├── flownet │ ├── config.yaml │ ├── dataset.py │ ├── models.py │ └── train.py ├── picknet │ ├── config.yaml │ ├── dataset.py │ ├── models.py │ └── train.py └── utils.py ├── prepare_1.0.sh └── setup.py /.gitignore: -------------------------------------------------------------------------------- 1 | _bak/ 2 | outputs/ 3 | data/ 4 | 5 | # Byte-compiled / optimized / DLL files 6 | __pycache__/ 7 | *.py[cod] 8 | *$py.class 9 | 10 | # C extensions 11 | *.so 12 | 13 | # Distribution / packaging 14 | .Python 15 | build/ 16 | develop-eggs/ 17 | dist/ 18 | downloads/ 19 | eggs/ 20 | .eggs/ 21 | lib/ 22 | lib64/ 23 | parts/ 24 | sdist/ 25 | var/ 26 | wheels/ 27 | share/python-wheels/ 28 | *.egg-info/ 29 | .installed.cfg 30 | *.egg 31 | MANIFEST 32 | 33 | # PyInstaller 34 | # Usually these files are written by a python script from a template 35 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 36 | *.manifest 37 | *.spec 38 | 39 | # Installer logs 40 | pip-log.txt 41 | pip-delete-this-directory.txt 42 | 43 | # Unit test / coverage reports 44 | htmlcov/ 45 | .tox/ 46 | .nox/ 47 | .coverage 48 | .coverage.* 49 | .cache 50 | nosetests.xml 51 | coverage.xml 52 | *.cover 53 | *.py,cover 54 | .hypothesis/ 55 | .pytest_cache/ 56 | cover/ 57 | 58 | # Translations 59 | *.mo 60 | *.pot 61 | 62 | # Django stuff: 63 | *.log 64 | local_settings.py 65 | db.sqlite3 66 | db.sqlite3-journal 67 | 68 | # Flask stuff: 69 | instance/ 70 | .webassets-cache 71 | 72 | # Scrapy stuff: 73 | .scrapy 74 | 75 | # Sphinx documentation 76 | docs/_build/ 77 | 78 | # PyBuilder 79 | .pybuilder/ 80 | target/ 81 | 82 | # Jupyter Notebook 83 | .ipynb_checkpoints 84 | 85 | # IPython 86 | profile_default/ 87 | ipython_config.py 88 | 89 | # pyenv 90 | # For a library or package, you might want to ignore these files since the code is 91 | # intended to run in multiple environments; otherwise, check them in: 92 | # .python-version 93 | 94 | # pipenv 95 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 96 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 97 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 98 | # install all needed dependencies. 99 | #Pipfile.lock 100 | 101 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 102 | __pypackages__/ 103 | 104 | # Celery stuff 105 | celerybeat-schedule 106 | celerybeat.pid 107 | 108 | # SageMath parsed files 109 | *.sage.py 110 | 111 | # Environments 112 | .env 113 | .venv 114 | env/ 115 | venv/ 116 | ENV/ 117 | env.bak/ 118 | venv.bak/ 119 | 120 | # Spyder project settings 121 | .spyderproject 122 | .spyproject 123 | 124 | # Rope project settings 125 | .ropeproject 126 | 127 | # mkdocs documentation 128 | /site 129 | 130 | # mypy 131 | .mypy_cache/ 132 | .dmypy.json 133 | dmypy.json 134 | 135 | # Pyre type checker 136 | .pyre/ 137 | 138 | # pytype static type analyzer 139 | .pytype/ 140 | 141 | # Cython debug symbols 142 | cython_debug/ 143 | -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "softgym"] 2 | path = softgym 3 | url = git@github.com:Xingyu-Lin/softgym.git 4 | branch = fabricflownet 5 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # FabricFlowNet: Bimanual Cloth Manipulation with a Flow-based Policy 2 | Thomas Weng, Sujay Bajracharya, Yufei Wang, Khush Agrawal, David Held 3 | Conference on Robot Learning 2021 4 | 5 | [[arXiv](https://arxiv.org/abs/2111.05623)] [[pdf](https://arxiv.org/pdf/2111.05623.pdf)] [[project page](https://sites.google.com/view/fabricflownet)] 6 | 7 | This repo contains the code for running the FabricFlowNet model used to report the simulation results in the paper. 8 | 9 | The ROS code for running FabricFlowNet on a real-world system will be released in a separate repository. 10 | 11 | ## File Structure 12 | ```angular2html 13 | ├── FabricFlowNet/ 14 | | |── data/ # Folder for datasets, saved runs, evaluation goals 15 | | └── fabricflownet/ 16 | | |── flownet/ # FlowNet model and training code 17 | | |── picknet/ # PickNet model and training code 18 | | |── eval.py # Evaluation script 19 | | └── utils.py 20 | └── softgym/ # SoftGym submodule 21 | ``` 22 | 23 | ## Installation 24 | 25 | These instructions have been tested on Ubuntu 18.04 with NVIDIA GTX 3080/3090 GPUs. 26 | SoftGym requires CUDA 9.2+, FFN training and inference have been tested on CUDA 11.1. 27 | 28 | * Clone this repo 29 | * Set up the [softgym](https://github.com/Xingyu-Lin/softgym) submodule, tracking the `fabricflownet` branch: `git submodule update --init` 30 | * Follow the softgym instructions to create a softygm conda environment. Then compile PyFlex: `. prepare_1.0.sh` and `. compile_1.0.sh`. Check the compile script to make sure that the `CUDA_BIN_PATH` env variable is set to the path of the CUDA library you installed SoftGym with. 31 | * In the FabricFlowNet directory, activate the conda environment and set environment variables: `. prepare_1.0.sh` 32 | * `conda env update -f environment.yml --prune`. Ensure that you install a PyTorch version that is suitable for your CUDA version. 33 | * Install the repo as a python package: `pip install -e .` 34 | 35 | # Evaluation 36 | * Download the evaluation set and model weights into `./data/`: 37 | * Download and extract the [evaluation set](https://drive.google.com/file/d/1A9GUPXuVIC1K-LCCvzrK95-m9_UVSbPd/view?usp=sharing) 38 | * Download the [FlowNet weights](https://drive.google.com/file/d/1P7Upskczb-iqOsPjgcjsd4cnQQEuf-uY/view?usp=sharing), this does not need to be extracted 39 | * Download and extract the [PickNet weights](https://drive.google.com/file/d/1dCuSpMyvzkPU3AL7MeXeL7knP5ngyKvq/view?usp=sharing) 40 | * Run the evaluation script: `python fabricflownet/eval.py --run_path=data/picknet_run --ckpt=105000 --cloth_type=square_towel` 41 | * To run in headless mode, add the `--headless` flag; use the `-h` flag to see other available flags. 42 | * Performance on square towel goals from the paper (in mm): 43 | ``` 44 | Square Towel (mm) 45 | all: 6.803 +/- 10.413 46 | one-step: 4.262 +/- 2.287 47 | mul-step: 23.741 +/- 21.597 48 | ``` 49 | ``` 50 | Rectangular Towel (mm) 51 | all: 9.270 +/- 7.001 52 | one-step: 4.254 +/- 1.086 53 | mul-step: 16.793 +/- 5.143 54 | ``` 55 | ``` 56 | Tshirt (mm) 57 | all: 31.628 +/- 10.892 58 | one-step: 24.459 +/- 4.876 59 | mul-step: 45.965 +/- 0.000 60 | ``` 61 | 62 | ## Citation 63 | If you find this code useful in your research, please feel free to cite: 64 | ``` 65 | @inproceedings{weng2021fabricflownet, 66 | title={FabricFlowNet: Bimanual Cloth Manipulation with a Flow-based Policy}, 67 | author={Weng, Thomas and Bajracharya, Sujay and Wang, Yufei and Agrawal, Khush and Held, David}, 68 | booktitle={Conference on Robot Learning}, 69 | year={2021} 70 | } 71 | ``` 72 | -------------------------------------------------------------------------------- /environment.yml: -------------------------------------------------------------------------------- 1 | name: softgym_ffn_release 2 | channels: 3 | - defaults 4 | dependencies: 5 | - numpy=1.17.2 6 | - imageio=2.6.1 7 | - glob2=0.7 8 | - cmake=3.14.0 9 | - pybind11=2.4.3 10 | - click 11 | - matplotlib 12 | - joblib 13 | - Pillow=6.1 14 | - plotly=2.0 15 | - pip 16 | - pip: 17 | - gtimer 18 | - gym==0.14.0 19 | - moviepy 20 | - opencv-python==4.1.1.26 21 | - Shapely==1.6.4.post2 22 | - pyquaternion==0.9.5 23 | - sk-video==1.1.10 24 | 25 | # FFN specific 26 | # - torch==1.9.0 27 | # - torchvision==0.10.0 28 | - scipy==1.5.4 29 | - pytorch-lightning==1.3.0 30 | - torchmetrics==0.7.0 31 | - hydra-core==1.1.1 32 | - matplotlib 33 | - omegaconf==2.1.1 34 | -------------------------------------------------------------------------------- /fabricflownet/eval.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | import argparse 4 | import numpy as np 5 | import matplotlib.pyplot as plt 6 | import cv2 7 | from omegaconf import OmegaConf 8 | 9 | import torch 10 | from torch.utils.data import DataLoader 11 | 12 | import pyflex 13 | from softgym.envs.bimanual_env import BimanualEnv 14 | from softgym.envs.bimanual_tshirt import BimanualTshirtEnv 15 | 16 | from fabricflownet.flownet.models import FlowNet 17 | from fabricflownet.picknet.models import FlowPickSplitModel 18 | from fabricflownet.picknet.dataset import Goals 19 | 20 | class EnvRollout(object): 21 | def __init__(self, args): 22 | self.args = args 23 | 24 | np.random.seed(args.seed) 25 | torch.manual_seed(args.seed) 26 | torch.cuda.manual_seed(args.seed) 27 | 28 | if 'towel' in args.cloth_type: 29 | self.env = BimanualEnv(use_depth=True, 30 | use_cached_states=False, 31 | horizon=1, 32 | action_repeat=1, 33 | headless=args.headless, 34 | shape='default' if 'square' in args.cloth_type else 'rect') 35 | elif args.cloth_type == 'tshirt': 36 | self.env = BimanualTshirtEnv(use_depth=True, 37 | use_cached_states=False, 38 | horizon=1, 39 | action_repeat=1, 40 | headless=args.headless) 41 | 42 | goal_data = Goals(cloth_type=args.cloth_type) 43 | self.goal_loader = DataLoader(goal_data, batch_size=1, shuffle=False, num_workers=0) 44 | self.save_dir = f'{args.run_path}/rollout' 45 | if not os.path.exists(self.save_dir): 46 | os.mkdir(self.save_dir) 47 | 48 | def load_model(self, load_iter=0): 49 | picknet_cfg = OmegaConf.load(f'{self.args.run_path}/config.yaml') 50 | first_path = f"{self.args.run_path}/weights/first_{load_iter}.pt" 51 | second_path = f"{self.args.run_path}/weights/second_{load_iter}.pt" 52 | 53 | self.picknet = FlowPickSplitModel( 54 | s_pick_thres=self.args.single_pick_thresh, 55 | a_len_thres=self.args.action_len_thresh).cuda() 56 | self.picknet.first.load_state_dict(torch.load(first_path)) 57 | self.picknet.second.load_state_dict(torch.load(second_path)) 58 | self.picknet.eval() 59 | 60 | # flow model 61 | self.flownet = FlowNet(input_channels=2).cuda() 62 | checkpt = torch.load(f'{os.path.dirname(__file__)}/../data/{picknet_cfg.flow}') 63 | self.flownet.load_state_dict(checkpt['state_dict']) 64 | self.flownet.eval() 65 | 66 | def run(self, crumple_idx=None): 67 | """Main eval loop 68 | crumple_idx: which crumpled configuration to load whenever environment is reset. -1 is no crumpling 69 | """ 70 | actions = [] 71 | total_metrics = [] 72 | times = [] 73 | 74 | for goal_idx, b in enumerate(self.goal_loader): 75 | for step in range(len(b)): 76 | transition = b[step] 77 | goal_name = transition['goal_name'][0] 78 | coords_pre = transition['coords_pre'].squeeze() 79 | coords_post = transition['coords_post'].squeeze() 80 | goal_im = transition['goal_im'].cuda() 81 | 82 | if step == 0: 83 | self.env.reset() 84 | pyflex.set_positions(coords_pre) 85 | pyflex.step() 86 | 87 | for rep in range(self.args.goal_repeat): 88 | self.env.render(mode='rgb_array') 89 | _, depth = self.env.get_rgbd(cloth_only=True) 90 | depth = cv2.resize(depth, (200, 200)) 91 | curr_im = torch.FloatTensor(depth).unsqueeze(0).cuda() 92 | 93 | inp = torch.cat([curr_im, goal_im]).unsqueeze(0) 94 | flow_out = self.flownet(inp) 95 | 96 | # mask flow 97 | flow_out[0,0,:,:][inp[0,0,:,:] == 0] = 0 98 | flow_out[0,1,:,:][inp[0,0,:,:] == 0] = 0 99 | 100 | start = time.time() 101 | action, unmasked_pred = self.picknet.get_action(flow_out, curr_im.unsqueeze(0), goal_im.unsqueeze(0)) 102 | duration = time.time() - start 103 | times.append(duration) 104 | 105 | actions.append(action) 106 | next_state, reward, done, _ = self.env.step(action, pickplace=True, on_table=True) 107 | self.env.render(mode='rgb_array') 108 | 109 | img = cv2.cvtColor(next_state["color"], cv2.COLOR_RGB2BGR) 110 | img = self.action_viz(img, action, unmasked_pred) 111 | cv2.imwrite(f'{self.save_dir}/{goal_name}-step{step}-rep{rep}.png', img) 112 | 113 | metrics = self.env.compute_reward(goal_pos=coords_post[:,:3]) 114 | total_metrics.append(metrics) 115 | print(f"goal {goal_idx}: {metrics}") 116 | 117 | print(f"average action time: {np.mean(times)}") 118 | print("\nmean, std metrics: ",np.mean(total_metrics), np.std(total_metrics)) 119 | np.save(f"{self.save_dir}/actions.npy",actions) 120 | return total_metrics 121 | 122 | def action_viz(self, img, action, unmasked_pred): 123 | ''' img: cv2 image 124 | action: pick1, place1, pick2, place2 125 | unmasked_pred: pick1_pred, pick2_pred''' 126 | pick1, place1, pick2, place2 = action 127 | pick1_pred, pick2_pred = unmasked_pred 128 | 129 | # draw the masked action 130 | u1,v1 = pick1 131 | u2,v2 = place1 132 | cv2.circle(img, (int(v1),int(u1)), 6, (0,200,0), 2) 133 | cv2.arrowedLine(img, (int(v1),int(u1)), (int(v2),int(u2)), (0,200,0), 3) 134 | u1,v1 = pick2 135 | u2,v2 = place2 136 | cv2.circle(img, (int(v1),int(u1)), 6, (0,200,0), 2) 137 | cv2.arrowedLine(img, (int(v1),int(u1)), (int(v2),int(u2)), (0,200,0), 3) 138 | return img 139 | 140 | if __name__ == '__main__': 141 | parser = argparse.ArgumentParser() 142 | parser.add_argument('--run_path', help='Run to evaluate', required=True) 143 | parser.add_argument('--ckpt', help="checkpoint to evaluate, don't set to evaluate all checkpoints", type=int, default=-1) 144 | parser.add_argument('--cloth_type', help='cloth type to load', default='square_towel', choices=['square_towel', 'rect_towel', 'tshirt']) 145 | parser.add_argument('--single_pick_thresh', help='min px distance to switch dual pick to single pick', default=30) 146 | parser.add_argument('--action_len_thresh', help='min px distance for an action', default=10) 147 | parser.add_argument('--goal_repeat', help='Number of times to repeat one goal', default=1) 148 | parser.add_argument('--seed', help='random seed', default=0) 149 | parser.add_argument('--headless', help='Run headless evaluation', action='store_true') 150 | parser.add_argument('--crumple_idx', help='index for crumpled initial configuration, set to -1 for no crumpling', type=int, default=-1) 151 | args = parser.parse_args() 152 | 153 | run_name = args.run_path.split('/')[-1] 154 | output_dir = '/'.join(args.run_path.split('/')[:-1]) 155 | 156 | avg_metrics = [] 157 | full_mean = [] 158 | 159 | env = EnvRollout(args) 160 | 161 | # loop through the checkpoints to evaluate 162 | rng = range(0, 300001, 5000) if args.ckpt==-1 else range(args.ckpt, args.ckpt+1, 5000) 163 | for i in rng: 164 | print(f"loading {i}") 165 | try: 166 | env.load_model(load_iter=i) 167 | fold_mean = env.run() 168 | full_mean.append(fold_mean) 169 | print(f"mean: {np.mean(fold_mean)}") 170 | avg_metrics.append([i, np.mean(fold_mean), np.std(fold_mean)]) 171 | except EOFError: 172 | print("EOFError. skipping...") 173 | 174 | avg_metrics = np.array(avg_metrics) 175 | idx = avg_metrics[:,1].argmin() 176 | print(f"\nmin: {avg_metrics[idx,0]} fold mean/std: {avg_metrics[idx,1]*1000.0:.3f} {avg_metrics[idx, 2]*1000.0:.3f}") 177 | 178 | 179 | if args.cloth_type == 'square_towel': 180 | ms_idx = 40 181 | elif args.cloth_type == 'rect_towel': 182 | ms_idx = 3 183 | elif args.cloth_type == 'tshirt': 184 | ms_idx = 2 185 | 186 | full_mean = np.array(full_mean)*1000.0 187 | print(f"\nall: {np.mean(full_mean):.3f} {np.std(full_mean):.3f}") 188 | print(f"one-step: {np.mean(full_mean[idx,:ms_idx]):.3f} {np.std(full_mean[idx,:ms_idx]):.3f}") 189 | print(f"mul-step: {np.mean(full_mean[idx,ms_idx:]):.3f} {np.std(full_mean[idx,ms_idx:]):.3f}") 190 | 191 | with open(f'{env.save_dir}/metrics.txt', "w") as f: 192 | f.write(f'all mean/std: {np.mean(full_mean):.3f} {np.std(full_mean):.3f}') 193 | f.write(f'one-step mean/std: {np.mean(full_mean[idx,:ms_idx]):.3f} {np.std(full_mean[idx,:ms_idx]):.3f}') 194 | f.write(f'mul-step mean/std: {np.mean(full_mean[idx,ms_idx:]):.3f} {np.std(full_mean[idx,ms_idx:]):.3f}') 195 | 196 | np.save(f'{env.save_dir}/fold_metrics.npy', avg_metrics) 197 | np.save(f'{env.save_dir}/all_goals.npy', full_mean) -------------------------------------------------------------------------------- /fabricflownet/flownet/config.yaml: -------------------------------------------------------------------------------- 1 | hydra: 2 | run: 3 | dir: outputs/${experiment}/${now:%Y-%m-%d_%H%M%S}_flownet 4 | sweep: 5 | dir: multirun/${experiment}/${now:%Y-%m-%d_%H%M%S}_flownet 6 | subdir: ${hydra.job.num} 7 | 8 | experiment: dbg 9 | 10 | seed: 20 11 | 12 | base_path: path/to/base/dir # base directory containing train and val data directories 13 | train_name: null # train data directory name 14 | val_name: null # validation data directory name 15 | 16 | max_train_samples: null 17 | max_val_samples: null 18 | 19 | workers: 6 20 | batch_size: 32 21 | 22 | epochs: 300 23 | check_val_every_n_epoch: 1 24 | tboard_log_dir: tb 25 | csv_log_dir: csv 26 | 27 | net_cfg: 28 | input_channels: 2 29 | weight_decay: 0.0001 30 | lr: 0.001 31 | batchNorm: True 32 | full_upsample: False 33 | 34 | spatial_aug: 0.9 35 | spatial_trans: 5 36 | spatial_rot: 5 37 | 38 | debug_viz: 39 | remove_occlusions: False 40 | data_sample: False 41 | -------------------------------------------------------------------------------- /fabricflownet/flownet/dataset.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | import cv2 4 | import random 5 | import numpy as np 6 | import matplotlib.pyplot as plt 7 | from copy import deepcopy 8 | 9 | import torch 10 | import torchvision.transforms as T 11 | import torchvision.transforms.functional as TF 12 | from torch.utils.data import Dataset 13 | 14 | from fabricflownet.utils import remove_occluded_knots, Flow, plot_flow 15 | 16 | class FlowDataset(Dataset): 17 | def __init__(self, cfg, ids, camera_params, stage='train'): 18 | self.cfg = cfg 19 | self.camera_params = camera_params 20 | self.transform = T.Compose([T.ToTensor()]) 21 | 22 | self.data_path = f'{cfg.base_path}/{cfg.train_name}' if stage == 'train' else f'{cfg.base_path}/{cfg.val_name}' 23 | self.ids = ids 24 | self.flow = Flow() 25 | self.stage = stage 26 | 27 | def __len__(self): 28 | return len(self.ids) 29 | 30 | def _get_uv_pre(self, depth_pre, id): 31 | coords_pre = np.load(f'{self.data_path}/coords/{id}_coords_before.npy') 32 | uv_pre_float = np.load(f'{self.data_path}/knots/{id}_knots_before.npy') 33 | 34 | # Remove occluded points and save 35 | depth_pre_resized = cv2.resize(depth_pre, (720, 720)) 36 | uv_pre = remove_occluded_knots(self.camera_params, uv_pre_float, coords_pre, depth_pre_resized, 37 | zthresh=0.005, debug_viz=self.cfg.debug_viz.remove_occlusions) 38 | np.save(f'{self.data_path}/knots/{id}_visibleknots_before.npy', uv_pre) 39 | return uv_pre 40 | 41 | def __getitem__(self, index): 42 | id = self.ids[index] 43 | # Load depth before action, cloth mask, knots 44 | depth_pre = np.load(f'{self.data_path}/rendered_images/{id}_depth_before.npy') 45 | cloth_mask = (depth_pre != 0).astype(float) # 200 x 200 46 | if not os.path.exists(f'{self.data_path}/knots/{id}_visibleknots_before.npy'): 47 | uv_pre = self._get_uv_pre(depth_pre, id) 48 | else: 49 | uv_pre = np.load(f'{self.data_path}/knots/{id}_visibleknots_before.npy', allow_pickle=True) 50 | depth_pre = self.transform(depth_pre) 51 | cloth_mask = self.transform(cloth_mask) 52 | 53 | # Load depth after action and knots 54 | depth_post = np.load(f'{self.data_path}/rendered_images/{id}_depth_after.npy') 55 | uv_post_float = np.load(f'{self.data_path}/knots/{id}_knots_after.npy') 56 | depth_post = self.transform(depth_post) 57 | 58 | # Spatial augmentation 59 | if self.stage == 'train' and torch.rand(1) < self.cfg.spatial_aug: 60 | depth_pre, depth_post, cloth_mask, uv_pre, uv_post_float = \ 61 | self._spatial_aug(depth_pre, depth_post, cloth_mask, uv_pre, uv_post_float) 62 | 63 | # Remove out of bounds 64 | uv_pre[uv_pre < 0] = float('NaN') 65 | uv_pre[uv_pre >= 720] = float('NaN') 66 | 67 | # Get flow image 68 | flow_gt = self.flow.get_flow_image(uv_pre, uv_post_float) 69 | flow_gt = self.transform(flow_gt) 70 | 71 | # Get loss mask 72 | loss_mask = torch.zeros((flow_gt.shape[1], flow_gt.shape[2]), dtype=torch.float32) 73 | non_nan_idxs = np.rint(uv_pre[~np.isnan(uv_pre).any(axis=1)]/719*199).astype(int) 74 | loss_mask[non_nan_idxs[:, 1], non_nan_idxs[:, 0]] = 1 75 | loss_mask = loss_mask.unsqueeze(0) 76 | 77 | # Construct sample 78 | depths = torch.cat([depth_pre, depth_post], axis=0) 79 | sample = {'depths': depths, 'flow_gt': flow_gt, 'loss_mask': loss_mask, 'cloth_mask': cloth_mask} 80 | 81 | # Debug plotting 82 | if self.cfg.debug_viz.data_sample and self.stage == 'train': 83 | depth_pre_np = depth_pre.squeeze().numpy() 84 | depth_post_np = depth_post.squeeze().numpy() 85 | flow_gt_np = flow_gt.permute(1, 2, 0).numpy() 86 | loss_mask_np = loss_mask.squeeze().numpy() 87 | self._plot(depth_pre_np, depth_post_np, flow_gt_np, loss_mask_np) 88 | return sample 89 | 90 | def _aug_uv(self, uv, angle, dx, dy): 91 | rad = np.deg2rad(-angle) 92 | R = np.array([ 93 | [np.cos(rad), -np.sin(rad)], 94 | [np.sin(rad), np.cos(rad)]]) 95 | uv -= 719 / 2 96 | uv = np.dot(R, uv.T).T 97 | uv += 719 / 2 98 | uv[:, 0] += dx 99 | uv[:, 1] += dy 100 | return uv 101 | 102 | def _spatial_aug(self, depth_pre, depth_post, cloth_mask, uv_pre, uv_post_float): 103 | spatial_rot = self.cfg.spatial_rot 104 | spatial_trans = self.cfg.spatial_trans 105 | angle = torch.randint(low=-spatial_rot, high=spatial_rot+1, size=(1,), dtype=torch.float32) 106 | dx = np.random.randint(-spatial_trans, spatial_trans+1) 107 | dy = np.random.randint(-spatial_trans, spatial_trans+1) 108 | depth_pre = TF.affine(depth_pre, angle=angle.item(), translate=(dx, dy), scale=1.0, shear=0) 109 | depth_post = TF.affine(depth_post, angle=angle.item(), translate=(dx, dy), scale=1.0, shear=0) 110 | cloth_mask = TF.affine(cloth_mask, angle=angle.item(), translate=(dx, dy), scale=1.0, shear=0) 111 | uv_pre = self._aug_uv(uv_pre, -angle, dx/199*719, dy/199*719) 112 | uv_post_float = self._aug_uv(uv_post_float, -angle, dx/199*719, dy/199*719) 113 | return depth_pre, depth_post, cloth_mask, uv_pre, uv_post_float 114 | 115 | def _plot(self, depth_pre, depth_post, flow_gt, loss_mask): 116 | fig, ax = plt.subplots(1, 4, figsize=(12, 3)) 117 | ax[0].set_title("depth before") 118 | ax[0].imshow(depth_pre) 119 | ax[1].set_title("depth after") 120 | ax[1].imshow(depth_post) 121 | ax[2].set_title("ground-truth flow") 122 | plot_flow(ax[2], flow_gt) 123 | ax[3].set_title("loss mask") 124 | ax[3].imshow(loss_mask) 125 | plt.tight_layout() 126 | plt.show() 127 | -------------------------------------------------------------------------------- /fabricflownet/flownet/models.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.nn import init 4 | 5 | import math 6 | import numpy as np 7 | import pytorch_lightning as pl 8 | 9 | import matplotlib as mpl 10 | import matplotlib.pyplot as plt 11 | from fabricflownet.utils import plot_flow 12 | 13 | def conv(batchNorm, in_planes, out_planes, kernel_size=3, stride=1): 14 | if batchNorm: 15 | return nn.Sequential( 16 | nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=False), 17 | nn.BatchNorm2d(out_planes), 18 | nn.LeakyReLU(0.1,inplace=True) 19 | ) 20 | else: 21 | return nn.Sequential( 22 | nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=True), 23 | nn.LeakyReLU(0.1,inplace=True) 24 | ) 25 | 26 | def predict_flow(in_planes): 27 | return nn.Conv2d(in_planes,2,kernel_size=3,stride=1,padding=1,bias=True) 28 | 29 | def deconv(in_planes, out_planes, ksize=3): 30 | return nn.Sequential( 31 | nn.ConvTranspose2d(in_planes, out_planes, kernel_size=ksize, stride=2, padding=1, bias=True), 32 | nn.LeakyReLU(0.1,inplace=True) 33 | ) 34 | 35 | class FlowNet(pl.LightningModule): 36 | def __init__(self, 37 | input_channels = 2, 38 | batchNorm=True, 39 | lr=0.0001, 40 | weight_decay=0.0001): 41 | super(FlowNet,self).__init__() 42 | 43 | self.lr = lr 44 | self.weight_decay = weight_decay 45 | 46 | fs = [8, 16, 32, 64, 128] # filter sizes 47 | self.batchNorm = batchNorm 48 | self.conv1 = conv(self.batchNorm, input_channels, fs[0], kernel_size=7, stride=2) # 384 -> (384 - 7 + 2*3)/2 + 1 = 377 49 | self.conv2 = conv(self.batchNorm, fs[0], fs[1], kernel_size=5, stride=2) 50 | self.conv3 = conv(self.batchNorm, fs[1], fs[2], kernel_size=5, stride=2) 51 | self.conv3_1 = conv(self.batchNorm, fs[2], fs[2]) 52 | self.conv4 = conv(self.batchNorm, fs[2], fs[3], stride=2) 53 | self.conv4_1 = conv(self.batchNorm, fs[3], fs[3]) 54 | self.conv5 = conv(self.batchNorm, fs[3], fs[3], stride=2) 55 | self.conv5_1 = conv(self.batchNorm, fs[3], fs[3]) 56 | self.conv6 = conv(self.batchNorm, fs[3], fs[4], stride=2) 57 | self.conv6_1 = conv(self.batchNorm, fs[4], fs[4]) 58 | 59 | self.deconv5 = deconv(fs[4],fs[3]) 60 | self.deconv4 = deconv(fs[3]+fs[3]+2,fs[2]) 61 | self.deconv3 = deconv(fs[3]+fs[2]+2,fs[1]) 62 | self.deconv2 = deconv(fs[2]+fs[1]+2,fs[0], ksize=4) 63 | 64 | self.predict_flow6 = predict_flow(fs[4]) 65 | self.predict_flow5 = predict_flow(fs[3]+fs[3]+2) 66 | self.predict_flow4 = predict_flow(fs[3]+fs[2]+2) 67 | self.predict_flow3 = predict_flow(fs[2]+fs[1]+2) 68 | self.predict_flow2 = predict_flow(fs[1]+fs[0]+2) 69 | 70 | self.upsampled_flow6_to_5 = nn.ConvTranspose2d(2, 2, 3, 2, 1, bias=False) # (H_in-1)*stride - 2*padding + (kernel-1) + 1 71 | self.upsampled_flow5_to_4 = nn.ConvTranspose2d(2, 2, 3, 2, 1, bias=False) 72 | self.upsampled_flow4_to_3 = nn.ConvTranspose2d(2, 2, 3, 2, 1, bias=False) 73 | self.upsampled_flow3_to_2 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=False) 74 | 75 | for m in self.modules(): 76 | if isinstance(m, nn.Conv2d): 77 | if m.bias is not None: 78 | init.uniform_(m.bias) 79 | init.xavier_uniform_(m.weight) 80 | 81 | if isinstance(m, nn.ConvTranspose2d): 82 | if m.bias is not None: 83 | init.uniform_(m.bias) 84 | init.xavier_uniform_(m.weight) 85 | self.upsample1 = nn.Upsample(scale_factor=4, mode='bilinear') 86 | 87 | def forward(self, x): 88 | out_conv1 = self.conv1(x) 89 | out_conv2 = self.conv2(out_conv1) 90 | out_conv3 = self.conv3_1(self.conv3(out_conv2)) 91 | out_conv4 = self.conv4_1(self.conv4(out_conv3)) 92 | out_conv5 = self.conv5_1(self.conv5(out_conv4)) 93 | out_conv6 = self.conv6_1(self.conv6(out_conv5)) 94 | 95 | flow6 = self.predict_flow6(out_conv6) 96 | flow6_up = self.upsampled_flow6_to_5(flow6) 97 | out_deconv5 = self.deconv5(out_conv6) 98 | 99 | concat5 = torch.cat((out_conv5,out_deconv5,flow6_up),1) 100 | flow5 = self.predict_flow5(concat5) 101 | flow5_up = self.upsampled_flow5_to_4(flow5) 102 | out_deconv4 = self.deconv4(concat5) 103 | 104 | concat4 = torch.cat((out_conv4,out_deconv4,flow5_up),1) 105 | flow4 = self.predict_flow4(concat4) 106 | flow4_up = self.upsampled_flow4_to_3(flow4) 107 | out_deconv3 = self.deconv3(concat4) 108 | 109 | concat3 = torch.cat((out_conv3,out_deconv3,flow4_up),1) 110 | flow3 = self.predict_flow3(concat3) 111 | flow3_up = self.upsampled_flow3_to_2(flow3) 112 | out_deconv2 = self.deconv2(concat3) 113 | 114 | concat2 = torch.cat((out_conv2,out_deconv2,flow3_up),1) 115 | flow2 = self.predict_flow2(concat2) 116 | 117 | out = self.upsample1(flow2) 118 | return out 119 | 120 | def loss(self, input_flow, target_flow, mask): 121 | b, c, h, w = input_flow.size() 122 | diff_flow = torch.reshape(target_flow - input_flow*mask, (b, c, h*w)) 123 | mask = torch.reshape(mask, (b, h*w)) 124 | norm_diff_flow = torch.linalg.norm(diff_flow, ord=2, dim=1) # B x 40000 get norm of flow vector diff 125 | mean_norm_diff_flow = norm_diff_flow.sum(dim=1) / mask.sum(dim=1) # B x 1 get average norm for each image 126 | batch_mean_diff_flow = mean_norm_diff_flow.mean() # mean over the batch 127 | return batch_mean_diff_flow 128 | 129 | def training_step(self, batch, batch_idx): 130 | depth_input = batch['depths'] 131 | flow_gt = batch['flow_gt'] 132 | loss_mask = batch['loss_mask'] 133 | flow_out = self.forward(depth_input) 134 | train_loss = self.loss(flow_out, flow_gt, loss_mask) 135 | loss = train_loss 136 | 137 | if batch_idx == 0: 138 | self.plot(depth_input, loss_mask, flow_gt, flow_out, stage="train") 139 | self.log('loss/train', loss) 140 | return {'loss': loss} 141 | 142 | def validation_step(self, batch, batch_idx): 143 | depth_input = batch['depths'] 144 | flow_gt = batch['flow_gt'] 145 | loss_mask = batch['loss_mask'] 146 | flow_out = self.forward(depth_input) 147 | loss = self.loss(flow_out, flow_gt, loss_mask) 148 | 149 | if batch_idx == 0: 150 | self.plot(depth_input, loss_mask, flow_gt, flow_out, stage="val") 151 | self.log('loss/val', loss) 152 | return {'loss': loss} 153 | 154 | def configure_optimizers(self): 155 | optimizer = torch.optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay) 156 | return optimizer 157 | 158 | def plot(self, depth_input, loss_mask, flow_gt, flow_out, stage): 159 | im1 = depth_input[0, 0].detach().cpu().numpy() 160 | im2 = depth_input[0, 1].detach().cpu().numpy() 161 | loss_mask = loss_mask[0].detach().squeeze().cpu().numpy() 162 | flow_gt = flow_gt[0].detach().permute(1, 2, 0).cpu().numpy() 163 | flow_out = flow_out[0].detach().permute(1, 2, 0).cpu().numpy() 164 | 165 | fig, ax = plt.subplots(1, 5, figsize=(12, 3), dpi=300) 166 | ax[0].set_title("depth before") 167 | ax[0].imshow(im1) 168 | ax[1].set_title("depth after") 169 | ax[1].imshow(im2) 170 | 171 | ax[2].set_title("ground-truth flow") 172 | plot_flow(ax[2], flow_gt) 173 | 174 | ax[3].set_title("predicted flow (masked)") 175 | flow_out[loss_mask == 0, :] = 0 176 | plot_flow(ax[3], flow_out) 177 | 178 | ax[4].set_title("loss mask") 179 | ax[4].imshow(loss_mask) 180 | 181 | plt.tight_layout() 182 | self.logger[1].experiment.add_figure(stage, fig, self.global_step) # tensorboard 183 | plt.close() 184 | -------------------------------------------------------------------------------- /fabricflownet/flownet/train.py: -------------------------------------------------------------------------------- 1 | from fabricflownet.flownet.dataset import FlowDataset 2 | 3 | import os 4 | import sys 5 | import numpy as np 6 | import matplotlib.pyplot as plt 7 | from copy import deepcopy 8 | from torch.utils.data import DataLoader 9 | import hydra 10 | import pytorch_lightning.utilities.seed as seed_utils 11 | import pytorch_lightning as pl 12 | from pytorch_lightning import loggers as pl_loggers 13 | from pytorch_lightning.callbacks import ModelCheckpoint 14 | 15 | from fabricflownet.flownet.models import FlowNet 16 | 17 | def get_flownet_dataloaders(cfg): 18 | # Get training samples 19 | trainpath = f'{cfg.base_path}/{cfg.train_name}' 20 | trainfs = sorted(['_'.join(fn.split('_')[0:2]) 21 | for fn in os.listdir(f'{trainpath}/actions')]) 22 | if cfg['max_train_samples'] != None: 23 | trainfs = trainfs[:cfg['max_train_samples']] 24 | print(f"Max training set: {len(trainfs)}") 25 | 26 | # Get validation samples 27 | valpath = f'{cfg.base_path}/{cfg.val_name}' 28 | valfs = sorted(['_'.join(fn.split('_')[0:2]) 29 | for fn in os.listdir(f'{valpath}/actions')]) 30 | if cfg['max_val_samples'] != None: 31 | valfs = valfs[:cfg['max_val_samples']] 32 | print(f"Max val set: {len(valfs)}") 33 | 34 | # Get camera params 35 | train_camera_params = np.load(f"{trainpath}/camera_params.npy", allow_pickle=True)[()] 36 | val_camera_params = np.load(f"{valpath}/camera_params.npy", allow_pickle=True)[()] 37 | np.testing.assert_equal(val_camera_params, train_camera_params) 38 | camera_params = train_camera_params 39 | 40 | # Get datasets 41 | train_data = FlowDataset(cfg, trainfs, camera_params, stage='train') 42 | val_data = FlowDataset(cfg, valfs, camera_params, stage='val') 43 | train_loader = DataLoader(train_data, batch_size=cfg.batch_size, shuffle=True, num_workers=cfg.workers, persistent_workers=cfg.workers>0) 44 | val_loader = DataLoader(val_data, batch_size=cfg.batch_size, shuffle=False, num_workers=cfg.workers, persistent_workers=cfg.workers>0) 45 | return train_loader, val_loader 46 | 47 | @hydra.main(config_name="config") 48 | def main(cfg): 49 | seed_utils.seed_everything(cfg.seed) 50 | with open('.hydra/command.txt', 'w') as f: 51 | f.write('python ' + ' '.join(sys.argv)) 52 | 53 | train_loader, val_loader = get_flownet_dataloaders(cfg) 54 | csv_logger = pl_loggers.CSVLogger(save_dir=cfg.csv_log_dir) 55 | tb_logger = pl_loggers.TensorBoardLogger(save_dir=cfg.tboard_log_dir) 56 | chkpt_cb = ModelCheckpoint(monitor='loss/val', save_last=True, save_top_k=-1, every_n_val_epochs=10) 57 | trainer = pl.Trainer(gpus=[0], 58 | logger=[csv_logger, tb_logger], 59 | max_epochs=cfg.epochs, 60 | check_val_every_n_epoch=cfg.check_val_every_n_epoch, 61 | log_every_n_steps=len(train_loader) if len(train_loader) < 50 else 50, 62 | callbacks=[chkpt_cb]) 63 | 64 | flownet = FlowNet(**cfg.net_cfg) 65 | trainer.fit(flownet, train_loader, val_loader) 66 | 67 | if __name__ == '__main__': 68 | main() 69 | -------------------------------------------------------------------------------- /fabricflownet/picknet/config.yaml: -------------------------------------------------------------------------------- 1 | hydra: 2 | run: 3 | dir: outputs/${experiment}/${now:%Y-%m-%d_%H%M%S}_picknet 4 | sweep: 5 | dir: multirun/${experiment}/${now:%Y-%m-%d_%H%M%S}_picknet 6 | subdir: ${hydra.job.num} 7 | 8 | experiment: dbg 9 | 10 | seed: 0 11 | 12 | base_path: path/to/base/dir # base directory containing train and val dirs 13 | train_name: null # train data directory name 14 | val_name: null # validation data directory name 15 | 16 | epochs: 150 # max training epochs 17 | max_buf: 20000 # 7200 # max training buf 18 | workers: 6 19 | batch_size: 10 # batch size 20 | 21 | check_val_every_n_epoch: 1 22 | tboard_log_dir: tb 23 | csv_log_dir: csv 24 | 25 | net_cfg: 26 | lr: 0.0001 27 | input_mode: flowonly 28 | min_loss: True 29 | model_type: split 30 | pick: True # Whether it is pick network or place network 31 | im_width: 200 32 | 33 | flow: /path/to/flow/ckpt 34 | augment: False # dataset spatial aug flag 35 | 36 | debug_viz: 37 | remove_occlusions: False 38 | data_sample: False -------------------------------------------------------------------------------- /fabricflownet/picknet/dataset.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | from torch.utils.data import Dataset 4 | import torchvision.transforms.functional as TF 5 | import torchvision.transforms as transforms 6 | from PIL import Image 7 | from fabricflownet.utils import remove_occluded_knots, flow_affinewarp, Flow, plot_flow 8 | import cv2 9 | from copy import deepcopy 10 | import os.path as osp 11 | import os 12 | import matplotlib.pyplot as plt 13 | from fabricflownet.flownet.models import FlowNet 14 | 15 | class PickNetDataset(Dataset): 16 | def __init__(self, camera_params, config, ids, mode='train', pick_pt=True): 17 | self.cfg = config 18 | self.mode = mode 19 | self.ids = ids 20 | self.camera_params = camera_params 21 | self.pick_pt = pick_pt 22 | 23 | if mode == 'train': 24 | self.data_path = f"{self.cfg.base_path}/{self.cfg.train_name}" 25 | else: 26 | self.data_path = f"{self.cfg.base_path}/{self.cfg.val_name}" 27 | 28 | if self.cfg.flow == 'gt': 29 | self.gt_flow = True 30 | self.flow = Flow() 31 | if not osp.exists(osp.join(self.data_path, "flow_gt")): 32 | os.mkdir(osp.join(self.data_path, "flow_gt")) 33 | else: 34 | self.gt_flow = False 35 | self.flow = FlowNet(input_channels=2) 36 | checkpt = torch.load(self.cfg.flow) 37 | self.flow.load_state_dict(checkpt['state_dict']) 38 | self.flow.eval() 39 | if not osp.exists(osp.join(self.data_path, "flow_pred")): 40 | os.mkdir(osp.join(self.data_path, "flow_pred")) 41 | 42 | def __len__(self): 43 | return len(self.ids) 44 | 45 | def __getitem__(self, index): 46 | id = self.ids[index] 47 | action = np.load(f'{self.data_path}/actions/{id}_action.npy') 48 | pick_uv1 = action[0] 49 | place_uv1 = action[1] 50 | pick_uv2 = action[2] 51 | place_uv2 = action[3] 52 | 53 | depth_post = np.load(f'{self.data_path}/rendered_images/{id}_depth_after.npy') 54 | depth_pre = np.load(f'{self.data_path}/rendered_images/{id}_depth_before.npy') 55 | 56 | if self.gt_flow: 57 | uv_post_float = np.load(f'{self.data_path}/knots/{id}_knots_after.npy') 58 | 59 | # Compute and save flow if it does not exist already 60 | if not osp.exists(osp.join(self.data_path, "flow_gt", f"{id}_flow.npy")): 61 | coords_pre = np.load(f'{self.data_path}/coords/{id}_coords_before.npy') 62 | uv_pre_float = np.load(f'{self.data_path}/knots/{id}_knots_before.npy') 63 | 64 | # Remove occlusions 65 | depth_pre_rs = cv2.resize(depth_pre, (720, 720)) 66 | uv_pre, _ = remove_occluded_knots(self.camera_params, uv_pre_float, coords_pre, depth_pre_rs, 67 | zthresh=0.005, debug_viz=self.cfg.debug_viz.remove_occlusions) 68 | 69 | # Remove out of bounds 70 | uv_pre[uv_pre < 0] = float('NaN') 71 | uv_pre[uv_pre >= 720] = float('NaN') 72 | 73 | # Get flow image 74 | flow_im = self.flow.get_image(uv_pre, uv_post_float) 75 | 76 | # Save the flow 77 | np.save(osp.join(self.data_path, "flow_gt", f"{id}_flow.npy"), flow_im) 78 | else: 79 | flow_im = np.load(osp.join(self.data_path, "flow_gt", f"{id}_flow.npy")) 80 | else: 81 | if not osp.exists(osp.join(self.data_path, "flow_pred", f"{id}_flow.npy")): 82 | inp = torch.stack([torch.FloatTensor(depth_pre), torch.FloatTensor(depth_post)]).unsqueeze(0) 83 | flow_im = self.flow(inp) 84 | flow_im = flow_im.squeeze().cpu() 85 | np.save(osp.join(self.data_path, "flow_pred", f"{id}_flow.npy"), flow_im.detach().numpy()) 86 | else: 87 | flow_im = np.load(osp.join(self.data_path, "flow_pred", f"{id}_flow.npy"), allow_pickle=True) 88 | 89 | depth_pre = torch.FloatTensor(depth_pre).unsqueeze(0) 90 | depth_post = torch.FloatTensor(depth_post).unsqueeze(0) 91 | 92 | if self.gt_flow: 93 | flow_im = flow_im.transpose([2,0,1]) 94 | 95 | if not isinstance(flow_im, torch.Tensor): 96 | flow_im = torch.FloatTensor(flow_im) 97 | 98 | # mask flow 99 | flow_im[0,:,:][depth_pre[0] == 0] = 0 100 | flow_im[1,:,:][depth_pre[0] == 0] = 0 101 | 102 | if self.cfg.augment: 103 | angle = np.random.randint(-5, 6) 104 | dx = np.random.randint(-5, 6) 105 | dy = np.random.randint(-5, 6) 106 | depth_pre, depth_post, pick_uv1, pick_uv2, place_uv1, place_uv2 \ 107 | = self.spatial_aug(depth_pre, depth_post, pick_uv1, pick_uv2, place_uv1, place_uv2, angle, dx, dy) 108 | flow_im = flow_im.permute(1, 2, 0).detach().numpy() 109 | flow_im = flow_affinewarp(flow_im, -angle, 0, 0) 110 | flow_im = torch.FloatTensor(flow_im).permute(2, 0, 1) 111 | 112 | if self.pick_pt: 113 | uv1, uv2 = pick_uv1, pick_uv2 114 | else: 115 | uv1, uv2 = place_uv1, place_uv2 116 | 117 | if self.cfg.debug_viz.data_sample: 118 | self.plot(depth_pre, depth_post, flow_im, uv1, uv2) 119 | 120 | return depth_pre, depth_post, flow_im, uv1, uv2 121 | 122 | def spatial_aug(self, depth_pre, depth_post, pick_uv1, pick_uv2, place_uv1, place_uv2, angle, dx, dy): 123 | depth_pre = TF.affine(depth_pre, angle=angle, translate=(dx, dy), scale=1.0, shear=0) 124 | depth_post = TF.affine(depth_post, angle=angle, translate=(dx, dy), scale=1.0, shear=0) 125 | pick_uv1 = self.aug_uv(pick_uv1.astype(np.float64)[None,:], -angle, dx, dy, size=199) 126 | pick_uv1 = pick_uv1.squeeze().astype(int) 127 | pick_uv2 = self.aug_uv(pick_uv2.astype(np.float64)[None,:], -angle, dx, dy, size=199) 128 | pick_uv2 = pick_uv2.squeeze().astype(int) 129 | place_uv1 = self.aug_uv(place_uv1.astype(np.float64)[None,:], -angle, dx, dy, size=199) 130 | place_uv1 = place_uv1.squeeze().astype(int) 131 | place_uv2 = self.aug_uv(place_uv2.astype(np.float64)[None,:], -angle, dx, dy, size=199) 132 | place_uv2 = place_uv2.squeeze().astype(int) 133 | return depth_pre, depth_post, pick_uv1, pick_uv2, place_uv1, place_uv2 134 | 135 | def aug_uv(self, uv, angle, dx, dy, size=719): 136 | rad = np.deg2rad(-angle) 137 | R = np.array([ 138 | [np.cos(rad), -np.sin(rad)], 139 | [np.sin(rad), np.cos(rad)]]) 140 | uv -= size / 2 141 | uv = np.dot(R, uv.T).T 142 | uv += size / 2 143 | uv[:, 0] += dx 144 | uv[:, 1] += dy 145 | uv = np.clip(uv, 0, size) 146 | return uv 147 | 148 | def plot(self, depth_pre, depth_post, flow, uv1, uv2): 149 | depth_pre = depth_pre.squeeze(0).numpy() 150 | depth_post = depth_post.squeeze(0).numpy() 151 | flow = flow.detach().permute(1, 2, 0).numpy() 152 | fig, ax = plt.subplots(1, 3, figsize=(12, 3)) 153 | ax[0].set_title("depth before") 154 | ax[0].imshow(depth_pre) 155 | ax[0].scatter(uv1[0], uv1[1], label='pick_uv1' if self.pick_pt else 'place_uv1') 156 | ax[0].scatter(uv2[0], uv2[1], label='pick_uv2' if self.pick_pt else 'place_uv2') 157 | ax[1].set_title("depth after") 158 | ax[1].imshow(depth_post) 159 | ax[2].set_title(f"{'ground_truth' if self.gt_flow else 'predicted'} flow") 160 | plot_flow(ax[2], flow, skip=0.05) 161 | plt.tight_layout() 162 | plt.show() 163 | 164 | class Goals(Dataset): 165 | def __init__(self, cloth_type='square_towel'): 166 | self.clothtype_path = f'{os.path.abspath(os.path.dirname(__file__))}/../../data/goals/{cloth_type}' 167 | 168 | # Load start state 169 | self.coords_start = np.load(f'{self.clothtype_path}/start.npy') 170 | 171 | # Load goals 172 | self.goals = [] 173 | 174 | # Load single step goals 175 | singlestep_path = f'{self.clothtype_path}/single_step' 176 | goal_names = sorted([x.replace('.png', '') for x in os.listdir(f'{singlestep_path}/rgb')], key=lambda x: int(x.split('_')[-1])) 177 | for goal_name in goal_names: 178 | goal_im = cv2.imread(f'{singlestep_path}/depth/{goal_name}_depth.png')[:, :, 0] / 255 179 | coords_post = np.load(f'{singlestep_path}/coords/{goal_name}.npy') 180 | self.goals.append([ 181 | { 182 | "goal_name": goal_name, 183 | "goal_im": torch.FloatTensor(goal_im), 184 | "coords_pre": self.coords_start, 185 | "coords_post": coords_post, 186 | } 187 | ]) 188 | 189 | # Load multi step goals 190 | multistep_path = f'{self.clothtype_path}/multi_step' 191 | goal_names = sorted([x.replace('.png', '') for x in os.listdir(f'{multistep_path}/rgb')]) 192 | curr_goals = [] 193 | curr_multistep_goal = goal_names[0][:-2] # Get multi-step goal name without step number at end 194 | for goal_name in goal_names: 195 | multistep_goal = goal_name[:-2] 196 | goal_im = cv2.imread(f'{multistep_path}/depth/{goal_name}_depth.png')[:, :, 0] / 255 197 | coords_post = np.load(f'{multistep_path}/coords/{goal_name}.npy') 198 | goal = { 199 | "goal_name": goal_name, 200 | "goal_im": torch.FloatTensor(goal_im), 201 | "coords_pre": self.coords_start, 202 | "coords_post": coords_post, 203 | } 204 | 205 | if curr_multistep_goal == multistep_goal: # add new goal to curr_goals 206 | curr_goals.append(goal) 207 | else: # switch to new goal set 208 | self.goals.append(curr_goals) 209 | curr_goals = [goal] 210 | curr_multistep_goal = multistep_goal 211 | self.goals.append(curr_goals) 212 | 213 | def __len__(self): 214 | return len(self.goals) 215 | 216 | def __getitem__(self, index): 217 | goal_sequence = self.goals[index] 218 | return goal_sequence 219 | -------------------------------------------------------------------------------- /fabricflownet/picknet/models.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | from torch import nn 4 | import pytorch_lightning as pl 5 | import torch.nn.functional as F 6 | import matplotlib.pyplot as plt 7 | from fabricflownet.utils import plot_flow 8 | import cv2 9 | 10 | class FlowPickSplit(nn.Module): 11 | def __init__(self, inchannels, im_w, second=False): 12 | super(FlowPickSplit, self).__init__() 13 | self.trunk = nn.Sequential(nn.Conv2d(inchannels, 32, 5, 2), 14 | nn.ReLU(True), 15 | nn.Conv2d(32,32, 5, 2), 16 | nn.ReLU(True), 17 | nn.Conv2d(32,32, 5, 2), 18 | nn.ReLU(True), 19 | nn.Conv2d(32,32, 5, 1), 20 | nn.ReLU(True)) 21 | self.head = nn.Sequential(nn.Conv2d(32,32, 3, 1), 22 | nn.ReLU(True), 23 | nn.UpsamplingBilinear2d(scale_factor=2), 24 | nn.Conv2d(32,1, 3, 1)) 25 | 26 | self.im_w = im_w 27 | self.second = second 28 | self.upsample = nn.Upsample(size=(20,20), mode="bilinear") 29 | 30 | def forward(self, x): 31 | x = self.trunk(x) 32 | out = self.head(x) 33 | out = self.upsample(out) 34 | return out 35 | 36 | class FlowPickSplitModel(pl.LightningModule): 37 | def __init__(self, 38 | lr=0.0001, 39 | input_mode='flowonly', 40 | model_type='split', 41 | min_loss=True, 42 | pick=True, 43 | s_pick_thres = 30, 44 | a_len_thres = 10, 45 | im_width=200): 46 | super(FlowPickSplitModel,self).__init__() 47 | self.automatic_optimization = False 48 | self.lr = lr 49 | self.model_type = model_type 50 | self.im_width = im_width 51 | self.min_loss = min_loss 52 | self.pick = pick 53 | self.input_mode = input_mode 54 | self.s_pick_thres = s_pick_thres 55 | self.a_len_thres = a_len_thres 56 | self.init_models() 57 | 58 | def init_models(self): 59 | if self.model_type == 'split': 60 | self.first = FlowPickSplit(2, self.im_width) 61 | self.second = FlowPickSplit(3, self.im_width, second=True) 62 | else: 63 | raise NotImplementedError 64 | 65 | def nearest_to_mask(self, u, v, depth): 66 | mask_idx = np.argwhere(depth) 67 | nearest_idx = mask_idx[((mask_idx - [u,v])**2).sum(1).argmin()] 68 | return nearest_idx 69 | 70 | def get_flow_place_pt(self, u,v, flow): 71 | """ compute place point using flow 72 | """ 73 | flow_u_idxs = np.argwhere(flow[0,:,:]) 74 | flow_v_idxs = np.argwhere(flow[1,:,:]) 75 | nearest_u_idx = flow_u_idxs[((flow_u_idxs - [u,v])**2).sum(1).argmin()] 76 | nearest_v_idx = flow_v_idxs[((flow_v_idxs - [u,v])**2).sum(1).argmin()] 77 | 78 | flow_u = flow[0,nearest_u_idx[0],nearest_u_idx[1]] 79 | flow_v = flow[1,nearest_v_idx[0],nearest_v_idx[1]] 80 | 81 | new_u = np.clip(u + flow_u, 0, 199) 82 | new_v = np.clip(v + flow_v, 0, 199) 83 | 84 | return new_u,new_v 85 | 86 | def get_action(self, flow, depth_pre, depth_post): 87 | pick_uv1, pick_uv2, info = self.forward(flow, depth_pre, depth_post) 88 | depth_pre_np = depth_pre.detach().squeeze().cpu().numpy() 89 | pick1 = self.nearest_to_mask(pick_uv1[0], pick_uv1[1], depth_pre_np) 90 | pick2 = self.nearest_to_mask(pick_uv2[0], pick_uv2[1], depth_pre_np) 91 | pickmask_u1, pickmask_v1 = pick1 92 | pickmask_u2, pickmask_v2 = pick2 93 | 94 | flow_np = flow.detach().squeeze(0).cpu().numpy() 95 | place_u1,place_v1 = self.get_flow_place_pt(pickmask_u1,pickmask_v1,flow_np) 96 | place_u2,place_v2 = self.get_flow_place_pt(pickmask_u2,pickmask_v2,flow_np) 97 | place1 = np.array([place_u1, place_v1]) 98 | place2 = np.array([place_u2, place_v2]) 99 | 100 | pred_1 = np.array(pick_uv1) 101 | pred_2 = np.array(pick_uv2) 102 | 103 | # single action threshold 104 | if self.s_pick_thres > 0 and np.linalg.norm(pick1-pick2) < self.s_pick_thres: 105 | pick2 = np.array([0,0]) 106 | place2 = np.array([0,0]) 107 | 108 | # action size threshold 109 | if self.a_len_thres > 0 and np.linalg.norm(pick1-place1) < self.a_len_thres: 110 | pick1 = np.array([0,0]) 111 | place1 = np.array([0,0]) 112 | 113 | if self.a_len_thres > 0 and np.linalg.norm(pick2-place2) < self.a_len_thres: 114 | pick2 = np.array([0,0]) 115 | place2 = np.array([0,0]) 116 | 117 | return np.array([pick1, place1, pick2, place2]), np.array([pred_1, pred_2]) 118 | 119 | def forward(self, flow, depth_pre, depth_post): 120 | logits1 = self.first(flow) 121 | u1, v1 = self.get_pt(logits1) 122 | pick1_gau = self.get_gaussian(u1,v1) 123 | 124 | if self.input_mode == 'flowonly': 125 | x2 = torch.cat([flow, pick1_gau], dim=1) 126 | else: 127 | # x2 = torch.cat([depth_pre.detach().clone(), depth_post.detach().clone(), pick1_gau.detach().clone()], dim=1) 128 | raise NotImplementedError 129 | 130 | logits2 = self.second(x2) 131 | u2, v2 = self.get_pt(logits2) 132 | info = { 133 | 'logits1': logits1, 134 | 'logits2': logits2, 135 | 'pick1_gau': pick1_gau 136 | } 137 | return [u1, v1], [u2, v2], info 138 | 139 | def get_pt(self, logits): 140 | N = logits.size(0) 141 | W = logits.size(2) 142 | 143 | prdepth_pre = torch.sigmoid(logits) 144 | prdepth_pre = prdepth_pre.view(N,1,W*W) 145 | val,idx = torch.max(prdepth_pre[:,0], 1) 146 | 147 | # u = (idx % 20) * 10 148 | # v = (idx // 20) * 10 149 | u = (idx // 20) * 10 150 | v = (idx % 20) * 10 151 | return u.item(),v.item() 152 | 153 | def get_gaussian(self, u, v, sigma=5, size=None): 154 | if size is None: 155 | size = self.im_width 156 | 157 | x0, y0 = torch.Tensor([u]).cuda(), torch.Tensor([v]).cuda() 158 | x0 = x0[:, None] 159 | y0 = y0[:, None] 160 | # sigma = torch.tensor(sigma, dtype=torch.float32, device=self.device) 161 | 162 | # N = u.size(0) 163 | N = 1 # u.size(0) 164 | num = torch.arange(size).float() 165 | x, y = torch.vstack([num]*N).to(self.device), torch.vstack([num]*N).to(self.device) 166 | gx = torch.exp(-(x-x0)**2/(2*sigma**2)) 167 | gy = torch.exp(-(y-y0)**2/(2*sigma**2)) 168 | g = torch.einsum('ni,no->nio', gx, gy) 169 | 170 | gmin = g.amin(dim=(1,2)) 171 | gmax = g.amax(dim=(1,2)) 172 | g = (g - gmin[:,None,None])/(gmax[:,None,None] - gmin[:,None,None]) 173 | g = g.unsqueeze(1) 174 | return g 175 | 176 | def loss(self, logits1, logits2, pick_uv1, pick_uv2): 177 | N = logits1.size(0) 178 | W = logits1.size(2) 179 | 180 | pick_uv1 = pick_uv1.cuda() 181 | pick_uv2 = pick_uv2.cuda() 182 | label_a = self.get_gaussian(pick_uv1[:,0] // 10, pick_uv1[:,1] // 10, sigma=2, size=20) 183 | label_b = self.get_gaussian(pick_uv2[:,0] // 10, pick_uv2[:,1] // 10, sigma=2, size=20) 184 | 185 | if self.min_loss: 186 | loss_1a = torch.mean(F.binary_cross_entropy_with_logits(logits1, label_a, reduction='none'), dim=(1,2,3)) 187 | loss_1b = torch.mean(F.binary_cross_entropy_with_logits(logits1, label_b, reduction='none'), dim=(1,2,3)) 188 | loss_2a = torch.mean(F.binary_cross_entropy_with_logits(logits2, label_a, reduction='none'), dim=(1,2,3)) 189 | loss_2b = torch.mean(F.binary_cross_entropy_with_logits(logits2, label_b, reduction='none'), dim=(1,2,3)) 190 | 191 | loss1 = torch.where((loss_1a + loss_2b) < (loss_1b + loss_2a), loss_1a, loss_1b).mean() 192 | loss2 = torch.where((loss_1a + loss_2b) < (loss_1b + loss_2a), loss_2b, loss_2a).mean() 193 | else: 194 | loss1 = F.binary_cross_entropy_with_logits(logits1, label_a) 195 | loss2 = F.binary_cross_entropy_with_logits(logits2, label_b) 196 | 197 | return loss1, loss2 198 | 199 | def training_step(self, batch, batch_idx): 200 | opt1, opt2 = self.optimizers() 201 | depth_pre, depth_post, flow, pick_uv1, pick_uv2 = batch 202 | if self.input_mode == 'flowonly': 203 | x1 = flow 204 | else: 205 | raise NotImplementedError 206 | 207 | uv1, uv2, info = self.forward(x1, depth_pre, depth_post) 208 | loss1, loss2 = self.loss(info['logits1'], info['logits2'], pick_uv1, pick_uv2) 209 | 210 | opt1.zero_grad() 211 | self.manual_backward(loss1) 212 | opt1.step() 213 | opt2.zero_grad() 214 | self.manual_backward(loss2) 215 | opt2.step() 216 | 217 | if batch_idx == 0: 218 | self.plot(batch, uv1, uv2, info, stage='train') 219 | self.log_dict({"loss1/train": loss1, "loss2/train": loss2}, on_step=False, on_epoch=True, prog_bar=False) 220 | 221 | def validation_step(self, batch, batch_idx, log=True): 222 | depth_pre, depth_post, flow, pick_uv1, pick_uv2 = batch 223 | if self.input_mode == 'flowonly': 224 | x1 = flow 225 | else: 226 | raise NotImplementedError 227 | 228 | uv1, uv2, info = self.forward(x1, depth_pre, depth_post) 229 | loss1, loss2 = self.loss(info['logits1'], info['logits2'], pick_uv1, pick_uv2) 230 | 231 | if batch_idx == 0: 232 | self.plot(batch, uv1, uv2, info, stage='val') 233 | self.log_dict({"loss1/val": loss1, "loss2/val": loss2}, on_step=False, on_epoch=True, prog_bar=False) 234 | 235 | def configure_optimizers(self): 236 | if self.model_type == 'split': 237 | opt1 = torch.optim.Adam(self.first.parameters(), lr=self.lr) 238 | opt2 = torch.optim.Adam(self.second.parameters(), lr=self.lr) 239 | return opt1, opt2 240 | 241 | def plot(self, batch, pred_uv1, pred_uv2, info, stage): 242 | fig, ax = plt.subplots(2, 3, figsize=(16, 8)) 243 | 244 | # Row 1 245 | depth_pre, depth_post, flow, gt_uv1, gt_uv2 = batch 246 | depth_pre = depth_pre[0].squeeze().cpu().numpy() 247 | depth_post = depth_post[0].squeeze().cpu().numpy() 248 | flow = flow[0].squeeze().permute(1, 2, 0).cpu().numpy() 249 | gt_uv1 = gt_uv1[0].squeeze().cpu().numpy() 250 | gt_uv2 = gt_uv2[0].squeeze().cpu().numpy() 251 | ax[0][0].set_title("depth before") 252 | ax[0][0].imshow(depth_pre) 253 | ax[0][0].scatter(gt_uv1[0], gt_uv1[1], label='pick_uv1' if self.pick else 'place_uv1') 254 | ax[0][0].scatter(gt_uv2[0], gt_uv2[1], label='pick_uv2' if self.pick else 'place_uv2') 255 | ax[0][0].legend() 256 | ax[0][1].set_title("depth after") 257 | ax[0][1].imshow(depth_post) 258 | ax[0][2].set_title("flow") 259 | plot_flow(ax[0][2], flow, skip=0.05) 260 | 261 | # Row 2 262 | logits1 = info['logits1'][0].detach().squeeze().cpu().numpy() 263 | logits1 = cv2.resize(logits1, (200, 200)) 264 | pick1_gau = info['pick1_gau'][0].detach().cpu().squeeze().numpy() 265 | logits2 = info['logits2'][0].detach().cpu().squeeze().numpy() 266 | logits2 = cv2.resize(logits2, (200, 200)) 267 | ax[1][0].set_title("logits1") 268 | ax[1][0].imshow(logits1) 269 | ax[1][1].set_title("pick1_gaussian") 270 | ax[1][1].imshow(pick1_gau) 271 | ax[1][2].set_title("logits2") 272 | ax[1][2].imshow(logits2) 273 | plt.tight_layout() 274 | self.logger[1].experiment.add_figure(stage, fig, self.global_step) # tensorboard 275 | plt.close() 276 | -------------------------------------------------------------------------------- /fabricflownet/picknet/train.py: -------------------------------------------------------------------------------- 1 | from fabricflownet.picknet.models import FlowPickSplitModel 2 | from fabricflownet.picknet.dataset import PickNetDataset 3 | 4 | import os 5 | import sys 6 | import matplotlib.pyplot as plt 7 | import numpy as np 8 | from copy import deepcopy 9 | 10 | import torch 11 | from torch.utils.data import DataLoader 12 | from collections import namedtuple 13 | import time 14 | 15 | import hydra 16 | import pytorch_lightning.utilities.seed as seed_utils 17 | import pytorch_lightning as pl 18 | from pytorch_lightning import loggers as pl_loggers 19 | from pytorch_lightning.callbacks import ModelCheckpoint 20 | 21 | Experience = namedtuple('Experience', ('obs', 'goal', 'act', 'rew', 'nobs', 'done')) 22 | 23 | @hydra.main(config_name="config") 24 | def main(cfg): 25 | seed_utils.seed_everything(cfg.seed) 26 | with open('.hydra/command.txt', 'w') as f: 27 | f.write('python ' + ' '.join(sys.argv)) 28 | 29 | # Get training samples 30 | trainpath = f'{cfg.base_path}/{cfg.train_name}' 31 | trainfs = sorted(['_'.join(fn.split('_')[0:2]) 32 | for fn in os.listdir(f'{trainpath}/actions')]) 33 | if cfg.max_buf != None: 34 | trainfs = trainfs[:cfg.max_buf] 35 | print(f"Max training set: {len(trainfs)}") 36 | 37 | # Get validation samples 38 | valpath = f'{cfg.base_path}/{cfg.val_name}' 39 | valfs = sorted(['_'.join(fn.split('_')[0:2]) 40 | for fn in os.listdir(f'{valpath}/actions')]) 41 | if cfg.max_buf != None: 42 | valfs = valfs[:cfg.max_buf] 43 | print(f"Max val set: {len(valfs)}") 44 | 45 | # buf_train = torch.load(f"{cfg.base_path}/{cfg.train_name}/{cfg.train_name}_idx.buf") 46 | # if cfg.max_buf is not None: 47 | # buf_train = buf_train[:cfg.max_buf] 48 | 49 | # buf_test = torch.load(f"{cfg.base_path}/{cfg.val_name}/{cfg.val_name}_idx.buf") 50 | # if cfg.max_buf is not None: 51 | # buf_test = buf_test[:cfg.max_buf] 52 | 53 | # Get camera params 54 | train_camera_params = np.load(f"{cfg.base_path}/{cfg.train_name}/camera_params.npy", allow_pickle=True)[()] 55 | val_camera_params = np.load(f"{cfg.base_path}/{cfg.val_name}/camera_params.npy", allow_pickle=True)[()] 56 | np.testing.assert_equal(val_camera_params, train_camera_params) 57 | camera_params = train_camera_params 58 | 59 | train_data = PickNetDataset(camera_params, cfg, trainfs, mode='train', pick_pt=cfg.net_cfg.pick) 60 | val_data = PickNetDataset(camera_params, cfg, valfs, mode='test', pick_pt=cfg.net_cfg.pick) 61 | train_loader = DataLoader(train_data, batch_size=cfg.batch_size, shuffle=True, num_workers=cfg.workers, persistent_workers=cfg.workers>0) 62 | val_loader = DataLoader(val_data, batch_size=1, shuffle=False, num_workers=cfg.workers, persistent_workers=cfg.workers>0) 63 | 64 | # Init model 65 | if cfg.net_cfg.model_type == 'split': 66 | model = FlowPickSplitModel(**cfg.net_cfg) 67 | else: 68 | raise NotImplementedError 69 | 70 | csv_logger = pl_loggers.CSVLogger(save_dir=cfg.csv_log_dir) 71 | tb_logger = pl_loggers.TensorBoardLogger(save_dir=cfg.tboard_log_dir) 72 | chkpt_cb = ModelCheckpoint(monitor='loss1/val', save_last=True, save_top_k=-1, every_n_val_epochs=10) 73 | trainer = pl.Trainer(gpus=[0], 74 | logger=[csv_logger, tb_logger], 75 | max_epochs=cfg.epochs, 76 | check_val_every_n_epoch=cfg.check_val_every_n_epoch, # TODO change to every k steps 77 | log_every_n_steps=len(train_loader) if len(train_loader) < 50 else 50, 78 | callbacks=[chkpt_cb]) 79 | trainer.fit(model, train_loader, val_loader) 80 | 81 | if __name__ == '__main__': 82 | main() 83 | 84 | -------------------------------------------------------------------------------- /fabricflownet/utils.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | from softgym.envs.bimanual_env import uv_to_world_pos 3 | import numpy as np 4 | from copy import deepcopy 5 | import torch 6 | import torch.nn.functional as F 7 | import matplotlib.pyplot as plt 8 | import matplotlib as mpl 9 | 10 | def remove_occluded_knots(camera_params, knots, coords, depth, zthresh=0.001, debug_viz=False): 11 | if depth.shape[0] < camera_params['default_camera']['height']: 12 | print('Warning: resizing depth') 13 | depth = cv2.resize(depth, (camera_params['default_camera']['height'], camera_params['default_camera']['width'])) 14 | 15 | unoccluded_knots = [] 16 | occluded_knots = [] 17 | for i, uv in enumerate(knots): 18 | u_float, v_float = uv[0], uv[1] 19 | if np.isnan(u_float) or np.isnan(v_float): 20 | continue 21 | u, v = int(np.rint(u_float)), int(np.rint(v_float)) 22 | 23 | # check if pixel is outside of image bounds 24 | if u < 0 or v < 0 or u >= depth.shape[1] or v >= depth.shape[0]: 25 | knots[i] = [float('NaN'), float('NaN')] 26 | continue 27 | 28 | # Get depth into world coordinates 29 | d = depth[v, u] 30 | deproj_coords = uv_to_world_pos(camera_params, depth, u_float, v_float, particle_radius=0, on_table=False)[0:3] 31 | zdiff = deproj_coords[1] - coords[i][1] 32 | 33 | # Check is well projected xyz point 34 | if zdiff > zthresh: 35 | # invalidate u, v and continue 36 | occluded_knots.append(deepcopy(knots[i])) 37 | knots[i] = [float('NaN'), float('NaN')] 38 | continue 39 | 40 | unoccluded_knots.append(deepcopy(knots[i])) 41 | 42 | # Debug plotting 43 | if debug_viz: 44 | # 3D scatterplot 45 | # fig = plt.figure() 46 | # ax = fig.add_subplot(1, 1, 1, projection='3d') 47 | # for i, (u, v) in enumerate(knots[::3]): 48 | # c = 'r' if np.isnan(u) or np.isnan(v) else 'b' 49 | # ax.scatter(coords[i, 0], coords[i, 2], coords[i, 1], s=1, c=c) 50 | # plt.show() 51 | 52 | # 2D plot 53 | fig, ax = plt.subplots(1, 3, figsize=(8, 3)) 54 | ax[0].set_title('depth') 55 | ax[0].imshow(depth) 56 | ax[1].set_title('occluded points\nin red') 57 | ax[1].imshow(depth) 58 | if occluded_knots != []: 59 | occluded_knots = np.array(occluded_knots) 60 | ax[1].scatter(occluded_knots[:, 0], occluded_knots[:, 1], marker='.', s=1, c='r', alpha=0.4) 61 | ax[2].imshow(depth) 62 | ax[2].set_title('unoccluded points\nin blue') 63 | unoccluded_knots = np.array(unoccluded_knots) 64 | ax[2].scatter(unoccluded_knots[:, 0], unoccluded_knots[:, 1], marker='.', s=1, alpha=0.4) 65 | plt.show() 66 | 67 | return knots 68 | 69 | def get_harris(mask, thresh=0.2): 70 | """Harris corner detector 71 | Params 72 | ------ 73 | - mask: np.float32 image of 0.0 and 1.0 74 | - thresh: threshold for filtering small harris values Returns 75 | ------- 76 | - harris: np.float32 array of 77 | """ 78 | # Params for cornerHarris: 79 | # mask - Input image, it should be grayscale and float32 type. 80 | # blockSize - It is the size of neighbourhood considered for corner detection 81 | # ksize - Aperture parameter of Sobel derivative used. 82 | # k - Harris detector free parameter in the equation. 83 | # https://docs.opencv.org/master/dd/d1a/group__imgproc__feature.html#gac1fc3598018010880e370e2f709b4345 84 | harris = cv2.cornerHarris(mask, blockSize=5, ksize=5, k=0.01) 85 | harris[harris