├── .gitignore ├── README.md ├── create_binary_masks.py ├── dockerfiles ├── Dockerfile └── Dockerfile2 ├── mask_df.csv ├── metadata.csv ├── requirements.txt ├── satellites.jpg ├── satellites_eda.ipynb ├── satellites_eda_d.ipynb ├── scripts ├── cnn_inference_env.sh ├── create_8bit_test_images.py ├── create_binary_masks.py ├── download.sh ├── env.sh └── install-dependencies.sh ├── sknw.py └── src ├── .ipynb_checkpoints ├── experiments-checkpoint.ipynb ├── pipeline_experiments-checkpoint.ipynb ├── play_w_stuff-checkpoint.ipynb └── readme.md-checkpoint.ipynb ├── DilatedResnet.py ├── InceptionResnetv2.py ├── LRScheduler.py ├── LinkNet.py ├── Loss.py ├── MaskUtils.py ├── ResNeXt.py ├── SatellitesAugs.py ├── SatellitesDataset.py ├── TbLogger.py ├── UNet.py ├── __pycache__ ├── DilatedResnet.cpython-35.pyc ├── InceptionResnetv2.cpython-35.pyc ├── LinkNet.cpython-35.pyc ├── Loss.cpython-35.pyc ├── LossSemSeg.cpython-35.pyc ├── MaskUtils.cpython-35.pyc ├── ResNeXt.cpython-35.pyc ├── SatellitesAugs.cpython-35.pyc ├── SatellitesDataset.cpython-35.pyc ├── TbLogger.cpython-35.pyc ├── UNet.cpython-35.pyc ├── presets.cpython-35.pyc └── sknw.cpython-35.pyc ├── bit8_test.csv ├── create_8bit_test_images.py ├── experiments.ipynb ├── final_model_lstrs.py ├── geojson_df.csv ├── geojson_df_full.csv ├── new_masks.csv ├── new_masks_layered.csv ├── pipeline_experiments.ipynb ├── play_w_stuff.ipynb ├── presets.py ├── readme.md.ipynb ├── resnext_features ├── __pycache__ │ ├── resnext101_32x4d_features.cpython-35.pyc │ └── resnext101_64x4d_features.cpython-35.pyc ├── resnext101_32x4d_features.py └── resnext101_64x4d_features.py ├── script.sh ├── sknw.py ├── test.jpg ├── train_gapnet.py ├── train_satellites.py ├── train_satellites_folds.py └── untitled.txt /.gitignore: -------------------------------------------------------------------------------- 1 | # Weights, logs and other folders # 2 | ###################### 3 | src/weights/* 4 | src/logs/* 5 | src/tb_logs/* 6 | data/* 7 | apls/* 8 | .ipynb_checkpoints/* 9 | __pycache__/* 10 | 11 | # Logs and databases # 12 | ###################### 13 | *.log 14 | *.npy 15 | 16 | # OS generated files # 17 | ###################### 18 | .DS_Store 19 | .DS_Store? 20 | ._* 21 | .Spotlight-V100 22 | .Trashes 23 | ehthumbs.db 24 | Thumbs.db 25 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ![Architecture](satellites.jpg) 2 | 3 | **More stuff from us** 4 | - [Telegram](https://t.me/snakers4) 5 | - [Twitter](https://twitter.com/AlexanderVeysov) 6 | - [Blog](https://spark-in.me/tag/data-science) 7 | 8 | # 1 Hardware requirements 9 | 10 | **Training** 11 | 12 | - 6+ core modern CPU (Xeon, i7) for fast image pre-processing; 13 | - The models were trained on 2 * GeForce 1080 Ti; 14 | - Training time on my setup ~ **3 hours** for models with 8-bit images as inputs; 15 | - Disk space - 40GB should be more than enough; 16 | 17 | **Inference** 18 | 19 | - 6+ core modern CPU (Xeon, i7) for fast image pre-processing; 20 | - On 2 * GeForce 1080 Ti inference takes **3-5 minutes**; 21 | - Graph creation takes **5-10 minutes**; 22 | 23 | # 2 Preparing and launching the Docker environment 24 | 25 | **Clone the repository** 26 | 27 | `git clone https://github.com/snakers4/spacenet-three .` 28 | 29 | 30 | **This repository contains 2 Dockerfiles** 31 | - `/dockerfiles/Dockerfile` - this is the main Dockerfile which was used as environment to run the training and inference scripts. This one also includes the majority of the tested legact spacenet libraries used by their APLS repository 32 | - `/dockerfiles/Dockerfile2`- this is an additional backup Docker file with newer versions of the drivers and PyTorch, just in case. Does not contain spacenet libraries 33 | 34 | **Build a Docker image** 35 | 36 | ` 37 | cd dockerfiles 38 | docker build -t aveysov . 39 | ` 40 | 41 | **Install the latest nvidia docker** 42 | 43 | Follow instructions from [here](https://github.com/NVIDIA/nvidia-docker). 44 | Please prefer nvidia-docker2 for more stable performance. 45 | 46 | 47 | To test all works fine run: 48 | 49 | 50 | `docker run --runtime=nvidia --rm nvidia/cuda nvidia-smi` 51 | 52 | **(IMPORTANT) Run docker container (IMPORTANT)** 53 | 54 | Unless you use this exact command (with --shm-size flag) (you can change ports and mounted volumes, of course), then the PyTorch generators **WILL NOT WORK**. 55 | 56 | 57 | - nvidia-docker 2: `docker run --runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=all -it -v /path/to/cloned/repository:/home/keras/notebook -p 8888:8888 -p 6006:6006 --shm-size 8G aveysov` 58 | - nvidia-docker: `nvidia-docker -it -v /path/to/cloned/repository:/home/keras/notebook -p 8888:8888 -p 6006:6006 --shm-size 8G aveysov` 59 | 60 | **Installing project specific software** 61 | 62 | 63 | 0. Already done for `/dockerfiles/Dockerfile`. You can do it yourself for `/dockerfiles/Dockerfile2` 64 | 1. Exec into the docker machine via `docker exec -it --user root YOUR_CONTAINER_ID /bin/bash` 65 | 2. Run these scripts one after another 66 | ``` 67 | conda install -y -c conda-forge cython 68 | conda install -y -c conda-forge rasterio 69 | conda install -y -c conda-forge libgdal 70 | conda install -y -c conda-forge gdal 71 | conda install -y -c conda-forge scikit-image 72 | conda install -y -c conda-forge pyproj 73 | conda install -y -c conda-forge geopandas 74 | conda install -y -c conda-forge tqdm 75 | conda install -y -c conda-forge shapely=1.5.16 76 | conda install -y -c conda-forge scipy 77 | conda install -y -c conda-forge networkx=1.11 78 | conda install -y -c conda-forge fiona 79 | pip3 install utm 80 | pip3 install osmnx==0.5.1 81 | ``` 82 | 3. Run these scripts one after another 83 | ``` 84 | pip3 install numba 85 | conda install -y -c conda-forge scikit-image 86 | ``` 87 | 4. Sometimes `apt install libgl1-mesa-glx` is necessary due to some Docker [quirks](https://github.com/ContinuumIO/docker-images/issues/49) 88 | 89 | Steps 2-3 are to ensure compatibility with legacy software from APLS [repository](https://github.com/CosmiQ/apls). 90 | An alternative to that - is to use pip's requirements.txt in the same order. 91 | These steps are required to run 8-bit mask creation step from APLS repository and mask creation step. 92 | If you will be trying to re-do this step - reserve 5-6 hours for experiments. 93 | 94 | **To start the stopped container** 95 | 96 | 97 | `docker start -i YOUR_CONTAINER_ID` 98 | 99 | 100 | # 3 Preparing the data and the machine for running scripts using APLS code 101 | 102 | - Dowload the data into `data/` 103 | - Data Download [guide](https://drive.google.com/open?id=1WJh8Q1Oj38Ahn0FiwUv1WZXoIwnY08_kL4XO8cikPo0) from the authors 104 | - Run these commands to create 8-bit images, mask and test 8-bit images: 105 | ``` 106 | docker exec -it YOUR_CONTAINER_ID sh -c "python3 create_binary_masks.py && \ 107 | cd src && \ 108 | python3 create_8bit_test_images.py && \ 109 | " 110 | ``` 111 | 112 | After all of your manipulations your directory should look like: 113 | 114 | ``` 115 | ├── README.md <- The top-level README for developers using this project. 116 | ├── data 117 | │ ├── AOI_2_Vegas_Roads_Test_Public <- Test set for each city 118 | │ └── AOI_2_Vegas_Roads_Train <- Train set for each city 119 | │ ├─ geojson 120 | │ ├─ summaryData 121 | │ ├─ MUL 122 | │ ├─ RGB-PanSharpen 123 | │ ├─ PAN 124 | │ ├─ MUL-PanSharpen 125 | │ ├─ MUL-PanSharpen_8bit 126 | │ ├─ RGB-PanSharpen_8bit 127 | │ ├─ PAN_8bit 128 | │ ├─ MUL_8bit 129 | │ ├─ MUL-PanSharpen_mask 130 | │ ├─ RGB-PanSharpen_mask 131 | │ ├─ PAN_mask 132 | │ ├─ MUL_mask 133 | │ └─ RGB-PanSharpen_mask 134 | │ │ 135 | │ ... 136 | │ │ 137 | │ ├── AOI_5_Khartoum_Roads_Test_Public <- Test set for each city 138 | │ └── AOI_5_Khartoum_Roads_Train <- Train set for each city 139 | │ 140 | ├── dockerfiles <- A folder with Dockerfiles 141 | │ 142 | ├── src <- Source code 143 | │ 144 | └── scripts <- One-off preparation scripts 145 | ``` 146 | 147 | # 4 Training the model from scratch 148 | 149 | If all is ok, then use the following command to train the model 150 | 151 | - optional - turn on tensorboard for monitoring progress `tensorboard --logdir='satellites_roads/src/tb_logs' --port=6006` via jupyter notebook console or via tmux + docker exec (model converges in 30-40 epochs) 152 | - then 153 | ``` 154 | docker exec -it YOUR_CONTAINER_ID sh -c "echo 'python3 train_satellites.py \ 155 | --arch linknet34 --batch-size 6 \ 156 | --imsize 1280 --preset mul_ps_vegetation --augs True \ 157 | --workers 6 --epochs 40 --start-epoch 0 \ 158 | --seed 42 --print-freq 20 \ 159 | --lr 1e-3 --optimizer adam \ 160 | --tensorboard True --lognumber ln34_mul_ps_vegetation_aug_dice' > train.sh && \ 161 | sh train.sh"" 162 | ``` 163 | 164 | # 5 Predicting masks 165 | 166 | - You can load the [pretrained weights](https://drive.google.com/open?id=19Zd4RG0P0YwwRZ_xgip7dAmb4Bp8WqTI) to the `src/weights` folder for the above definition of the model 167 | - Run this script 168 | ``` 169 | docker exec -it YOUR_CONTAINER_ID sh -c "cd src && \ 170 | echo 'python3 train_satellites.py \ 171 | --arch linknet34 --batch-size 12 \ 172 | --imsize 1312 --preset mul_ps_vegetation --augs True \ 173 | --workers 6 --epochs 50 --start-epoch 0 \ 174 | --seed 42 --print-freq 10 \ 175 | --lr 1e-3 --optimizer adam \ 176 | --lognumber norm_ln34_mul_ps_vegetation_aug_dice_predict \ 177 | --predict --resume weights/norm_ln34_mul_ps_vegetation_aug_dice_best.pth.tar\' > predict.sh && \ 178 | sh predict.sh 179 | ``` 180 | 181 | 182 | 183 | # 6 Creating graphs and submission files 184 | 185 | 186 | Execute this script: 187 | ``` 188 | docker exec -it YOUR_CONTAINER_ID sh -c "cd src && \ 189 | python3 final_model_lstrs.py --folder norm_ln34_mul_ps_vegetation_aug_dice_predict" 190 | ``` 191 | 192 | 193 | `folder` argument is for masks containing folder name, default is `norm_ln34_mul_ps_vegetation_aug_dice_predict`. 194 | Scipt saves a file called `norm_test.csv` into `../solutions` directory. 195 | 196 | The resulting file is used then as a submission file. 197 | 198 | 199 | 200 | # 7 Additional notes 201 | 202 | - You can run training and inference on the presets from `/src/presets.py`; 203 | - So the model can be evaluated on RGB-PS images and / or 8-channel images as well; 204 | - This script, for example will train an 8-channel model: 205 | ``` 206 | python3 train_satellites.py \ 207 | --arch linknet34 --batch-size 6 \ 208 | --imsize 1280 --preset mul_ps_vegetation --augs True \ 209 | --workers 6 --epochs 40 --start-epoch 0 \ 210 | --seed 42 --print-freq 20 \ 211 | --lr 1e-3 --optimizer adam \ 212 | --tensorboard True --lognumber ln34_mul_ps_vegetation_aug_dice 213 | ``` 214 | 215 | To train an 8-channel model you should also replace mean and std settings in the `src/SatellitesAugs.py` 216 | 217 | ``` 218 | # 8-channel settings 219 | # normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5], 220 | # std=[1, 1, 1, 1, 1, 1, 1, 1]) 221 | 222 | # 3 channel settings 223 | normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], 224 | std=[0.229, 0.224, 0.225]) 225 | ``` 226 | 227 | - 16-bit images are also supported: 228 | 229 | This snippet is commented in `src/SatellitesAugs.py` 230 | 231 | ``` 232 | # version compatible with 16-bit images 233 | """ 234 | class ImgAugAugs(object): 235 | def __call__(self, 236 | image): 237 | global seed 238 | 239 | # poor man's flipping 240 | if seed%2==0: 241 | image = np.fliplr(image) 242 | elif seed%4==0: 243 | image = np.fliplr(image) 244 | image = np.flipud(image) 245 | 246 | # poor man's affine transformations 247 | image = rotate(image, 248 | angle=seed, 249 | resize=False, 250 | clip=True, 251 | preserve_range=True) 252 | 253 | return image 254 | """ 255 | 256 | But be careful with normalization (`div(255)`), as 16-bit images actually are 11-bit (divide by 2048). 257 | 258 | ``` 259 | 260 | - Also the following models are supported 261 | - `unet11` (VGG11 + Unet) 262 | - `linknet50` (ResNet50 + LinkNet, 3 layers) 263 | - `linknet50_full` (ResNet50 + LinkNet, 4 layers) 264 | - `linknext` (ResNext-101-32 + LinkNet, 4 layers) 265 | 266 | - Also in the repo you can find scripts to generate wide masks (i.e. wide roads have varying width) and layered masks (paved / non-paved). There are scripts in the `src/SatellitesDataset.py` that support that. They basically just replace some paths; 267 | 268 | # 8 Juputer notebooks 269 | 270 | Use these notebooks on your own risk! 271 | 272 | - `src/experiments.ipynb` - general debugging notebook with new models / generators / etc 273 | - `src/play_w_stuff.ipynb` - visualizing the solutions 274 | - `src/pipeline_experiments.ipynb`- some minor experiments with the graph creation script 275 | -------------------------------------------------------------------------------- /create_binary_masks.py: -------------------------------------------------------------------------------- 1 | import pandas as pd 2 | from multiprocessing import Pool 3 | import tqdm 4 | import numpy as np 5 | import os 6 | import glob as glob 7 | from skimage.io import imread, imsave 8 | import osmnx as ox 9 | import numpy as np 10 | import matplotlib 11 | import matplotlib.pyplot as plt 12 | import geopandas as gpd 13 | from osgeo import gdal, ogr, osr 14 | import cv2 15 | import subprocess 16 | import shapely 17 | from shapely.geometry import MultiLineString 18 | from matplotlib.patches import PathPatch 19 | import matplotlib.path 20 | 21 | imgs = [] 22 | 23 | # change this to your data prefix 24 | path_prefix = 'data' 25 | 26 | # default variables from the hosts of the challenge 27 | buffer_meters = 2 28 | burnValue = 150 29 | 30 | # only train folders 31 | folders = ['AOI_2_Vegas_Roads_Train', 32 | 'AOI_5_Khartoum_Roads_Train', 33 | 'AOI_3_Paris_Roads_Train', 34 | 'AOI_4_Shanghai_Roads_Train'] 35 | 36 | # image types 37 | prefix_dict = { 38 | 'mul': 'MUL', 39 | 'muls': 'MUL-PanSharpen', 40 | 'pan': 'PAN', 41 | 'rgbps': 'RGB-PanSharpen', 42 | } 43 | 44 | for folder in folders: 45 | for prefix in prefix_dict.items(): 46 | g = glob.glob(path_prefix+'/{}/{}/*.tif'.format(folder,prefix[1])) 47 | imgs.extend(g) 48 | 49 | img_folders = [(img.split('/')[1]) for img in imgs] 50 | img_subfolders = [(img.split('/')[2]) for img in imgs] 51 | img_files = [(img.split('/')[3]) for img in imgs] 52 | 53 | def create_binary_mask(input_data): 54 | img_path = input_data[0] 55 | img_folder = input_data[1] 56 | img_subfolder = input_data[2] 57 | img_file = input_data[3] 58 | 59 | # create paths for masks and 8bit images 60 | label_file = os.path.join(path_prefix,img_folder,'geojson/spacenetroads','spacenetroads_AOI'+img_file.split('AOI')[1][0:-3]+'geojson') 61 | bit8_folder = os.path.join(path_prefix,img_folder,img_subfolder+'_8bit') 62 | bit8_path = os.path.join(bit8_folder,img_file) 63 | mask_folder = os.path.join(path_prefix,img_folder,img_subfolder+'_mask') 64 | mask_path = os.path.join(mask_folder,img_file[:-3])+'png' 65 | # vis_folder = os.path.join(path_prefix,img_folder,img_subfolder+'_vis') 66 | # vis_path = os.path.join(vis_folder,img_file[:-3])+'png' 67 | 68 | # print(label_file) 69 | # create the necessary folders and remove the existing files 70 | 71 | if not os.path.exists(bit8_folder): 72 | os.mkdir(bit8_folder) 73 | if not os.path.exists(mask_folder): 74 | os.mkdir(mask_folder) 75 | # if not os.path.exists(vis_folder): 76 | # os.mkdir(vis_folder) 77 | # if os.path.isfile(vis_path): 78 | # os.remove(vis_path) 79 | if os.path.isfile(bit8_path): 80 | os.remove(bit8_path) 81 | if os.path.isfile(mask_path): 82 | os.remove(mask_path) 83 | if os.path.isfile(mask_path[:-3]+'jpg'): 84 | os.remove(mask_path[:-3]+'jpg') 85 | 86 | try: 87 | # convert images to 8-bit 88 | convert_to_8Bit(img_path, 89 | bit8_path, 90 | outputPixType='Byte', 91 | outputFormat='GTiff', 92 | rescale_type='rescale', 93 | percentiles=[2,98]) 94 | 95 | # create masks 96 | # note that though the output raster file has .png extension 97 | # in reality I delete this file and save only jpg version later 98 | 99 | mask, gdf_buffer = get_road_buffer(geoJson = label_file, 100 | im_vis_file = bit8_path, 101 | output_raster = mask_path, 102 | buffer_meters= buffer_meters, 103 | burnValue= burnValue, 104 | bufferRoundness=6, 105 | plot_file='', # this indicates that no visualization plot is required 106 | figsize= (6,6), 107 | fontsize=8, 108 | dpi=200, 109 | show_plot=False, 110 | verbose=False) 111 | 112 | # read the png file, save it as jpeg and 113 | mask = imread(mask_path) 114 | imsave(fname=mask_path[:-3]+'jpg',arr = mask) 115 | mask_max = np.max(mask) 116 | del mask 117 | # remove the png file, but keep the 8-bit mask 118 | os.remove(mask_path) 119 | except BaseException as e: 120 | print(str(e)) 121 | mask_max = -1 122 | 123 | return [label_file,bit8_folder,bit8_path,mask_folder,mask_path[:-3]+'jpg',img_path,img_folder,img_subfolder,img_file,mask_max] 124 | 125 | def get_road_buffer(geoJson, im_vis_file, output_raster, 126 | buffer_meters=2, burnValue=1, 127 | bufferRoundness=6, 128 | plot_file='', figsize=(6,6), fontsize=6, 129 | dpi=800, show_plot=False, 130 | verbose=False): 131 | ''' 132 | Get buffer around roads defined by geojson and image files. 133 | Calls create_buffer_geopandas() and gdf_to_array(). 134 | Assumes in_vis_file is an 8-bit RGB file. 135 | Returns geodataframe and ouptut mask. 136 | ''' 137 | 138 | gdf_buffer = create_buffer_geopandas(geoJson, 139 | bufferDistanceMeters=buffer_meters, 140 | bufferRoundness=bufferRoundness, 141 | projectToUTM=True) 142 | 143 | 144 | # create label image 145 | if len(gdf_buffer) == 0: 146 | mask_gray = np.zeros(cv2.imread(im_vis_file,0).shape) 147 | cv2.imwrite(output_raster, mask_gray) 148 | else: 149 | gdf_to_array(gdf_buffer, im_vis_file, output_raster, 150 | burnValue=burnValue) 151 | 152 | # load mask 153 | mask_gray = cv2.imread(output_raster, 0) 154 | 155 | # make plots 156 | if plot_file: 157 | # plot all in a line 158 | if (figsize[0] != figsize[1]): 159 | fig, (ax0, ax1, ax2, ax3) = plt.subplots(1,4, figsize=figsize)#(13,4)) 160 | # else, plot a 2 x 2 grid 161 | else: 162 | fig, ((ax0, ax1), (ax2, ax3)) = plt.subplots(2,2, figsize=figsize) 163 | 164 | # road lines 165 | try: 166 | gdfRoadLines = gpd.read_file(geoJson) 167 | gdfRoadLines.plot(ax=ax0, marker='o', color='red') 168 | except: 169 | ax0.imshow(mask_gray) 170 | ax0.axis('off') 171 | ax0.set_aspect('equal') 172 | ax0.set_title('Roads from GeoJson', fontsize=fontsize) 173 | 174 | # first show raw image 175 | im_vis = cv2.imread(im_vis_file, 1) 176 | img_mpl = cv2.cvtColor(im_vis, cv2.COLOR_BGR2RGB) 177 | ax1.imshow(img_mpl) 178 | ax1.axis('off') 179 | ax1.set_title('8-bit RGB Image', fontsize=fontsize) 180 | 181 | # plot mask 182 | ax2.imshow(mask_gray) 183 | ax2.axis('off') 184 | ax2.set_title('Roads Mask (' + str(np.round(buffer_meters)) \ 185 | + ' meter buffer)', fontsize=fontsize) 186 | 187 | # plot combined 188 | ax3.imshow(img_mpl) 189 | # overlay mask 190 | # set zeros to nan 191 | z = mask_gray.astype(float) 192 | z[z==0] = np.nan 193 | # change palette to orange 194 | palette = plt.cm.gray 195 | #palette.set_over('yellow', 0.9) 196 | palette.set_over('lime', 0.9) 197 | ax3.imshow(z, cmap=palette, alpha=0.66, 198 | norm=matplotlib.colors.Normalize(vmin=0.5, vmax=0.9, clip=False)) 199 | ax3.set_title('8-bit RGB Image + Buffered Roads', fontsize=fontsize) 200 | ax3.axis('off') 201 | 202 | #plt.axes().set_aspect('equal', 'datalim') 203 | 204 | plt.tight_layout() 205 | plt.savefig(plot_file, dpi=dpi) 206 | if not show_plot: 207 | plt.close() 208 | 209 | return mask_gray, gdf_buffer 210 | 211 | def create_buffer_geopandas(geoJsonFileName, 212 | bufferDistanceMeters=2, 213 | bufferRoundness=1, 214 | projectToUTM=True): 215 | ''' 216 | Create a buffer around the lines of the geojson. 217 | Return a geodataframe. 218 | ''' 219 | 220 | inGDF = gpd.read_file(geoJsonFileName) 221 | 222 | # set a few columns that we will need later 223 | inGDF['type'] = inGDF['road_type'].values 224 | inGDF['class'] = 'highway' 225 | inGDF['highway'] = 'highway' 226 | 227 | if len(inGDF) == 0: 228 | return [], [] 229 | 230 | # Transform gdf Roadlines into UTM so that Buffer makes sense 231 | if projectToUTM: 232 | tmpGDF = ox.project_gdf(inGDF) 233 | else: 234 | tmpGDF = inGDF 235 | 236 | gdf_utm_buffer = tmpGDF 237 | 238 | # perform Buffer to produce polygons from Line Segments 239 | gdf_utm_buffer['geometry'] = tmpGDF.buffer(bufferDistanceMeters, 240 | bufferRoundness) 241 | 242 | gdf_utm_dissolve = gdf_utm_buffer.dissolve(by='class') 243 | gdf_utm_dissolve.crs = gdf_utm_buffer.crs 244 | 245 | if projectToUTM: 246 | gdf_buffer = gdf_utm_dissolve.to_crs(inGDF.crs) 247 | else: 248 | gdf_buffer = gdf_utm_dissolve 249 | 250 | return gdf_buffer 251 | 252 | def gdf_to_array(gdf, im_file, output_raster, burnValue=150): 253 | 254 | ''' 255 | Turn geodataframe to array, save as image file with non-null pixels 256 | set to burnValue 257 | ''' 258 | 259 | NoData_value = 0 # -9999 260 | 261 | gdata = gdal.Open(im_file) 262 | 263 | # set target info 264 | target_ds = gdal.GetDriverByName('GTiff').Create(output_raster, 265 | gdata.RasterXSize, 266 | gdata.RasterYSize, 1, gdal.GDT_Byte) 267 | target_ds.SetGeoTransform(gdata.GetGeoTransform()) 268 | 269 | # set raster info 270 | raster_srs = osr.SpatialReference() 271 | raster_srs.ImportFromWkt(gdata.GetProjectionRef()) 272 | target_ds.SetProjection(raster_srs.ExportToWkt()) 273 | 274 | band = target_ds.GetRasterBand(1) 275 | band.SetNoDataValue(NoData_value) 276 | 277 | outdriver=ogr.GetDriverByName('MEMORY') 278 | outDataSource=outdriver.CreateDataSource('memData') 279 | tmp=outdriver.Open('memData',1) 280 | outLayer = outDataSource.CreateLayer("states_extent", raster_srs, 281 | geom_type=ogr.wkbMultiPolygon) 282 | # burn 283 | burnField = "burn" 284 | idField = ogr.FieldDefn(burnField, ogr.OFTInteger) 285 | outLayer.CreateField(idField) 286 | featureDefn = outLayer.GetLayerDefn() 287 | for geomShape in gdf['geometry'].values: 288 | 289 | outFeature = ogr.Feature(featureDefn) 290 | outFeature.SetGeometry(ogr.CreateGeometryFromWkt(geomShape.wkt)) 291 | outFeature.SetField(burnField, burnValue) 292 | outLayer.CreateFeature(outFeature) 293 | outFeature = 0 294 | 295 | gdal.RasterizeLayer(target_ds, [1], outLayer, burn_values=[burnValue]) 296 | outLayer = 0 297 | outDatSource = 0 298 | tmp = 0 299 | 300 | return 301 | 302 | def convert_to_8Bit(inputRaster, outputRaster, 303 | outputPixType='Byte', 304 | outputFormat='GTiff', 305 | rescale_type='rescale', 306 | percentiles=[2, 98]): 307 | ''' 308 | Convert 16bit image to 8bit 309 | rescale_type = [clip, rescale] 310 | if clip, scaling is done strictly between 0 65535 311 | if rescale, each band is rescaled to a min and max 312 | set by percentiles 313 | ''' 314 | 315 | srcRaster = gdal.Open(inputRaster) 316 | cmd = ['gdal_translate', '-ot', outputPixType, '-of', 317 | outputFormat] 318 | 319 | # iterate through bands 320 | for bandId in range(srcRaster.RasterCount): 321 | bandId = bandId+1 322 | band = srcRaster.GetRasterBand(bandId) 323 | if rescale_type == 'rescale': 324 | bmin = band.GetMinimum() 325 | bmax = band.GetMaximum() 326 | # if not exist minimum and maximum values 327 | if bmin is None or bmax is None: 328 | (bmin, bmax) = band.ComputeRasterMinMax(1) 329 | # else, rescale 330 | band_arr_tmp = band.ReadAsArray() 331 | bmin = np.percentile(band_arr_tmp.flatten(), 332 | percentiles[0]) 333 | bmax= np.percentile(band_arr_tmp.flatten(), 334 | percentiles[1]) 335 | 336 | else: 337 | bmin, bmax = 0, 65535 338 | 339 | cmd.append('-scale_{}'.format(bandId)) 340 | cmd.append('{}'.format(bmin)) 341 | cmd.append('{}'.format(bmax)) 342 | cmd.append('{}'.format(0)) 343 | cmd.append('{}'.format(255)) 344 | 345 | cmd.append(inputRaster) 346 | cmd.append(outputRaster) 347 | # print("Conversin command:", cmd) 348 | subprocess.call(cmd) 349 | 350 | return 351 | 352 | 353 | input_data = zip(imgs,img_folders,img_subfolders,img_files) 354 | input_data = [item for item in input_data] 355 | 356 | with Pool(10) as p: 357 | mask_data = list(tqdm.tqdm(p.imap(create_binary_mask, input_data), 358 | total=len(input_data))) 359 | # transpose the list 360 | mask_data = list(map(list, zip(*mask_data))) 361 | 362 | mask_df = pd.DataFrame() 363 | 364 | for i,key in enumerate(['label_file','bit8_folder','bit8_path','mask_folder','mask_path','img_path','img_folder','img_subfolder','img_file', 'mask_max']): 365 | mask_df[key] = mask_data[i] 366 | 367 | mask_df.to_csv('mask_df_run.csv') 368 | -------------------------------------------------------------------------------- /dockerfiles/Dockerfile: -------------------------------------------------------------------------------- 1 | # add 7z tar and zip archivers 2 | FROM nvidia/cuda:8.0-cudnn6-devel 3 | 4 | RUN apt-get update && apt-get install -y openssh-server 5 | RUN mkdir /var/run/sshd 6 | RUN echo 'root:Ubuntu@41' | chpasswd 7 | RUN sed -i 's/PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config 8 | RUN apt -y install libgl1-mesa-glx 9 | 10 | # SSH login fix. Otherwise user is kicked off after login 11 | RUN sed 's@session\s*required\s*pam_loginuid.so@session optional pam_loginuid.so@g' -i /etc/pam.d/sshd 12 | 13 | ENV NOTVISIBLE "in users profile" 14 | RUN echo "export VISIBLE=now" >> /etc/profile 15 | 16 | ENV CONDA_DIR /opt/conda 17 | ENV PATH $CONDA_DIR/bin:$PATH 18 | 19 | # writing env variables to /etc/profile as mentioned here https://docs.docker.com/engine/examples/running_ssh_service/#run-a-test_sshd-container 20 | RUN echo "export CONDA_DIR=/opt/conda" >> /etc/profile 21 | RUN echo "export PATH=$CONDA_DIR/bin:$PATH" >> /etc/profile 22 | 23 | RUN mkdir -p $CONDA_DIR && \ 24 | echo export PATH=$CONDA_DIR/bin:'$PATH' > /etc/profile.d/conda.sh && \ 25 | apt-get update && \ 26 | apt-get install -y wget git libhdf5-dev g++ graphviz openmpi-bin nano && \ 27 | wget --quiet https://repo.continuum.io/miniconda/Miniconda3-4.2.12-Linux-x86_64.sh && \ 28 | echo "c59b3dd3cad550ac7596e0d599b91e75d88826db132e4146030ef471bb434e9a *Miniconda3-4.2.12-Linux-x86_64.sh" | sha256sum -c - && \ 29 | /bin/bash /Miniconda3-4.2.12-Linux-x86_64.sh -f -b -p $CONDA_DIR && \ 30 | ln /usr/lib/x86_64-linux-gnu/libcudnn.so /usr/local/cuda/lib64/libcudnn.so && \ 31 | ln /usr/lib/x86_64-linux-gnu/libcudnn.so.6 /usr/local/cuda/lib64/libcudnn.so.6 && \ 32 | ln /usr/include/cudnn.h /usr/local/cuda/include/cudnn.h && \ 33 | rm Miniconda3-4.2.12-Linux-x86_64.sh 34 | 35 | ENV NB_USER keras 36 | ENV NB_UID 1000 37 | 38 | RUN echo "export NB_USER=keras" >> /etc/profile 39 | RUN echo "export NB_UID=1000" >> /etc/profile 40 | 41 | RUN echo "export LD_LIBRARY_PATH=/usr/local/cuda/lib64:/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH" >> /etc/profile 42 | RUN echo "export CPATH=/usr/include:/usr/include/x86_64-linux-gnu:/usr/local/cuda/include:$CPATH" >> /etc/profile 43 | RUN echo "export LIBRARY_PATH=/usr/local/cuda/lib64:/lib/x86_64-linux-gnu:$LIBRARY_PATH" >> /etc/profile 44 | RUN echo "export CUDA_HOME=/usr/local/cuda" >> /etc/profile 45 | RUN echo "export CPLUS_INCLUDE_PATH=$CPATH" >> /etc/profile 46 | RUN echo "export KERAS_BACKEND=tensorflow" >> /etc/profile 47 | 48 | RUN useradd -m -s /bin/bash -N -u $NB_UID $NB_USER && \ 49 | mkdir -p $CONDA_DIR && \ 50 | chown keras $CONDA_DIR -R 51 | 52 | USER keras 53 | 54 | RUN mkdir -p /home/keras/notebook 55 | 56 | # Python 57 | ARG python_version=3.5 58 | 59 | RUN conda install -y python=${python_version} && \ 60 | pip install --upgrade pip && \ 61 | pip install tensorflow-gpu==1.4 && \ 62 | conda install Pillow scikit-learn notebook pandas matplotlib mkl nose pyyaml six h5py && \ 63 | conda install theano pygpu bcolz && \ 64 | pip install keras kaggle-cli lxml opencv-python requests scipy tqdm visdom && \ 65 | conda install pytorch torchvision cuda80 -c soumith && \ 66 | pip install imgaug && \ 67 | conda clean -yt 68 | 69 | RUN pip install git+https://github.com/ipython-contrib/jupyter_contrib_nbextensions && \ 70 | jupyter contrib nbextension install --user 71 | 72 | RUN conda install -y -c conda-forge cython 73 | RUN conda install -y -c conda-forge libgdal 74 | RUN conda install -y -c conda-forge gdal 75 | RUN conda install -y -c conda-forge scikit-image 76 | RUN conda install -y -c conda-forge pyproj 77 | RUN conda install -y -c conda-forge geopandas 78 | RUN conda install -y -c conda-forge tqdm 79 | RUN conda install -y -c conda-forge shapely=1.5.16 80 | RUN conda install -y -c conda-forge scipy 81 | RUN conda install -y -c conda-forge networkx=1.11 82 | RUN conda install -y -c conda-forge fiona 83 | RUN pip install utm 84 | RUN pip install osmnx==0.5.1 85 | RUN pip install numba 86 | RUN conda install -y -c conda-forge scikit-image 87 | 88 | ENV LD_LIBRARY_PATH /usr/local/cuda/lib64:/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH 89 | ENV CPATH /usr/include:/usr/include/x86_64-linux-gnu:/usr/local/cuda/include:$CPATH 90 | ENV LIBRARY_PATH /usr/local/cuda/lib64:/lib/x86_64-linux-gnu:$LIBRARY_PATH 91 | ENV CUDA_HOME /usr/local/cuda 92 | ENV CPLUS_INCLUDE_PATH $CPATH 93 | ENV KERAS_BACKEND tensorflow 94 | 95 | WORKDIR /home/keras/notebook 96 | 97 | EXPOSE 8888 6006 22 8097 98 | 99 | CMD jupyter notebook --port=8888 --ip=0.0.0.0 --no-browser 100 | -------------------------------------------------------------------------------- /dockerfiles/Dockerfile2: -------------------------------------------------------------------------------- 1 | # add 7z tar and zip archivers 2 | FROM nvidia/cuda:9.0-cudnn7-devel 3 | 4 | RUN apt-get update && apt-get install -y openssh-server 5 | RUN mkdir /var/run/sshd 6 | RUN echo 'root:Ubuntu@41' | chpasswd 7 | RUN sed -i 's/PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config 8 | 9 | # SSH login fix. Otherwise user is kicked off after login 10 | RUN sed 's@session\s*required\s*pam_loginuid.so@session optional pam_loginuid.so@g' -i /etc/pam.d/sshd 11 | 12 | ENV NOTVISIBLE "in users profile" 13 | RUN echo "export VISIBLE=now" >> /etc/profile 14 | 15 | ENV CONDA_DIR /opt/conda 16 | ENV PATH $CONDA_DIR/bin:$PATH 17 | 18 | # writing env variables to /etc/profile as mentioned here https://docs.docker.com/engine/examples/running_ssh_service/#run-a-test_sshd-container 19 | RUN echo "export CONDA_DIR=/opt/conda" >> /etc/profile 20 | RUN echo "export PATH=$CONDA_DIR/bin:$PATH" >> /etc/profile 21 | 22 | RUN mkdir -p $CONDA_DIR && \ 23 | echo export PATH=$CONDA_DIR/bin:'$PATH' > /etc/profile.d/conda.sh && \ 24 | apt-get update && \ 25 | apt-get install -y wget git libhdf5-dev g++ graphviz openmpi-bin nano && \ 26 | wget --quiet https://repo.continuum.io/miniconda/Miniconda3-4.2.12-Linux-x86_64.sh && \ 27 | echo "c59b3dd3cad550ac7596e0d599b91e75d88826db132e4146030ef471bb434e9a *Miniconda3-4.2.12-Linux-x86_64.sh" | sha256sum -c - && \ 28 | /bin/bash /Miniconda3-4.2.12-Linux-x86_64.sh -f -b -p $CONDA_DIR && \ 29 | ln /usr/lib/x86_64-linux-gnu/libcudnn.so /usr/local/cuda/lib64/libcudnn.so && \ 30 | ln /usr/lib/x86_64-linux-gnu/libcudnn.so.7 /usr/local/cuda/lib64/libcudnn.so.7 && \ 31 | ln /usr/include/cudnn.h /usr/local/cuda/include/cudnn.h && \ 32 | rm Miniconda3-4.2.12-Linux-x86_64.sh 33 | 34 | ENV NB_USER keras 35 | ENV NB_UID 1000 36 | 37 | RUN echo "export NB_USER=keras" >> /etc/profile 38 | RUN echo "export NB_UID=1000" >> /etc/profile 39 | 40 | RUN echo "export LD_LIBRARY_PATH=/usr/local/cuda/lib64:/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH" >> /etc/profile 41 | RUN echo "export CPATH=/usr/include:/usr/include/x86_64-linux-gnu:/usr/local/cuda/include:$CPATH" >> /etc/profile 42 | RUN echo "export LIBRARY_PATH=/usr/local/cuda/lib64:/lib/x86_64-linux-gnu:$LIBRARY_PATH" >> /etc/profile 43 | RUN echo "export CUDA_HOME=/usr/local/cuda" >> /etc/profile 44 | RUN echo "export CPLUS_INCLUDE_PATH=$CPATH" >> /etc/profile 45 | RUN echo "export KERAS_BACKEND=tensorflow" >> /etc/profile 46 | 47 | RUN useradd -m -s /bin/bash -N -u $NB_UID $NB_USER && \ 48 | mkdir -p $CONDA_DIR && \ 49 | chown keras $CONDA_DIR -R 50 | 51 | USER keras 52 | 53 | RUN mkdir -p /home/keras/notebook 54 | 55 | # Python 56 | ARG python_version=3.5 57 | 58 | RUN conda install -y python=${python_version} && \ 59 | pip install --upgrade pip && \ 60 | pip install tensorflow-gpu && \ 61 | conda install Pillow scikit-learn notebook pandas matplotlib mkl nose pyyaml six h5py && \ 62 | conda install theano pygpu bcolz && \ 63 | pip install keras kaggle-cli lxml opencv-python requests scipy tqdm visdom && \ 64 | conda install pytorch torchvision cuda90 -c pytorch && \ 65 | conda clean -yt 66 | 67 | RUN pip install jupyter_contrib_nbextensions && \ 68 | pip install 'html5lib==0.9999999' && \ 69 | jupyter contrib nbextension install --user 70 | 71 | 72 | ENV LD_LIBRARY_PATH /usr/local/cuda/lib64:/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH 73 | ENV CPATH /usr/include:/usr/include/x86_64-linux-gnu:/usr/local/cuda/include:$CPATH 74 | ENV LIBRARY_PATH /usr/local/cuda/lib64:/lib/x86_64-linux-gnu:$LIBRARY_PATH 75 | ENV CUDA_HOME /usr/local/cuda 76 | ENV CPLUS_INCLUDE_PATH $CPATH 77 | ENV KERAS_BACKEND tensorflow 78 | 79 | WORKDIR /home/keras/notebook 80 | 81 | EXPOSE 8888 6006 22 8097 82 | 83 | CMD jupyter notebook --port=8888 --ip=0.0.0.0 --no-browser 84 | 85 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | ###### Requirements without Version Specifiers ###### 2 | cython 3 | libgdal 4 | gdal 5 | numpy 6 | opencv 7 | pip 8 | pyproj 9 | matplotlib 10 | 11 | fiona 12 | scikit-image 13 | scikit-learn 14 | scipy 15 | geopandas 16 | tqdm 17 | utm 18 | 19 | ###### Requirements with Version Specifiers ###### 20 | shapely=1.5.16 21 | smnx==0.5.1 22 | networkx=1.11 23 | -------------------------------------------------------------------------------- /satellites.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snakers4/spacenet-three/5e1b202b824c6d976298488225078d850c908a8e/satellites.jpg -------------------------------------------------------------------------------- /scripts/cnn_inference_env.sh: -------------------------------------------------------------------------------- 1 | conda install -y -c conda-forge numba 2 | conda install -y -c conda-forge rasterio 3 | conda install -y -c conda-forge cython 4 | conda install -y -c conda-forge rasterio 5 | conda install -y -c conda-forge libgdal 6 | conda install -y -c conda-forge gdal 7 | 8 | -------------------------------------------------------------------------------- /scripts/create_8bit_test_images.py: -------------------------------------------------------------------------------- 1 | import pandas as pd 2 | from multiprocessing import Pool 3 | import tqdm 4 | import numpy as np 5 | import os 6 | import glob as glob 7 | from skimage.io import imread, imsave 8 | import osmnx as ox 9 | import numpy as np 10 | import matplotlib 11 | import matplotlib.pyplot as plt 12 | import geopandas as gpd 13 | from osgeo import gdal, ogr, osr 14 | import cv2 15 | import subprocess 16 | import shapely 17 | from shapely.geometry import MultiLineString 18 | from matplotlib.patches import PathPatch 19 | import matplotlib.path 20 | 21 | imgs = [] 22 | 23 | # change this to your data prefix 24 | path_prefix = '../data' 25 | 26 | # default variables from the hosts of the challenge 27 | buffer_meters = 2 28 | burnValue = 150 29 | 30 | # only test folders 31 | folders = ['AOI_2_Vegas_Roads_Test_Public', 32 | 'AOI_3_Paris_Roads_Test_Public', 33 | 'AOI_4_Shanghai_Roads_Test_Public', 34 | 'AOI_5_Khartoum_Roads_Test_Public'] 35 | 36 | # image types 37 | prefix_dict = { 38 | 'mul': 'MUL', 39 | 'muls': 'MUL-PanSharpen', 40 | 'pan': 'PAN', 41 | 'rgbps': 'RGB-PanSharpen', 42 | } 43 | 44 | for folder in folders: 45 | for prefix in prefix_dict.items(): 46 | g = glob.glob(path_prefix+'/{}/{}/*.tif'.format(folder,prefix[1])) 47 | imgs.extend(g) 48 | 49 | img_folders = [(img.split('/')[2]) for img in imgs] 50 | img_subfolders = [(img.split('/')[3]) for img in imgs] 51 | img_files = [(img.split('/')[4]) for img in imgs] 52 | 53 | def create_8bit_test_images(input_data): 54 | img_path = input_data[0] 55 | img_folder = input_data[1] 56 | img_subfolder = input_data[2] 57 | img_file = input_data[3] 58 | 59 | # create paths for masks and 8bit images 60 | bit8_folder = os.path.join(path_prefix,img_folder,img_subfolder+'_8bit') 61 | bit8_path = os.path.join(bit8_folder,img_file) 62 | 63 | if not os.path.exists(bit8_folder): 64 | os.mkdir(bit8_folder) 65 | if os.path.isfile(bit8_path): 66 | os.remove(bit8_path) 67 | 68 | try: 69 | # convert images to 8-bit 70 | convert_to_8Bit(img_path, 71 | bit8_path, 72 | outputPixType='Byte', 73 | outputFormat='GTiff', 74 | rescale_type='rescale', 75 | percentiles=[2,98]) 76 | 77 | except BaseException as e: 78 | print(str(e)) 79 | 80 | return [bit8_folder,bit8_path,img_path,img_folder,img_subfolder,img_file] 81 | 82 | def convert_to_8Bit(inputRaster, outputRaster, 83 | outputPixType='Byte', 84 | outputFormat='GTiff', 85 | rescale_type='rescale', 86 | percentiles=[2, 98]): 87 | ''' 88 | Convert 16bit image to 8bit 89 | rescale_type = [clip, rescale] 90 | if clip, scaling is done strictly between 0 65535 91 | if rescale, each band is rescaled to a min and max 92 | set by percentiles 93 | ''' 94 | 95 | srcRaster = gdal.Open(inputRaster) 96 | cmd = ['gdal_translate', '-ot', outputPixType, '-of', 97 | outputFormat] 98 | 99 | # iterate through bands 100 | for bandId in range(srcRaster.RasterCount): 101 | bandId = bandId+1 102 | band = srcRaster.GetRasterBand(bandId) 103 | if rescale_type == 'rescale': 104 | bmin = band.GetMinimum() 105 | bmax = band.GetMaximum() 106 | # if not exist minimum and maximum values 107 | if bmin is None or bmax is None: 108 | (bmin, bmax) = band.ComputeRasterMinMax(1) 109 | # else, rescale 110 | band_arr_tmp = band.ReadAsArray() 111 | bmin = np.percentile(band_arr_tmp.flatten(), 112 | percentiles[0]) 113 | bmax= np.percentile(band_arr_tmp.flatten(), 114 | percentiles[1]) 115 | 116 | else: 117 | bmin, bmax = 0, 65535 118 | 119 | cmd.append('-scale_{}'.format(bandId)) 120 | cmd.append('{}'.format(bmin)) 121 | cmd.append('{}'.format(bmax)) 122 | cmd.append('{}'.format(0)) 123 | cmd.append('{}'.format(255)) 124 | 125 | cmd.append(inputRaster) 126 | cmd.append(outputRaster) 127 | # print("Conversin command:", cmd) 128 | subprocess.call(cmd) 129 | 130 | return 131 | 132 | input_data = zip(imgs,img_folders,img_subfolders,img_files) 133 | input_data = [item for item in input_data] 134 | 135 | with Pool(10) as p: 136 | bit8_data = list(tqdm.tqdm(p.imap(create_8bit_test_images, input_data), 137 | total=len(input_data))) 138 | # transpose the list 139 | bit8_data = list(map(list, zip(*bit8_data))) 140 | 141 | bit8_df = pd.DataFrame() 142 | 143 | for i,key in enumerate(['bit8_folder','bit8_path','img_path','img_folder','img_subfolder','img_file']): 144 | bit8_df[key] = bit8_data[i] 145 | 146 | bit8_df.to_csv('bit8_test.csv') 147 | 148 | -------------------------------------------------------------------------------- /scripts/create_binary_masks.py: -------------------------------------------------------------------------------- 1 | import pandas as pd 2 | from multiprocessing import Pool 3 | import tqdm 4 | import numpy as np 5 | import os 6 | import glob as glob 7 | from skimage.io import imread, imsave 8 | import osmnx as ox 9 | import numpy as np 10 | import matplotlib 11 | import matplotlib.pyplot as plt 12 | import geopandas as gpd 13 | from osgeo import gdal, ogr, osr 14 | import cv2 15 | import subprocess 16 | import shapely 17 | from shapely.geometry import MultiLineString 18 | from matplotlib.patches import PathPatch 19 | import matplotlib.path 20 | 21 | imgs = [] 22 | 23 | # change this to your data prefix 24 | path_prefix = 'data' 25 | 26 | # default variables from the hosts of the challenge 27 | buffer_meters = 2 28 | burnValue = 150 29 | 30 | # only train folders 31 | folders = ['AOI_2_Vegas_Roads_Train', 32 | 'AOI_5_Khartoum_Roads_Train', 33 | 'AOI_3_Paris_Roads_Train', 34 | 'AOI_4_Shanghai_Roads_Train'] 35 | 36 | # image types 37 | prefix_dict = { 38 | 'mul': 'MUL', 39 | 'muls': 'MUL-PanSharpen', 40 | 'pan': 'PAN', 41 | 'rgbps': 'RGB-PanSharpen', 42 | } 43 | 44 | for folder in folders: 45 | for prefix in prefix_dict.items(): 46 | g = glob.glob(path_prefix+'/{}/{}/*.tif'.format(folder,prefix[1])) 47 | imgs.extend(g) 48 | 49 | img_folders = [(img.split('/')[1]) for img in imgs] 50 | img_subfolders = [(img.split('/')[2]) for img in imgs] 51 | img_files = [(img.split('/')[3]) for img in imgs] 52 | 53 | def create_binary_mask(input_data): 54 | img_path = input_data[0] 55 | img_folder = input_data[1] 56 | img_subfolder = input_data[2] 57 | img_file = input_data[3] 58 | 59 | # create paths for masks and 8bit images 60 | label_file = os.path.join(path_prefix,img_folder,'geojson/spacenetroads','spacenetroads_AOI'+img_file.split('AOI')[1][0:-3]+'geojson') 61 | bit8_folder = os.path.join(path_prefix,img_folder,img_subfolder+'_8bit') 62 | bit8_path = os.path.join(bit8_folder,img_file) 63 | mask_folder = os.path.join(path_prefix,img_folder,img_subfolder+'_mask') 64 | mask_path = os.path.join(mask_folder,img_file[:-3])+'png' 65 | # vis_folder = os.path.join(path_prefix,img_folder,img_subfolder+'_vis') 66 | # vis_path = os.path.join(vis_folder,img_file[:-3])+'png' 67 | 68 | # print(label_file) 69 | # create the necessary folders and remove the existing files 70 | 71 | if not os.path.exists(bit8_folder): 72 | os.mkdir(bit8_folder) 73 | if not os.path.exists(mask_folder): 74 | os.mkdir(mask_folder) 75 | # if not os.path.exists(vis_folder): 76 | # os.mkdir(vis_folder) 77 | # if os.path.isfile(vis_path): 78 | # os.remove(vis_path) 79 | if os.path.isfile(bit8_path): 80 | os.remove(bit8_path) 81 | if os.path.isfile(mask_path): 82 | os.remove(mask_path) 83 | if os.path.isfile(mask_path[:-3]+'jpg'): 84 | os.remove(mask_path[:-3]+'jpg') 85 | 86 | try: 87 | # convert images to 8-bit 88 | convert_to_8Bit(img_path, 89 | bit8_path, 90 | outputPixType='Byte', 91 | outputFormat='GTiff', 92 | rescale_type='rescale', 93 | percentiles=[2,98]) 94 | 95 | # create masks 96 | # note that though the output raster file has .png extension 97 | # in reality I delete this file and save only jpg version later 98 | 99 | mask, gdf_buffer = get_road_buffer(geoJson = label_file, 100 | im_vis_file = bit8_path, 101 | output_raster = mask_path, 102 | buffer_meters= buffer_meters, 103 | burnValue= burnValue, 104 | bufferRoundness=6, 105 | plot_file='', # this indicates that no visualization plot is required 106 | figsize= (6,6), 107 | fontsize=8, 108 | dpi=200, 109 | show_plot=False, 110 | verbose=False) 111 | 112 | # read the png file, save it as jpeg and 113 | mask = imread(mask_path) 114 | imsave(fname=mask_path[:-3]+'jpg',arr = mask) 115 | mask_max = np.max(mask) 116 | del mask 117 | # remove the png file, but keep the 8-bit mask 118 | os.remove(mask_path) 119 | except BaseException as e: 120 | print(str(e)) 121 | mask_max = -1 122 | 123 | return [label_file,bit8_folder,bit8_path,mask_folder,mask_path[:-3]+'jpg',img_path,img_folder,img_subfolder,img_file,mask_max] 124 | 125 | def get_road_buffer(geoJson, im_vis_file, output_raster, 126 | buffer_meters=2, burnValue=1, 127 | bufferRoundness=6, 128 | plot_file='', figsize=(6,6), fontsize=6, 129 | dpi=800, show_plot=False, 130 | verbose=False): 131 | ''' 132 | Get buffer around roads defined by geojson and image files. 133 | Calls create_buffer_geopandas() and gdf_to_array(). 134 | Assumes in_vis_file is an 8-bit RGB file. 135 | Returns geodataframe and ouptut mask. 136 | ''' 137 | 138 | gdf_buffer = create_buffer_geopandas(geoJson, 139 | bufferDistanceMeters=buffer_meters, 140 | bufferRoundness=bufferRoundness, 141 | projectToUTM=True) 142 | 143 | 144 | # create label image 145 | if len(gdf_buffer) == 0: 146 | mask_gray = np.zeros(cv2.imread(im_vis_file,0).shape) 147 | cv2.imwrite(output_raster, mask_gray) 148 | else: 149 | gdf_to_array(gdf_buffer, im_vis_file, output_raster, 150 | burnValue=burnValue) 151 | 152 | # load mask 153 | mask_gray = cv2.imread(output_raster, 0) 154 | 155 | # make plots 156 | if plot_file: 157 | # plot all in a line 158 | if (figsize[0] != figsize[1]): 159 | fig, (ax0, ax1, ax2, ax3) = plt.subplots(1,4, figsize=figsize)#(13,4)) 160 | # else, plot a 2 x 2 grid 161 | else: 162 | fig, ((ax0, ax1), (ax2, ax3)) = plt.subplots(2,2, figsize=figsize) 163 | 164 | # road lines 165 | try: 166 | gdfRoadLines = gpd.read_file(geoJson) 167 | gdfRoadLines.plot(ax=ax0, marker='o', color='red') 168 | except: 169 | ax0.imshow(mask_gray) 170 | ax0.axis('off') 171 | ax0.set_aspect('equal') 172 | ax0.set_title('Roads from GeoJson', fontsize=fontsize) 173 | 174 | # first show raw image 175 | im_vis = cv2.imread(im_vis_file, 1) 176 | img_mpl = cv2.cvtColor(im_vis, cv2.COLOR_BGR2RGB) 177 | ax1.imshow(img_mpl) 178 | ax1.axis('off') 179 | ax1.set_title('8-bit RGB Image', fontsize=fontsize) 180 | 181 | # plot mask 182 | ax2.imshow(mask_gray) 183 | ax2.axis('off') 184 | ax2.set_title('Roads Mask (' + str(np.round(buffer_meters)) \ 185 | + ' meter buffer)', fontsize=fontsize) 186 | 187 | # plot combined 188 | ax3.imshow(img_mpl) 189 | # overlay mask 190 | # set zeros to nan 191 | z = mask_gray.astype(float) 192 | z[z==0] = np.nan 193 | # change palette to orange 194 | palette = plt.cm.gray 195 | #palette.set_over('yellow', 0.9) 196 | palette.set_over('lime', 0.9) 197 | ax3.imshow(z, cmap=palette, alpha=0.66, 198 | norm=matplotlib.colors.Normalize(vmin=0.5, vmax=0.9, clip=False)) 199 | ax3.set_title('8-bit RGB Image + Buffered Roads', fontsize=fontsize) 200 | ax3.axis('off') 201 | 202 | #plt.axes().set_aspect('equal', 'datalim') 203 | 204 | plt.tight_layout() 205 | plt.savefig(plot_file, dpi=dpi) 206 | if not show_plot: 207 | plt.close() 208 | 209 | return mask_gray, gdf_buffer 210 | 211 | def create_buffer_geopandas(geoJsonFileName, 212 | bufferDistanceMeters=2, 213 | bufferRoundness=1, 214 | projectToUTM=True): 215 | ''' 216 | Create a buffer around the lines of the geojson. 217 | Return a geodataframe. 218 | ''' 219 | 220 | inGDF = gpd.read_file(geoJsonFileName) 221 | 222 | # set a few columns that we will need later 223 | inGDF['type'] = inGDF['road_type'].values 224 | inGDF['class'] = 'highway' 225 | inGDF['highway'] = 'highway' 226 | 227 | if len(inGDF) == 0: 228 | return [], [] 229 | 230 | # Transform gdf Roadlines into UTM so that Buffer makes sense 231 | if projectToUTM: 232 | tmpGDF = ox.project_gdf(inGDF) 233 | else: 234 | tmpGDF = inGDF 235 | 236 | gdf_utm_buffer = tmpGDF 237 | 238 | # perform Buffer to produce polygons from Line Segments 239 | gdf_utm_buffer['geometry'] = tmpGDF.buffer(bufferDistanceMeters, 240 | bufferRoundness) 241 | 242 | gdf_utm_dissolve = gdf_utm_buffer.dissolve(by='class') 243 | gdf_utm_dissolve.crs = gdf_utm_buffer.crs 244 | 245 | if projectToUTM: 246 | gdf_buffer = gdf_utm_dissolve.to_crs(inGDF.crs) 247 | else: 248 | gdf_buffer = gdf_utm_dissolve 249 | 250 | return gdf_buffer 251 | 252 | def gdf_to_array(gdf, im_file, output_raster, burnValue=150): 253 | 254 | ''' 255 | Turn geodataframe to array, save as image file with non-null pixels 256 | set to burnValue 257 | ''' 258 | 259 | NoData_value = 0 # -9999 260 | 261 | gdata = gdal.Open(im_file) 262 | 263 | # set target info 264 | target_ds = gdal.GetDriverByName('GTiff').Create(output_raster, 265 | gdata.RasterXSize, 266 | gdata.RasterYSize, 1, gdal.GDT_Byte) 267 | target_ds.SetGeoTransform(gdata.GetGeoTransform()) 268 | 269 | # set raster info 270 | raster_srs = osr.SpatialReference() 271 | raster_srs.ImportFromWkt(gdata.GetProjectionRef()) 272 | target_ds.SetProjection(raster_srs.ExportToWkt()) 273 | 274 | band = target_ds.GetRasterBand(1) 275 | band.SetNoDataValue(NoData_value) 276 | 277 | outdriver=ogr.GetDriverByName('MEMORY') 278 | outDataSource=outdriver.CreateDataSource('memData') 279 | tmp=outdriver.Open('memData',1) 280 | outLayer = outDataSource.CreateLayer("states_extent", raster_srs, 281 | geom_type=ogr.wkbMultiPolygon) 282 | # burn 283 | burnField = "burn" 284 | idField = ogr.FieldDefn(burnField, ogr.OFTInteger) 285 | outLayer.CreateField(idField) 286 | featureDefn = outLayer.GetLayerDefn() 287 | for geomShape in gdf['geometry'].values: 288 | 289 | outFeature = ogr.Feature(featureDefn) 290 | outFeature.SetGeometry(ogr.CreateGeometryFromWkt(geomShape.wkt)) 291 | outFeature.SetField(burnField, burnValue) 292 | outLayer.CreateFeature(outFeature) 293 | outFeature = 0 294 | 295 | gdal.RasterizeLayer(target_ds, [1], outLayer, burn_values=[burnValue]) 296 | outLayer = 0 297 | outDatSource = 0 298 | tmp = 0 299 | 300 | return 301 | 302 | def convert_to_8Bit(inputRaster, outputRaster, 303 | outputPixType='Byte', 304 | outputFormat='GTiff', 305 | rescale_type='rescale', 306 | percentiles=[2, 98]): 307 | ''' 308 | Convert 16bit image to 8bit 309 | rescale_type = [clip, rescale] 310 | if clip, scaling is done strictly between 0 65535 311 | if rescale, each band is rescaled to a min and max 312 | set by percentiles 313 | ''' 314 | 315 | srcRaster = gdal.Open(inputRaster) 316 | cmd = ['gdal_translate', '-ot', outputPixType, '-of', 317 | outputFormat] 318 | 319 | # iterate through bands 320 | for bandId in range(srcRaster.RasterCount): 321 | bandId = bandId+1 322 | band = srcRaster.GetRasterBand(bandId) 323 | if rescale_type == 'rescale': 324 | bmin = band.GetMinimum() 325 | bmax = band.GetMaximum() 326 | # if not exist minimum and maximum values 327 | if bmin is None or bmax is None: 328 | (bmin, bmax) = band.ComputeRasterMinMax(1) 329 | # else, rescale 330 | band_arr_tmp = band.ReadAsArray() 331 | bmin = np.percentile(band_arr_tmp.flatten(), 332 | percentiles[0]) 333 | bmax= np.percentile(band_arr_tmp.flatten(), 334 | percentiles[1]) 335 | 336 | else: 337 | bmin, bmax = 0, 65535 338 | 339 | cmd.append('-scale_{}'.format(bandId)) 340 | cmd.append('{}'.format(bmin)) 341 | cmd.append('{}'.format(bmax)) 342 | cmd.append('{}'.format(0)) 343 | cmd.append('{}'.format(255)) 344 | 345 | cmd.append(inputRaster) 346 | cmd.append(outputRaster) 347 | # print("Conversin command:", cmd) 348 | subprocess.call(cmd) 349 | 350 | return 351 | 352 | 353 | input_data = zip(imgs,img_folders,img_subfolders,img_files) 354 | input_data = [item for item in input_data] 355 | 356 | with Pool(10) as p: 357 | mask_data = list(tqdm.tqdm(p.imap(create_binary_mask, input_data), 358 | total=len(input_data))) 359 | # transpose the list 360 | mask_data = list(map(list, zip(*mask_data))) 361 | 362 | mask_df = pd.DataFrame() 363 | 364 | for i,key in enumerate(['label_file','bit8_folder','bit8_path','mask_folder','mask_path','img_path','img_folder','img_subfolder','img_file', 'mask_max']): 365 | mask_df[key] = mask_data[i] 366 | 367 | mask_df.to_csv('mask_df.csv') 368 | -------------------------------------------------------------------------------- /scripts/download.sh: -------------------------------------------------------------------------------- 1 | aws s3api get-object --bucket spacenet-dataset --key SpaceNet_Roads_Competition/AOI_2_Vegas_Roads_Test_Public.tar.gz --request-payer requester /media/server/13_DS_BACKUP/satellites_roads/AOI_2_Vegas_Roads_Test_Public.tar.gz 2 | aws s3api get-object --bucket spacenet-dataset --key SpaceNet_Roads_Competition/AOI_2_Vegas_Roads_Train.tar.gz --request-payer requester /media/server/13_DS_BACKUP/satellites_roads/AOI_2_Vegas_Roads_Train.tar.gz 3 | aws s3api get-object --bucket spacenet-dataset --key SpaceNet_Roads_Competition/AOI_3_Paris_Roads_Test_Public.tar.gz --request-payer requester /media/server/13_DS_BACKUP/satellites_roads/AOI_3_Paris_Roads_Test_Public.tar.gz 4 | aws s3api get-object --bucket spacenet-dataset --key SpaceNet_Roads_Competition/AOI_3_Paris_Roads_Train.tar.gz --request-payer requester /media/server/13_DS_BACKUP/satellites_roads/AOI_3_Paris_Roads_Train.tar.gz 5 | aws s3api get-object --bucket spacenet-dataset --key SpaceNet_Roads_Competition/AOI_4_Shanghai_Roads_Test_Public.tar.gz --request-payer requester /media/server/13_DS_BACKUP/satellites_roads/AOI_4_Shanghai_Roads_Test_Public.tar.gz 6 | aws s3api get-object --bucket spacenet-dataset --key SpaceNet_Roads_Competition/AOI_4_Shanghai_Roads_Train.tar.gz --request-payer requester /media/server/13_DS_BACKUP/satellites_roads/AOI_4_Shanghai_Roads_Train.tar.gz 7 | aws s3api get-object --bucket spacenet-dataset --key SpaceNet_Roads_Competition/AOI_5_Khartoum_Roads_Test_Public.tar.gz --request-payer requester /media/server/13_DS_BACKUP/satellites_roads/AOI_5_Khartoum_Roads_Test_Public.tar.gz 8 | aws s3api get-object --bucket spacenet-dataset --key SpaceNet_Roads_Competition/AOI_5_Khartoum_Roads_Train.tar.gz --request-payer requester /media/server/13_DS_BACKUP/satellites_roads/AOI_5_Khartoum_Roads_Train.tar.gz 9 | -------------------------------------------------------------------------------- /scripts/env.sh: -------------------------------------------------------------------------------- 1 | conda install -y -c conda-forge cython 2 | conda install -y -c conda-forge rasterio 3 | conda install -y -c conda-forge libgdal 4 | conda install -y -c conda-forge gdal 5 | conda install -y -c conda-forge scikit-image 6 | conda install -y -c conda-forge pyproj 7 | conda install -y -c conda-forge geopandas 8 | conda install -y -c conda-forge tqdm 9 | conda install -y -c conda-forge shapely=1.5.16 10 | conda install -y -c conda-forge scipy 11 | conda install -y -c conda-forge networkx=1.11 12 | conda install -y -c conda-forge fiona 13 | pip3 install utm 14 | pip3 install osmnx==0.5.1 15 | -------------------------------------------------------------------------------- /scripts/install-dependencies.sh: -------------------------------------------------------------------------------- 1 | conda install -y scikit-image 2 | conda install -y tqdm 3 | conda install -y shapely=1.5.16 4 | conda install -y pyproj 5 | conda install -y cython 6 | conda install -y networkx=1.11 7 | conda install -y fiona 8 | conda install -c -y conda-forge geopandas 9 | conda install -c -y conda-forge rasterio 10 | conda install -y geos=3.5.0 11 | 12 | -------------------------------------------------------------------------------- /sknw.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from numba import jit 3 | import networkx as nx 4 | 5 | # author https://github.com/yxdragon/sknw/blob/master/sknw/sknw.py 6 | 7 | # get neighbors d index 8 | def neighbors(shape): 9 | dim = len(shape) 10 | block = np.ones([3]*dim) 11 | block[tuple([1]*dim)] = 0 12 | idx = np.where(block>0) 13 | idx = np.array(idx, dtype=np.uint8).T 14 | idx = np.array(idx-[1]*dim) 15 | acc = np.cumprod((1,)+shape[::-1][:-1]) 16 | return np.dot(idx, acc[::-1]) 17 | 18 | @jit # my mark 19 | def mark(img): # mark the array use (0, 1, 2) 20 | nbs = neighbors(img.shape) 21 | img = img.ravel() 22 | for p in range(len(img)): 23 | if img[p]==0:continue 24 | s = 0 25 | for dp in nbs: 26 | if img[p+dp]!=0:s+=1 27 | if s==2:img[p]=1 28 | else:img[p]=2 29 | 30 | @jit # trans index to r, c... 31 | def idx2rc(idx, acc): 32 | rst = np.zeros((len(idx), len(acc)), dtype=np.int16) 33 | for i in range(len(idx)): 34 | for j in range(len(acc)): 35 | rst[i,j] = idx[i]//acc[j] 36 | idx[i] -= rst[i,j]*acc[j] 37 | rst -= 1 38 | return rst 39 | 40 | @jit # fill a node (may be two or more points) 41 | def fill(img, p, num, nbs, acc, buf): 42 | back = img[p] 43 | img[p] = num 44 | buf[0] = p 45 | cur = 0; s = 1; 46 | 47 | while True: 48 | p = buf[cur] 49 | for dp in nbs: 50 | cp = p+dp 51 | if img[cp]==back: 52 | img[cp] = num 53 | buf[s] = cp 54 | s+=1 55 | cur += 1 56 | if cur==s:break 57 | return idx2rc(buf[:s], acc) 58 | 59 | @jit # trace the edge and use a buffer, then buf.copy, if use [] numba not works 60 | def trace(img, p, nbs, acc, buf): 61 | c1 = 0; c2 = 0; 62 | newp = 0 63 | cur = 0 64 | 65 | while True: 66 | buf[cur] = p 67 | img[p] = 0 68 | cur += 1 69 | for dp in nbs: 70 | cp = p + dp 71 | if img[cp] >= 10: 72 | if c1==0:c1=img[cp] 73 | else: c2 = img[cp] 74 | if img[cp] == 1: 75 | newp = cp 76 | p = newp 77 | if c2!=0:break 78 | return (c1-10, c2-10, idx2rc(buf[:cur], acc)) 79 | 80 | @jit # parse the image then get the nodes and edges 81 | def parse_struc(img): 82 | nbs = neighbors(img.shape) 83 | acc = np.cumprod((1,)+img.shape[::-1][:-1])[::-1] 84 | img = img.ravel() 85 | pts = np.array(np.where(img==2))[0] 86 | buf = np.zeros(131072, dtype=np.int64) 87 | num = 10 88 | nodes = [] 89 | for p in pts: 90 | if img[p] == 2: 91 | nds = fill(img, p, num, nbs, acc, buf) 92 | num += 1 93 | nodes.append(nds) 94 | 95 | edges = [] 96 | for p in pts: 97 | for dp in nbs: 98 | if img[p+dp]==1: 99 | edge = trace(img, p+dp, nbs, acc, buf) 100 | edges.append(edge) 101 | return nodes, edges 102 | 103 | # use nodes and edges build a networkx graph 104 | def build_graph(nodes, edges, multi=False): 105 | graph = nx.MultiGraph() if multi else nx.Graph() 106 | for i in range(len(nodes)): 107 | graph.add_node(i, pts=nodes[i], o=nodes[i].mean(axis=0)) 108 | for s,e,pts in edges: 109 | l = np.linalg.norm(pts[1:]-pts[:-1], axis=1).sum() 110 | graph.add_edge(s,e, pts=pts, weight=l) 111 | return graph 112 | 113 | def buffer(ske): 114 | buf = np.zeros(tuple(np.array(ske.shape)+2), dtype=np.uint16) 115 | buf[tuple([slice(1,-1)]*buf.ndim)] = ske 116 | return buf 117 | 118 | def build_sknw(ske, multi=False): 119 | buf = buffer(ske) 120 | mark(buf) 121 | nodes, edges = parse_struc(buf) 122 | return build_graph(nodes, edges, multi) 123 | 124 | # draw the graph 125 | def draw_graph(img, graph, cn=255, ce=128): 126 | acc = np.cumprod((1,)+img.shape[::-1][:-1])[::-1] 127 | img = img.ravel() 128 | for idx in graph.nodes(): 129 | pts = graph.node[idx]['pts'] 130 | img[np.dot(pts, acc)] = cn 131 | for (s, e) in graph.edges(): 132 | eds = graph[s][e] 133 | for i in eds: 134 | pts = eds[i]['pts'] 135 | img[np.dot(pts, acc)] = ce 136 | 137 | if __name__ == '__main__': 138 | g = nx.MultiGraph() 139 | g.add_nodes_from([1,2,3,4,5]) 140 | g.add_edges_from([(1,2),(1,3),(2,3),(4,5),(5,4)]) 141 | print(g.nodes()) 142 | print(g.edges()) 143 | a = g.subgraph(1) 144 | print('d') 145 | print(a) 146 | print('d') -------------------------------------------------------------------------------- /src/.ipynb_checkpoints/readme.md-checkpoint.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "metadata": { 6 | "ExecuteTime": { 7 | "end_time": "2018-02-04T07:16:18.892637Z", 8 | "start_time": "2018-02-04T07:16:18.882802Z" 9 | } 10 | }, 11 | "source": [ 12 | "# 1 Hardware requirements\n", 13 | "\n", 14 | "**Training**\n", 15 | "\n", 16 | "- 6+ core modern CPU (Xeon, i7) for fast image pre-processing;\n", 17 | "- The models were trained on 2 * GeForce 1080 Ti;\n", 18 | "- Training time on my setup ~ **3 hours** for models with 8-bit images as inputs;\n", 19 | "- Disk space - 40GB should be more than enough;\n", 20 | "\n", 21 | "**Inference**\n", 22 | "\n", 23 | "- 6+ core modern CPU (Xeon, i7) for fast image pre-processing;\n", 24 | "- On 2 * GeForce 1080 Ti inference takes **3-5 minutes**;\n", 25 | "- Graph creation takes **5-10 minutes**;\n", 26 | "\n", 27 | "# 2 Preparing and launching the Docker environment\n", 28 | "\n", 29 | "**Clone the repository**\n", 30 | "\n", 31 | "`git clone https://github.com/snakers4/spacenet-three .`\n", 32 | "\n", 33 | "\n", 34 | "**This repository contains 2 Dockerfiles**\n", 35 | "- `/dockerfiles/Dockerfile` - this is the main Dockerfile which was used as environment to run the training and inference scripts\n", 36 | "- `/dockerfiles/Dockerfile2`- this is an additional backup Docker file with newer versions of the drivers and PyTorch, just in case\n", 37 | "\n", 38 | "**Build a Docker image**\n", 39 | "\n", 40 | "`\n", 41 | "cd dockerfiles\n", 42 | "docker build -t aveysov .\n", 43 | "`\n", 44 | "\n", 45 | "**Install the latest nvidia docker**\n", 46 | "\n", 47 | "Follow instructions from [here](https://github.com/NVIDIA/nvidia-docker).\n", 48 | "Please prefer nvidia-docker2 for more stable performance.\n", 49 | "\n", 50 | "\n", 51 | "To test all works fine run:\n", 52 | "\n", 53 | "\n", 54 | "`docker run --runtime=nvidia --rm nvidia/cuda nvidia-smi`\n", 55 | "\n", 56 | "**(IMPORTANT) Run docker container (IMPORTANT)**\n", 57 | "\n", 58 | "Unless you use this exact command (with --shm-size flag) (you can change ports and mounted volumes, of course), then the PyTorch generators **WILL NOT WORK**. \n", 59 | "\n", 60 | "\n", 61 | "- nvidia-docker 2: `docker run --runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=all -it -v /path/to/cloned/repository:/home/keras/notebook -p 8888:8888 -p 6006:6006 --shm-size 8G aveysov`\n", 62 | "- nvidia-docker: `nvidia-docker -it -v /path/to/cloned/repository:/home/keras/notebook -p 8888:8888 -p 6006:6006 --shm-size 8G aveysov`\n", 63 | "\n", 64 | "**Installing project specific software**\n", 65 | "\n", 66 | "1. Exec into the docker machine via `docker exec -it --user root YOUR_CONTAINER_ID /bin/bash`\n", 67 | "2. Run these scripts one after another\n", 68 | "```\n", 69 | "conda install -y -c conda-forge cython\n", 70 | "conda install -y -c conda-forge rasterio\n", 71 | "conda install -y -c conda-forge libgdal\n", 72 | "conda install -y -c conda-forge gdal\n", 73 | "conda install -y -c conda-forge scikit-image\n", 74 | "conda install -y -c conda-forge pyproj\n", 75 | "conda install -y -c conda-forge geopandas\n", 76 | "conda install -y -c conda-forge tqdm\n", 77 | "conda install -y -c conda-forge shapely=1.5.16\n", 78 | "conda install -y -c conda-forge scipy\n", 79 | "conda install -y -c conda-forge networkx=1.11\n", 80 | "conda install -y -c conda-forge fiona\n", 81 | "pip3 install utm\n", 82 | "pip3 install osmnx==0.5.1\n", 83 | "```\n", 84 | "3. Run these scripts one after another\n", 85 | "```\n", 86 | "pip3 install numba\n", 87 | "conda install -y -c conda-forge scikit-image\n", 88 | "```\n", 89 | "\n", 90 | "Steps 2-3 are to ensure compatibility with legacy software from APLS [repository](https://github.com/CosmiQ/apls).\n", 91 | "An alternative to that - is to use pip's requirements.txt in the same order.\n", 92 | "These steps are required to run 8-bit mask creation step from APLS repository and mask creation step.\n", 93 | "If you will be trying to re-do this step - reserve 5-6 hours for experiments.\n", 94 | "\n", 95 | "**To start the stopped container**\n", 96 | "\n", 97 | "\n", 98 | "`docker start -i YOUR_CONTAINER_ID`\n", 99 | "\n", 100 | "\n", 101 | "# 3 Preparing the data and the machine for running scripts\n", 102 | "\n", 103 | "- Ssh into the docker container via `docker exec -it YOUR_CONTAINER_ID`\n", 104 | "- Cd to the root folder of thre repo\n", 105 | "- Dowload the data into `data/`\n", 106 | "- Run these commands:\n", 107 | " - `mkdir src/weights`\n", 108 | " - `mkdir src/tb_logs`\n", 109 | " - `cd scripts` \n", 110 | " - `python3 create_binary_masks.py` \n", 111 | " - `python3 create_8bit_test_images.py` \n", 112 | " \n", 113 | " \n", 114 | "After all of your manipulations your directory should look like:\n", 115 | "\n", 116 | "```\n", 117 | "├── README.md <- The top-level README for developers using this project.\n", 118 | "├── data\n", 119 | "│ ├── AOI_2_Vegas_Roads_Test_Public <- Test set for each city\n", 120 | "│ └── AOI_2_Vegas_Roads_Train <- Train set for each city\n", 121 | "│ ├─ geojson\n", 122 | "│ ├─ summaryData\n", 123 | "│ ├─ MUL\n", 124 | "│ ├─ RGB-PanSharpen\n", 125 | "│ ├─ PAN\n", 126 | "│ ├─ MUL-PanSharpen\n", 127 | "│ ├─ MUL-PanSharpen_8bit\n", 128 | "│ ├─ RGB-PanSharpen_8bit\n", 129 | "│ ├─ PAN_8bit\n", 130 | "│ ├─ MUL_8bit\n", 131 | "│ ├─ MUL-PanSharpen_mask\n", 132 | "│ ├─ RGB-PanSharpen_mask\n", 133 | "│ ├─ PAN_mask\n", 134 | "│ ├─ MUL_mask\n", 135 | "│ └─ RGB-PanSharpen_mask\n", 136 | "│ │\n", 137 | "│ ...\n", 138 | "│ │\n", 139 | "│ ├── AOI_5_Khartoum_Roads_Test_Public <- Test set for each city\n", 140 | "│ └── AOI_5_Khartoum_Roads_Train <- Train set for each city\n", 141 | "│\n", 142 | "├── dockerfiles <- A folder with Dockerfiles\n", 143 | "│\n", 144 | "├── src <- Source code\n", 145 | "│\n", 146 | "└── scripts <- One-off preparation scripts\n", 147 | "```\n", 148 | "\n", 149 | "# 4 Training the model\n", 150 | "\n", 151 | "If all is ok, then use the following command to train the model\n", 152 | "\n", 153 | "- Ssh into the docker container via `docker exec -it YOUR_CONTAINER_ID`\n", 154 | "- Cd to the root folder of thre repo\n", 155 | "- `cd src`\n", 156 | "- optional - turn on tensorboard for monitoring progress `tensorboard --logdir='satellites_roads/src/tb_logs' --port=6006` via jupyter notebook console or via tmux + docker exec (model converges in 30-40 epochs)\n", 157 | "- then\n", 158 | "```\n", 159 | "echo 'python3 train_satellites.py \\\n", 160 | "\t--arch linknet34 --batch-size 6 \\\n", 161 | "\t--imsize 1280 --preset mul_ps_vegetation --augs True \\\n", 162 | "\t--workers 6 --epochs 40 --start-epoch 0 \\\n", 163 | "\t--seed 42 --print-freq 20 \\\n", 164 | "\t--lr 1e-3 --optimizer adam \\\n", 165 | "\t--tensorboard True --lognumber ln34_mul_ps_vegetation_aug_dice' > train.sh\n", 166 | "```\n", 167 | "- `sh train.sh`\n", 168 | "\n", 169 | "# 5 Predicting masks\n", 170 | "\n", 171 | "- Ssh into the docker container via `docker exec -it YOUR_CONTAINER_ID`\n", 172 | "- Cd to the root folder of thre repo\n", 173 | "- `cd src`\n", 174 | "- then\n", 175 | "``` \n", 176 | "echo 'python3 train_satellites.py\\\n", 177 | "\t--arch linknet34 --batch-size 12\\\n", 178 | "\t--imsize 1312 --preset mul_ps_vegetation --augs True\\\n", 179 | "\t--workers 6 --epochs 50 --start-epoch 0\\\n", 180 | "\t--seed 42 --print-freq 10\\\n", 181 | "\t--lr 1e-3 --optimizer adam\\\n", 182 | "\t--lognumber norm_ln34_mul_ps_vegetation_aug_dice_predict\\\n", 183 | "\t--predict --resume weights/norm_ln34_mul_ps_vegetation_aug_dice_best.pth.tar\\' > predict.sh\n", 184 | "```\n", 185 | "- `sh predict.sh`\n", 186 | "\n", 187 | "\n", 188 | "# 6 Creating graphs and submission files\n", 189 | "`cd` into `src` directory and execute `final_model_lstrs.py` script as follows:\n", 190 | "```\n", 191 | "docker exec -it YOUR_CONTAINER_ID sh -c \"cd path/to/src && python3 final_model_lstrs.py --folder norm_ln34_mul_ps_vegetation_aug_dice_predict\"\n", 192 | "```\n", 193 | "`folder` argument is for masks containing folder name, default is `norm_ln34_mul_ps_vegetation_aug_dice_predict`.\n", 194 | "Scipt saves a file called `norm_test.csv` into `../solutions` directory. The resulting file is used then as a submission file.\n", 195 | "\n", 196 | "# 7 Additional notes\n", 197 | "\n", 198 | "- You can run training and inference on the presets from `/src/presets.py`;\n", 199 | "- So the model can be evaluated on RGB-PS images and / or 8-channel images as well;\n", 200 | "- This script, for example will train an 8-channel model:\n", 201 | "```\n", 202 | "python3 train_satellites.py \\\n", 203 | "\t--arch linknet34 --batch-size 6 \\\n", 204 | "\t--imsize 1280 --preset mul_ps_vegetation --augs True \\\n", 205 | "\t--workers 6 --epochs 40 --start-epoch 0 \\\n", 206 | "\t--seed 42 --print-freq 20 \\\n", 207 | "\t--lr 1e-3 --optimizer adam \\\n", 208 | "\t--tensorboard True --lognumber ln34_mul_ps_vegetation_aug_dice\n", 209 | "```\n", 210 | "\n", 211 | "To train an 8-channel model you should also replace mean and std settings in the `src/SatellitesAugs.py`\n", 212 | "\n", 213 | "```\n", 214 | "# 8-channel settings\n", 215 | "# normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5],\n", 216 | "# std=[1, 1, 1, 1, 1, 1, 1, 1])\n", 217 | "\n", 218 | "# 3 channel settings\n", 219 | "normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n", 220 | " std=[0.229, 0.224, 0.225])\n", 221 | "```\n", 222 | "\n", 223 | "- 16-bit images are also supported:\n", 224 | "\n", 225 | "This snippet is commented in `src/SatellitesAugs.py`\n", 226 | "\n", 227 | "```\n", 228 | "# version compatible with 16-bit images\n", 229 | "\"\"\" \n", 230 | "class ImgAugAugs(object):\n", 231 | " def __call__(self,\n", 232 | " image):\n", 233 | " global seed \n", 234 | " \n", 235 | " # poor man's flipping\n", 236 | " if seed%2==0:\n", 237 | " image = np.fliplr(image)\n", 238 | " elif seed%4==0:\n", 239 | " image = np.fliplr(image)\n", 240 | " image = np.flipud(image)\n", 241 | " \n", 242 | " # poor man's affine transformations\n", 243 | " image = rotate(image,\n", 244 | " angle=seed,\n", 245 | " resize=False,\n", 246 | " clip=True,\n", 247 | " preserve_range=True) \n", 248 | "\n", 249 | " return image\n", 250 | "\"\"\"\n", 251 | "```\n", 252 | "\n", 253 | "- Also the following models are supported\n", 254 | " - `unet11` (VGG11 + Unet)\n", 255 | " - `linknet50` (ResNet50 + LinkNet, 3 layers)\n", 256 | " - `linknet50_full` (ResNet50 + LinkNet, 4 layers)\n", 257 | " - `linknext` (ResNext-101-32 + LinkNet, 4 layers)\n", 258 | " \n", 259 | "- Also in the repo you can find scripts to generate wide masks (i.e. wide roads have varying width) and layered masks (paved / non-paved). There are scripts in the `src/SatellitesDataset.py` that support that. They basically just replace some paths; \n", 260 | " \n", 261 | "# 8 Juputer notebooks\n", 262 | "\n", 263 | "Use these notebooks on your own risk!\n", 264 | "\n", 265 | "- `src/experiments.ipynb` - general debugging notebook with new models / generators / etc\n", 266 | "- `src/play_w_stuff.ipynb` - visualizing the solutions\n", 267 | "- `src/pipeline_experiments.ipynb`- some minor experiments with the graph creation script" 268 | ] 269 | }, 270 | { 271 | "cell_type": "code", 272 | "execution_count": null, 273 | "metadata": {}, 274 | "outputs": [], 275 | "source": [ 276 | "*" 277 | ] 278 | } 279 | ], 280 | "metadata": { 281 | "kernelspec": { 282 | "display_name": "Python 3", 283 | "language": "python", 284 | "name": "python3" 285 | }, 286 | "language_info": { 287 | "codemirror_mode": { 288 | "name": "ipython", 289 | "version": 3 290 | }, 291 | "file_extension": ".py", 292 | "mimetype": "text/x-python", 293 | "name": "python", 294 | "nbconvert_exporter": "python", 295 | "pygments_lexer": "ipython3", 296 | "version": "3.5.4" 297 | }, 298 | "toc": { 299 | "nav_menu": {}, 300 | "number_sections": true, 301 | "sideBar": true, 302 | "skip_h1_title": false, 303 | "toc_cell": false, 304 | "toc_position": {}, 305 | "toc_section_display": "block", 306 | "toc_window_display": false 307 | } 308 | }, 309 | "nbformat": 4, 310 | "nbformat_minor": 2 311 | } 312 | -------------------------------------------------------------------------------- /src/DilatedResnet.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import math 3 | import torch.utils.model_zoo as model_zoo 4 | import torch.nn.functional as F 5 | import torch 6 | 7 | nonlinearity = nn.ReLU 8 | 9 | class DecoderBlock(nn.Module): 10 | def __init__(self, in_channels, n_filters): 11 | super().__init__() 12 | 13 | # B, C, H, W -> B, C/4, H, W 14 | self.conv1 = nn.Conv2d(in_channels, in_channels // 4, 1) 15 | self.norm1 = nn.BatchNorm2d(in_channels // 4) 16 | self.relu1 = nonlinearity(inplace=True) 17 | 18 | # B, C/4, H, W -> B, C/4, H, W 19 | self.deconv2 = nn.ConvTranspose2d(in_channels // 4, in_channels // 4, 3, 20 | stride=2, padding=1, output_padding=1) 21 | self.norm2 = nn.BatchNorm2d(in_channels // 4) 22 | self.relu2 = nonlinearity(inplace=True) 23 | 24 | # B, C/4, H, W -> B, C, H, W 25 | self.conv3 = nn.Conv2d(in_channels // 4, n_filters, 1) 26 | self.norm3 = nn.BatchNorm2d(n_filters) 27 | self.relu3 = nonlinearity(inplace=True) 28 | 29 | def forward(self, x): 30 | x = self.conv1(x) 31 | x = self.norm1(x) 32 | x = self.relu1(x) 33 | x = self.deconv2(x) 34 | x = self.norm2(x) 35 | x = self.relu2(x) 36 | x = self.conv3(x) 37 | x = self.norm3(x) 38 | x = self.relu3(x) 39 | return x 40 | 41 | def conv3x3(in_planes, out_planes, stride=1, dilation=1, padding=1): 42 | """3x3 convolution with padding""" 43 | return nn.Conv2d(in_planes, 44 | out_planes, 45 | kernel_size=3, 46 | stride=stride, 47 | padding=padding, 48 | dilation=dilation, 49 | bias=False) 50 | 51 | class BasicBlock(nn.Module): 52 | expansion = 1 53 | 54 | def __init__(self, 55 | inplanes, 56 | planes, 57 | stride=1, 58 | downsample=None, 59 | dilation=1): 60 | super(BasicBlock, self).__init__() 61 | self.conv1 = conv3x3(inplanes, 62 | planes, 63 | stride, 64 | dilation=dilation, 65 | padding=dilation) 66 | self.bn1 = nn.BatchNorm2d(planes) 67 | self.relu = nn.ReLU(inplace=True) 68 | self.conv2 = conv3x3(planes, 69 | planes, 70 | dilation=dilation, 71 | padding=dilation) 72 | self.bn2 = nn.BatchNorm2d(planes) 73 | self.downsample = downsample 74 | self.stride = stride 75 | 76 | def forward(self, x): 77 | residual = x 78 | 79 | out = self.conv1(x) 80 | out = self.bn1(out) 81 | out = self.relu(out) 82 | 83 | out = self.conv2(out) 84 | out = self.bn2(out) 85 | 86 | if self.downsample is not None: 87 | residual = self.downsample(x) 88 | 89 | out += residual 90 | out = self.relu(out) 91 | 92 | return out 93 | 94 | class ResNet(nn.Module): 95 | 96 | def __init__(self, block, layers, dilation=1): 97 | 98 | self.inplanes = 64 99 | super(ResNet, self).__init__() 100 | self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, 101 | bias=False) 102 | self.bn1 = nn.BatchNorm2d(64) 103 | self.relu = nn.ReLU(inplace=True) 104 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 105 | self.layer1 = self._make_layer(block, 64, layers[0], dilation=dilation) 106 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2, dilation=dilation) 107 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2, dilation=dilation) 108 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2, dilation=dilation) 109 | 110 | for m in self.modules(): 111 | if isinstance(m, nn.Conv2d): 112 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 113 | m.weight.data.normal_(0, math.sqrt(2. / n)) 114 | elif isinstance(m, nn.BatchNorm2d): 115 | m.weight.data.fill_(1) 116 | m.bias.data.zero_() 117 | 118 | def _make_layer(self, 119 | block, 120 | planes, 121 | blocks, 122 | stride=1, 123 | dilation=1): 124 | downsample = None 125 | if stride != 1 or self.inplanes != planes * block.expansion: 126 | downsample = nn.Sequential( 127 | nn.Conv2d(self.inplanes, planes * block.expansion, 128 | kernel_size=1, stride=stride, bias=False), 129 | nn.BatchNorm2d(planes * block.expansion), 130 | ) 131 | 132 | layers = [] 133 | layers.append(block(self.inplanes, 134 | planes, 135 | stride, 136 | downsample, 137 | dilation=dilation)) 138 | self.inplanes = planes * block.expansion 139 | for i in range(1, blocks): 140 | layers.append(block(self.inplanes, planes, dilation=dilation)) 141 | 142 | return nn.Sequential(*layers) 143 | 144 | def forward(self, x): 145 | x = self.conv1(x) 146 | x = self.bn1(x) 147 | x = self.relu(x) 148 | x = self.maxpool(x) 149 | 150 | x = self.layer1(x) 151 | x = self.layer2(x) 152 | x = self.layer3(x) 153 | x = self.layer4(x) 154 | 155 | return x 156 | 157 | def dilated_resnet18(**kwargs): 158 | """Constructs a ResNet-18 model. 159 | Args: 160 | pretrained (bool): If True, returns a model pre-trained on ImageNet 161 | """ 162 | model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs) 163 | 164 | return model 165 | 166 | class GapNet18(nn.Module): 167 | def __init__(self, 168 | num_classes, 169 | num_channels=3, 170 | dilation=1): 171 | 172 | super().__init__() 173 | 174 | filters = [64*2, 128*2, 256*2, 512*2] 175 | resnet = dilated_resnet18(dilation=dilation) 176 | 177 | # self.firstconv = resnet.conv1 178 | # assert num_channels == 3, "num channels not used now. to use changle first conv layer to support num channels other then 3" 179 | # try to use 8-channels as first input 180 | if num_channels==3: 181 | self.firstconv = resnet.conv1 182 | else: 183 | self.firstconv = nn.Conv2d(num_channels, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3)) 184 | 185 | self.firstbn = resnet.bn1 186 | self.firstrelu = resnet.relu 187 | self.firstmaxpool = resnet.maxpool 188 | self.encoder1 = resnet.layer1 189 | self.encoder2 = resnet.layer2 190 | self.encoder3 = resnet.layer3 191 | self.encoder4 = resnet.layer4 192 | 193 | # Decoder 194 | self.decoder4 = DecoderBlock(filters[3], filters[2]) 195 | self.decoder3 = DecoderBlock(filters[2], filters[1]) 196 | self.decoder2 = DecoderBlock(filters[1], filters[0]) 197 | self.decoder1 = DecoderBlock(filters[0], filters[0]) 198 | 199 | # Final Classifier 200 | self.finaldeconv1 = nn.ConvTranspose2d(filters[0], 32, 3, stride=2) 201 | self.finalrelu1 = nonlinearity(inplace=True) 202 | self.finalconv2 = nn.Conv2d(32, 32, 3) 203 | self.finalrelu2 = nonlinearity(inplace=True) 204 | self.finalconv3 = nn.Conv2d(32, num_classes, 2, padding=1) 205 | 206 | # noinspection PyCallingNonCallable 207 | def forward(self, x1, x2): 208 | # Encoder 209 | x1 = self.firstconv(x1) 210 | x1 = self.firstbn(x1) 211 | x1 = self.firstrelu(x1) 212 | x1 = self.firstmaxpool(x1) 213 | e1_1 = self.encoder1(x1) 214 | e2_1 = self.encoder2(e1_1) 215 | e3_1 = self.encoder3(e2_1) 216 | e4_1 = self.encoder4(e3_1) 217 | 218 | x2 = self.firstconv(x2) 219 | x2 = self.firstbn(x2) 220 | x2 = self.firstrelu(x2) 221 | x2 = self.firstmaxpool(x2) 222 | e1_2 = self.encoder1(x2) 223 | e2_2 = self.encoder2(e1_2) 224 | e3_2 = self.encoder3(e2_2) 225 | e4_2 = self.encoder4(e3_2) 226 | 227 | e1 = torch.cat((e1_1,e1_2), dim=1) 228 | e2 = torch.cat((e2_1,e2_2), dim=1) 229 | e3 = torch.cat((e3_1,e3_2), dim=1) 230 | e4 = torch.cat((e4_1,e4_2), dim=1) 231 | 232 | # Decoder with Skip Connections 233 | d4 = self.decoder4(e4) + e3 234 | # d4 = e3 235 | d3 = self.decoder3(d4) + e2 236 | d2 = self.decoder2(d3) + e1 237 | d1 = self.decoder1(d2) 238 | 239 | # Final Classification 240 | f1 = self.finaldeconv1(d1) 241 | f2 = self.finalrelu1(f1) 242 | f3 = self.finalconv2(f2) 243 | f4 = self.finalrelu2(f3) 244 | f5 = self.finalconv3(f4) 245 | 246 | # return f5 247 | return F.sigmoid(f5) 248 | 249 | class GapNetImg18(nn.Module): 250 | def __init__(self, 251 | num_classes, 252 | num_channels=3, 253 | dilation=1): 254 | 255 | super().__init__() 256 | 257 | filters = [64*3, 128*3, 256*3, 512*3] 258 | resnet = dilated_resnet18(dilation=dilation) 259 | 260 | 261 | self.firstconv = nn.Conv2d(num_channels, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3)) 262 | self.firstconv_img = nn.Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3)) 263 | 264 | self.firstbn = resnet.bn1 265 | self.firstrelu = resnet.relu 266 | self.firstmaxpool = resnet.maxpool 267 | self.encoder1 = resnet.layer1 268 | self.encoder2 = resnet.layer2 269 | self.encoder3 = resnet.layer3 270 | self.encoder4 = resnet.layer4 271 | 272 | # Decoder 273 | self.decoder4 = DecoderBlock(filters[3], filters[2]) 274 | self.decoder3 = DecoderBlock(filters[2], filters[1]) 275 | self.decoder2 = DecoderBlock(filters[1], filters[0]) 276 | self.decoder1 = DecoderBlock(filters[0], filters[0]) 277 | 278 | # Final Classifier 279 | self.finaldeconv1 = nn.ConvTranspose2d(filters[0], 32, 3, stride=2) 280 | self.finalrelu1 = nonlinearity(inplace=True) 281 | self.finalconv2 = nn.Conv2d(32, 32, 3) 282 | self.finalrelu2 = nonlinearity(inplace=True) 283 | self.finalconv3 = nn.Conv2d(32, num_classes, 2, padding=1) 284 | 285 | # noinspection PyCallingNonCallable 286 | def forward(self, x, x1, x2): 287 | # Encoder 288 | x = self.firstconv_img(x) 289 | x = self.firstbn(x) 290 | x = self.firstrelu(x) 291 | x = self.firstmaxpool(x) 292 | e1_0 = self.encoder1(x) 293 | e2_0 = self.encoder2(e1_0) 294 | e3_0 = self.encoder3(e2_0) 295 | e4_0 = self.encoder4(e3_0) 296 | 297 | x1 = self.firstconv(x1) 298 | x1 = self.firstbn(x1) 299 | x1 = self.firstrelu(x1) 300 | x1 = self.firstmaxpool(x1) 301 | e1_1 = self.encoder1(x1) 302 | e2_1 = self.encoder2(e1_1) 303 | e3_1 = self.encoder3(e2_1) 304 | e4_1 = self.encoder4(e3_1) 305 | 306 | x2 = self.firstconv(x2) 307 | x2 = self.firstbn(x2) 308 | x2 = self.firstrelu(x2) 309 | x2 = self.firstmaxpool(x2) 310 | e1_2 = self.encoder1(x2) 311 | e2_2 = self.encoder2(e1_2) 312 | e3_2 = self.encoder3(e2_2) 313 | e4_2 = self.encoder4(e3_2) 314 | 315 | e1 = torch.cat((e1_0,e1_1,e1_2), dim=1) 316 | e2 = torch.cat((e2_0,e2_1,e2_2), dim=1) 317 | e3 = torch.cat((e3_0,e3_1,e3_2), dim=1) 318 | e4 = torch.cat((e4_0,e4_1,e4_2), dim=1) 319 | 320 | # Decoder with Skip Connections 321 | d4 = self.decoder4(e4) + e3 322 | # d4 = e3 323 | d3 = self.decoder3(d4) + e2 324 | d2 = self.decoder2(d3) + e1 325 | d1 = self.decoder1(d2) 326 | 327 | # Final Classification 328 | f1 = self.finaldeconv1(d1) 329 | f2 = self.finalrelu1(f1) 330 | f3 = self.finalconv2(f2) 331 | f4 = self.finalrelu2(f3) 332 | f5 = self.finalconv3(f4) 333 | 334 | # return f5 335 | return F.sigmoid(f5) -------------------------------------------------------------------------------- /src/InceptionResnetv2.py: -------------------------------------------------------------------------------- 1 | 2 | import torch 3 | import torch.nn as nn 4 | import torch.utils.model_zoo as model_zoo 5 | import os 6 | import sys 7 | 8 | __all__ = ['InceptionResNetV2', 'inceptionresnetv2'] 9 | 10 | pretrained_settings = { 11 | 'inceptionresnetv2': { 12 | 'imagenet': { 13 | 'url': 'http://data.lip6.fr/cadene/pretrainedmodels/inceptionresnetv2-520b38e4.pth', 14 | 'input_space': 'RGB', 15 | 'input_size': [3, 299, 299], 16 | 'input_range': [0, 1], 17 | 'mean': [0.5, 0.5, 0.5], 18 | 'std': [0.5, 0.5, 0.5], 19 | 'num_classes': 1000 20 | }, 21 | 'imagenet+background': { 22 | 'url': 'http://data.lip6.fr/cadene/pretrainedmodels/inceptionresnetv2-520b38e4.pth', 23 | 'input_space': 'RGB', 24 | 'input_size': [3, 299, 299], 25 | 'input_range': [0, 1], 26 | 'mean': [0.5, 0.5, 0.5], 27 | 'std': [0.5, 0.5, 0.5], 28 | 'num_classes': 1001 29 | } 30 | } 31 | } 32 | 33 | 34 | class BasicConv2d(nn.Module): 35 | 36 | def __init__(self, in_planes, out_planes, kernel_size, stride, padding=0): 37 | super(BasicConv2d, self).__init__() 38 | self.conv = nn.Conv2d(in_planes, out_planes, 39 | kernel_size=kernel_size, stride=stride, 40 | padding=padding, bias=False) # verify bias false 41 | self.bn = nn.BatchNorm2d(out_planes, 42 | eps=0.001, # value found in tensorflow 43 | momentum=0.1, # default pytorch value 44 | affine=True) 45 | self.relu = nn.ReLU(inplace=False) 46 | 47 | def forward(self, x): 48 | x = self.conv(x) 49 | x = self.bn(x) 50 | x = self.relu(x) 51 | return x 52 | 53 | 54 | class Mixed_5b(nn.Module): 55 | 56 | def __init__(self): 57 | super(Mixed_5b, self).__init__() 58 | 59 | self.branch0 = BasicConv2d(192, 96, kernel_size=1, stride=1) 60 | 61 | self.branch1 = nn.Sequential( 62 | BasicConv2d(192, 48, kernel_size=1, stride=1), 63 | BasicConv2d(48, 64, kernel_size=5, stride=1, padding=2) 64 | ) 65 | 66 | self.branch2 = nn.Sequential( 67 | BasicConv2d(192, 64, kernel_size=1, stride=1), 68 | BasicConv2d(64, 96, kernel_size=3, stride=1, padding=1), 69 | BasicConv2d(96, 96, kernel_size=3, stride=1, padding=1) 70 | ) 71 | 72 | self.branch3 = nn.Sequential( 73 | nn.AvgPool2d(3, stride=1, padding=1, count_include_pad=False), 74 | BasicConv2d(192, 64, kernel_size=1, stride=1) 75 | ) 76 | 77 | def forward(self, x): 78 | x0 = self.branch0(x) 79 | x1 = self.branch1(x) 80 | x2 = self.branch2(x) 81 | x3 = self.branch3(x) 82 | out = torch.cat((x0, x1, x2, x3), 1) 83 | return out 84 | 85 | 86 | class Block35(nn.Module): 87 | 88 | def __init__(self, scale=1.0): 89 | super(Block35, self).__init__() 90 | 91 | self.scale = scale 92 | 93 | self.branch0 = BasicConv2d(320, 32, kernel_size=1, stride=1) 94 | 95 | self.branch1 = nn.Sequential( 96 | BasicConv2d(320, 32, kernel_size=1, stride=1), 97 | BasicConv2d(32, 32, kernel_size=3, stride=1, padding=1) 98 | ) 99 | 100 | self.branch2 = nn.Sequential( 101 | BasicConv2d(320, 32, kernel_size=1, stride=1), 102 | BasicConv2d(32, 48, kernel_size=3, stride=1, padding=1), 103 | BasicConv2d(48, 64, kernel_size=3, stride=1, padding=1) 104 | ) 105 | 106 | self.conv2d = nn.Conv2d(128, 320, kernel_size=1, stride=1) 107 | self.relu = nn.ReLU(inplace=False) 108 | 109 | def forward(self, x): 110 | x0 = self.branch0(x) 111 | x1 = self.branch1(x) 112 | x2 = self.branch2(x) 113 | out = torch.cat((x0, x1, x2), 1) 114 | out = self.conv2d(out) 115 | out = out * self.scale + x 116 | out = self.relu(out) 117 | return out 118 | 119 | 120 | class Mixed_6a(nn.Module): 121 | 122 | def __init__(self): 123 | super(Mixed_6a, self).__init__() 124 | 125 | self.branch0 = BasicConv2d(320, 384, kernel_size=3, stride=2) 126 | 127 | self.branch1 = nn.Sequential( 128 | BasicConv2d(320, 256, kernel_size=1, stride=1), 129 | BasicConv2d(256, 256, kernel_size=3, stride=1, padding=1), 130 | BasicConv2d(256, 384, kernel_size=3, stride=2) 131 | ) 132 | 133 | self.branch2 = nn.MaxPool2d(3, stride=2) 134 | 135 | def forward(self, x): 136 | x0 = self.branch0(x) 137 | x1 = self.branch1(x) 138 | x2 = self.branch2(x) 139 | out = torch.cat((x0, x1, x2), 1) 140 | return out 141 | 142 | 143 | class Block17(nn.Module): 144 | 145 | def __init__(self, scale=1.0): 146 | super(Block17, self).__init__() 147 | 148 | self.scale = scale 149 | 150 | self.branch0 = BasicConv2d(1088, 192, kernel_size=1, stride=1) 151 | 152 | self.branch1 = nn.Sequential( 153 | BasicConv2d(1088, 128, kernel_size=1, stride=1), 154 | BasicConv2d(128, 160, kernel_size=(1,7), stride=1, padding=(0,3)), 155 | BasicConv2d(160, 192, kernel_size=(7,1), stride=1, padding=(3,0)) 156 | ) 157 | 158 | self.conv2d = nn.Conv2d(384, 1088, kernel_size=1, stride=1) 159 | self.relu = nn.ReLU(inplace=False) 160 | 161 | def forward(self, x): 162 | x0 = self.branch0(x) 163 | x1 = self.branch1(x) 164 | out = torch.cat((x0, x1), 1) 165 | out = self.conv2d(out) 166 | out = out * self.scale + x 167 | out = self.relu(out) 168 | return out 169 | 170 | 171 | class Mixed_7a(nn.Module): 172 | 173 | def __init__(self): 174 | super(Mixed_7a, self).__init__() 175 | 176 | self.branch0 = nn.Sequential( 177 | BasicConv2d(1088, 256, kernel_size=1, stride=1), 178 | BasicConv2d(256, 384, kernel_size=3, stride=2) 179 | ) 180 | 181 | self.branch1 = nn.Sequential( 182 | BasicConv2d(1088, 256, kernel_size=1, stride=1), 183 | BasicConv2d(256, 288, kernel_size=3, stride=2) 184 | ) 185 | 186 | self.branch2 = nn.Sequential( 187 | BasicConv2d(1088, 256, kernel_size=1, stride=1), 188 | BasicConv2d(256, 288, kernel_size=3, stride=1, padding=1), 189 | BasicConv2d(288, 320, kernel_size=3, stride=2) 190 | ) 191 | 192 | self.branch3 = nn.MaxPool2d(3, stride=2) 193 | 194 | def forward(self, x): 195 | x0 = self.branch0(x) 196 | x1 = self.branch1(x) 197 | x2 = self.branch2(x) 198 | x3 = self.branch3(x) 199 | out = torch.cat((x0, x1, x2, x3), 1) 200 | return out 201 | 202 | 203 | class Block8(nn.Module): 204 | 205 | def __init__(self, scale=1.0, noReLU=False): 206 | super(Block8, self).__init__() 207 | 208 | self.scale = scale 209 | self.noReLU = noReLU 210 | 211 | self.branch0 = BasicConv2d(2080, 192, kernel_size=1, stride=1) 212 | 213 | self.branch1 = nn.Sequential( 214 | BasicConv2d(2080, 192, kernel_size=1, stride=1), 215 | BasicConv2d(192, 224, kernel_size=(1,3), stride=1, padding=(0,1)), 216 | BasicConv2d(224, 256, kernel_size=(3,1), stride=1, padding=(1,0)) 217 | ) 218 | 219 | self.conv2d = nn.Conv2d(448, 2080, kernel_size=1, stride=1) 220 | if not self.noReLU: 221 | self.relu = nn.ReLU(inplace=False) 222 | 223 | def forward(self, x): 224 | x0 = self.branch0(x) 225 | x1 = self.branch1(x) 226 | out = torch.cat((x0, x1), 1) 227 | out = self.conv2d(out) 228 | out = out * self.scale + x 229 | if not self.noReLU: 230 | out = self.relu(out) 231 | return out 232 | 233 | 234 | class InceptionResNetV2(nn.Module): 235 | 236 | def __init__(self, num_classes=1001): 237 | super(InceptionResNetV2, self).__init__() 238 | # Special attributs 239 | self.input_space = None 240 | self.input_size = (299, 299, 3) 241 | self.mean = None 242 | self.std = None 243 | # Modules 244 | self.conv2d_1a = BasicConv2d(3, 32, kernel_size=3, stride=2) 245 | self.conv2d_2a = BasicConv2d(32, 32, kernel_size=3, stride=1) 246 | self.conv2d_2b = BasicConv2d(32, 64, kernel_size=3, stride=1, padding=1) 247 | self.maxpool_3a = nn.MaxPool2d(3, stride=2) 248 | self.conv2d_3b = BasicConv2d(64, 80, kernel_size=1, stride=1) 249 | self.conv2d_4a = BasicConv2d(80, 192, kernel_size=3, stride=1) 250 | self.maxpool_5a = nn.MaxPool2d(3, stride=2) 251 | self.mixed_5b = Mixed_5b() 252 | self.repeat = nn.Sequential( 253 | Block35(scale=0.17), 254 | Block35(scale=0.17), 255 | Block35(scale=0.17), 256 | Block35(scale=0.17), 257 | Block35(scale=0.17), 258 | Block35(scale=0.17), 259 | Block35(scale=0.17), 260 | Block35(scale=0.17), 261 | Block35(scale=0.17), 262 | Block35(scale=0.17) 263 | ) 264 | self.mixed_6a = Mixed_6a() 265 | self.repeat_1 = nn.Sequential( 266 | Block17(scale=0.10), 267 | Block17(scale=0.10), 268 | Block17(scale=0.10), 269 | Block17(scale=0.10), 270 | Block17(scale=0.10), 271 | Block17(scale=0.10), 272 | Block17(scale=0.10), 273 | Block17(scale=0.10), 274 | Block17(scale=0.10), 275 | Block17(scale=0.10), 276 | Block17(scale=0.10), 277 | Block17(scale=0.10), 278 | Block17(scale=0.10), 279 | Block17(scale=0.10), 280 | Block17(scale=0.10), 281 | Block17(scale=0.10), 282 | Block17(scale=0.10), 283 | Block17(scale=0.10), 284 | Block17(scale=0.10), 285 | Block17(scale=0.10) 286 | ) 287 | self.mixed_7a = Mixed_7a() 288 | self.repeat_2 = nn.Sequential( 289 | Block8(scale=0.20), 290 | Block8(scale=0.20), 291 | Block8(scale=0.20), 292 | Block8(scale=0.20), 293 | Block8(scale=0.20), 294 | Block8(scale=0.20), 295 | Block8(scale=0.20), 296 | Block8(scale=0.20), 297 | Block8(scale=0.20) 298 | ) 299 | self.block8 = Block8(noReLU=True) 300 | self.conv2d_7b = BasicConv2d(2080, 1536, kernel_size=1, stride=1) 301 | self.avgpool_1a = nn.AvgPool2d(8, count_include_pad=False) 302 | self.last_linear = nn.Linear(1536, num_classes) 303 | 304 | def features(self, input): 305 | x = self.conv2d_1a(input) 306 | x = self.conv2d_2a(x) 307 | x = self.conv2d_2b(x) 308 | x = self.maxpool_3a(x) 309 | x = self.conv2d_3b(x) 310 | x = self.conv2d_4a(x) 311 | x = self.maxpool_5a(x) 312 | x = self.mixed_5b(x) 313 | x = self.repeat(x) 314 | x = self.mixed_6a(x) 315 | x = self.repeat_1(x) 316 | x = self.mixed_7a(x) 317 | x = self.repeat_2(x) 318 | x = self.block8(x) 319 | x = self.conv2d_7b(x) 320 | return x 321 | 322 | def logits(self, features): 323 | x = self.avgpool_1a(features) 324 | x = x.view(x.size(0), -1) 325 | x = self.last_linear(x) 326 | return x 327 | 328 | def forward(self, input): 329 | x = self.features(input) 330 | x = self.logits(x) 331 | return x 332 | 333 | def inceptionresnetv2(num_classes=1001, pretrained='imagenet'): 334 | r"""InceptionResNetV2 model architecture from the 335 | `"InceptionV4, Inception-ResNet..." `_ paper. 336 | """ 337 | if pretrained: 338 | settings = pretrained_settings['inceptionresnetv2'][pretrained] 339 | assert num_classes == settings['num_classes'], \ 340 | "num_classes should be {}, but is {}".format(settings['num_classes'], num_classes) 341 | 342 | # both 'imagenet'&'imagenet+background' are loaded from same parameters 343 | model = InceptionResNetV2(num_classes=1001) 344 | model.load_state_dict(model_zoo.load_url(settings['url'])) 345 | 346 | if pretrained == 'imagenet': 347 | new_last_linear = nn.Linear(1536, 1000) 348 | new_last_linear.weight.data = model.last_linear.weight.data[1:] 349 | new_last_linear.bias.data = model.last_linear.bias.data[1:] 350 | model.last_linear = new_last_linear 351 | 352 | model.input_space = settings['input_space'] 353 | model.input_size = settings['input_size'] 354 | model.input_range = settings['input_range'] 355 | 356 | model.mean = settings['mean'] 357 | model.std = settings['std'] 358 | else: 359 | model = InceptionResNetV2(num_classes=num_classes) 360 | return model -------------------------------------------------------------------------------- /src/LRScheduler.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from torch.optim.optimizer import Optimizer 3 | 4 | class CyclicLR(object): 5 | """Sets the learning rate of each parameter group according to 6 | cyclical learning rate policy (CLR). The policy cycles the learning 7 | rate between two boundaries with a constant frequency, as detailed in 8 | the paper `Cyclical Learning Rates for Training Neural Networks`_. 9 | The distance between the two boundaries can be scaled on a per-iteration 10 | or per-cycle basis. 11 | Cyclical learning rate policy changes the learning rate after every batch. 12 | `batch_step` should be called after a batch has been used for training. 13 | To resume training, save `last_batch_iteration` and use it to instantiate `CycleLR`. 14 | This class has three built-in policies, as put forth in the paper: 15 | "triangular": 16 | A basic triangular cycle w/ no amplitude scaling. 17 | "triangular2": 18 | A basic triangular cycle that scales initial amplitude by half each cycle. 19 | "exp_range": 20 | A cycle that scales initial amplitude by gamma**(cycle iterations) at each 21 | cycle iteration. 22 | This implementation was adapted from the github repo: `bckenstler/CLR`_ 23 | Args: 24 | optimizer (Optimizer): Wrapped optimizer. 25 | base_lr (float or list): Initial learning rate which is the 26 | lower boundary in the cycle for eachparam groups. 27 | Default: 0.001 28 | max_lr (float or list): Upper boundaries in the cycle for 29 | each parameter group. Functionally, 30 | it defines the cycle amplitude (max_lr - base_lr). 31 | The lr at any cycle is the sum of base_lr 32 | and some scaling of the amplitude; therefore 33 | max_lr may not actually be reached depending on 34 | scaling function. Default: 0.006 35 | step_size (int): Number of training iterations per 36 | half cycle. Authors suggest setting step_size 37 | 2-8 x training iterations in epoch. Default: 2000 38 | mode (str): One of {triangular, triangular2, exp_range}. 39 | Values correspond to policies detailed above. 40 | If scale_fn is not None, this argument is ignored. 41 | Default: 'triangular' 42 | gamma (float): Constant in 'exp_range' scaling function: 43 | gamma**(cycle iterations) 44 | Default: 1.0 45 | scale_fn (function): Custom scaling policy defined by a single 46 | argument lambda function, where 47 | 0 <= scale_fn(x) <= 1 for all x >= 0. 48 | mode paramater is ignored 49 | Default: None 50 | scale_mode (str): {'cycle', 'iterations'}. 51 | Defines whether scale_fn is evaluated on 52 | cycle number or cycle iterations (training 53 | iterations since start of cycle). 54 | Default: 'cycle' 55 | last_batch_iteration (int): The index of the last batch. Default: -1 56 | Example: 57 | >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) 58 | >>> scheduler = torch.optim.CyclicLR(optimizer) 59 | >>> data_loader = torch.utils.data.DataLoader(...) 60 | >>> for epoch in range(10): 61 | >>> for batch in data_loader: 62 | >>> scheduler.batch_step() 63 | >>> train_batch(...) 64 | .. _Cyclical Learning Rates for Training Neural Networks: https://arxiv.org/abs/1506.01186 65 | .. _bckenstler/CLR: https://github.com/bckenstler/CLR 66 | """ 67 | 68 | def __init__(self, optimizer, base_lr=1e-3, max_lr=6e-3, 69 | step_size=2000, mode='triangular', gamma=1., 70 | scale_fn=None, scale_mode='cycle', last_batch_iteration=-1): 71 | 72 | if not isinstance(optimizer, Optimizer): 73 | raise TypeError('{} is not an Optimizer'.format( 74 | type(optimizer).__name__)) 75 | self.optimizer = optimizer 76 | 77 | if isinstance(base_lr, list) or isinstance(base_lr, tuple): 78 | if len(base_lr) != len(optimizer.param_groups): 79 | raise ValueError("expected {} base_lr, got {}".format( 80 | len(optimizer.param_groups), len(base_lr))) 81 | self.base_lrs = list(base_lr) 82 | else: 83 | self.base_lrs = [base_lr] * len(optimizer.param_groups) 84 | 85 | if isinstance(max_lr, list) or isinstance(max_lr, tuple): 86 | if len(max_lr) != len(optimizer.param_groups): 87 | raise ValueError("expected {} max_lr, got {}".format( 88 | len(optimizer.param_groups), len(max_lr))) 89 | self.max_lrs = list(max_lr) 90 | else: 91 | self.max_lrs = [max_lr] * len(optimizer.param_groups) 92 | 93 | self.step_size = step_size 94 | 95 | if mode not in ['triangular', 'triangular2', 'exp_range'] \ 96 | and scale_fn is None: 97 | raise ValueError('mode is invalid and scale_fn is None') 98 | 99 | self.mode = mode 100 | self.gamma = gamma 101 | 102 | if scale_fn is None: 103 | if self.mode == 'triangular': 104 | self.scale_fn = self._triangular_scale_fn 105 | self.scale_mode = 'cycle' 106 | elif self.mode == 'triangular2': 107 | self.scale_fn = self._triangular2_scale_fn 108 | self.scale_mode = 'cycle' 109 | elif self.mode == 'exp_range': 110 | self.scale_fn = self._exp_range_scale_fn 111 | self.scale_mode = 'iterations' 112 | else: 113 | self.scale_fn = scale_fn 114 | self.scale_mode = scale_mode 115 | 116 | self.batch_step(last_batch_iteration + 1) 117 | self.last_batch_iteration = last_batch_iteration 118 | 119 | def batch_step(self, batch_iteration=None): 120 | if batch_iteration is None: 121 | batch_iteration = self.last_batch_iteration + 1 122 | self.last_batch_iteration = batch_iteration 123 | for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()): 124 | param_group['lr'] = lr 125 | 126 | def _triangular_scale_fn(self, x): 127 | return 1. 128 | 129 | def _triangular2_scale_fn(self, x): 130 | return 1 / (2. ** (x - 1)) 131 | 132 | def _exp_range_scale_fn(self, x): 133 | return self.gamma**(x) 134 | 135 | def get_lr(self): 136 | step_size = float(self.step_size) 137 | cycle = np.floor(1 + self.last_batch_iteration / (2 * step_size)) 138 | x = np.abs(self.last_batch_iteration / step_size - 2 * cycle + 1) 139 | 140 | lrs = [] 141 | param_lrs = zip(self.optimizer.param_groups, self.base_lrs, self.max_lrs) 142 | for param_group, base_lr, max_lr in param_lrs: 143 | base_height = (max_lr - base_lr) * np.maximum(0, (1 - x)) 144 | if self.scale_mode == 'cycle': 145 | lr = base_lr + base_height * self.scale_fn(cycle) 146 | else: 147 | lr = base_lr + base_height * self.scale_fn(self.last_batch_iteration) 148 | lrs.append(lr) 149 | return lrs 150 | -------------------------------------------------------------------------------- /src/LinkNet.py: -------------------------------------------------------------------------------- 1 | """ 2 | all credits to ternaus and albu 3 | """ 4 | import torch 5 | import torch.nn as nn 6 | from torch.autograd import Variable 7 | from torchvision import models 8 | import torch.nn.functional as F 9 | from ResNeXt import resnext101_32x4d 10 | 11 | nonlinearity = nn.ReLU 12 | 13 | class DecoderBlock(nn.Module): 14 | def __init__(self, in_channels, n_filters): 15 | super().__init__() 16 | 17 | # B, C, H, W -> B, C/4, H, W 18 | self.conv1 = nn.Conv2d(in_channels, in_channels // 4, 1) 19 | self.norm1 = nn.BatchNorm2d(in_channels // 4) 20 | self.relu1 = nonlinearity(inplace=True) 21 | 22 | # B, C/4, H, W -> B, C/4, H, W 23 | self.deconv2 = nn.ConvTranspose2d(in_channels // 4, in_channels // 4, 3, 24 | stride=2, padding=1, output_padding=1) 25 | self.norm2 = nn.BatchNorm2d(in_channels // 4) 26 | self.relu2 = nonlinearity(inplace=True) 27 | 28 | # B, C/4, H, W -> B, C, H, W 29 | self.conv3 = nn.Conv2d(in_channels // 4, n_filters, 1) 30 | self.norm3 = nn.BatchNorm2d(n_filters) 31 | self.relu3 = nonlinearity(inplace=True) 32 | 33 | def forward(self, x): 34 | x = self.conv1(x) 35 | x = self.norm1(x) 36 | x = self.relu1(x) 37 | x = self.deconv2(x) 38 | x = self.norm2(x) 39 | x = self.relu2(x) 40 | x = self.conv3(x) 41 | x = self.norm3(x) 42 | x = self.relu3(x) 43 | return x 44 | 45 | class LinkNet34(nn.Module): 46 | def __init__(self, num_classes, num_channels=3): 47 | super().__init__() 48 | 49 | filters = [64, 128, 256, 512] 50 | resnet = models.resnet34(pretrained=True) 51 | 52 | # self.firstconv = resnet.conv1 53 | # assert num_channels == 3, "num channels not used now. to use changle first conv layer to support num channels other then 3" 54 | # try to use 8-channels as first input 55 | if num_channels==3: 56 | self.firstconv = resnet.conv1 57 | else: 58 | self.firstconv = nn.Conv2d(num_channels, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3)) 59 | 60 | self.firstbn = resnet.bn1 61 | self.firstrelu = resnet.relu 62 | self.firstmaxpool = resnet.maxpool 63 | self.encoder1 = resnet.layer1 64 | self.encoder2 = resnet.layer2 65 | self.encoder3 = resnet.layer3 66 | self.encoder4 = resnet.layer4 67 | 68 | # Decoder 69 | self.decoder4 = DecoderBlock(filters[3], filters[2]) 70 | self.decoder3 = DecoderBlock(filters[2], filters[1]) 71 | self.decoder2 = DecoderBlock(filters[1], filters[0]) 72 | self.decoder1 = DecoderBlock(filters[0], filters[0]) 73 | 74 | # Final Classifier 75 | self.finaldeconv1 = nn.ConvTranspose2d(filters[0], 32, 3, stride=2) 76 | self.finalrelu1 = nonlinearity(inplace=True) 77 | self.finalconv2 = nn.Conv2d(32, 32, 3) 78 | self.finalrelu2 = nonlinearity(inplace=True) 79 | self.finalconv3 = nn.Conv2d(32, num_classes, 2, padding=1) 80 | 81 | # noinspection PyCallingNonCallable 82 | def forward(self, x): 83 | # Encoder 84 | x = self.firstconv(x) 85 | x = self.firstbn(x) 86 | x = self.firstrelu(x) 87 | x = self.firstmaxpool(x) 88 | e1 = self.encoder1(x) 89 | e2 = self.encoder2(e1) 90 | e3 = self.encoder3(e2) 91 | e4 = self.encoder4(e3) 92 | 93 | # Decoder with Skip Connections 94 | d4 = self.decoder4(e4) + e3 95 | # d4 = e3 96 | d3 = self.decoder3(d4) + e2 97 | d2 = self.decoder2(d3) + e1 98 | d1 = self.decoder1(d2) 99 | 100 | # Final Classification 101 | f1 = self.finaldeconv1(d1) 102 | f2 = self.finalrelu1(f1) 103 | f3 = self.finalconv2(f2) 104 | f4 = self.finalrelu2(f3) 105 | f5 = self.finalconv3(f4) 106 | 107 | # return f5 108 | return F.sigmoid(f5) 109 | 110 | class LinkNet50(nn.Module): 111 | def __init__(self, num_classes, num_channels=3): 112 | super().__init__() 113 | 114 | filters = [256, 512, 1024] 115 | resnet = models.resnet50(pretrained=True) 116 | 117 | # self.firstconv = resnet.conv1 118 | # assert num_channels == 3, "num channels not used now. to use changle first conv layer to support num channels other then 3" 119 | # try to use 8-channels as first input 120 | if num_channels==3: 121 | self.firstconv = resnet.conv1 122 | else: 123 | self.firstconv = nn.Conv2d(num_channels, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3)) 124 | 125 | self.firstbn = resnet.bn1 126 | self.firstrelu = resnet.relu 127 | self.firstmaxpool = resnet.maxpool 128 | self.encoder1 = resnet.layer1 129 | self.encoder2 = resnet.layer2 130 | self.encoder3 = resnet.layer3 131 | 132 | # Decoder 133 | self.decoder3 = DecoderBlock(filters[2], filters[1]) 134 | self.decoder2 = DecoderBlock(filters[1], filters[0]) 135 | self.decoder1 = DecoderBlock(filters[0], filters[0]) 136 | 137 | # Final Classifier 138 | self.finaldeconv1 = nn.ConvTranspose2d(filters[0], 32, 3, stride=2) 139 | self.finalrelu1 = nonlinearity(inplace=True) 140 | self.finalconv2 = nn.Conv2d(32, 32, 3) 141 | self.finalrelu2 = nonlinearity(inplace=True) 142 | self.finalconv3 = nn.Conv2d(32, num_classes, 2, padding=1) 143 | 144 | # noinspection PyCallingNonCallable 145 | def forward(self, x): 146 | # Encoder 147 | x = self.firstconv(x) 148 | x = self.firstbn(x) 149 | x = self.firstrelu(x) 150 | x = self.firstmaxpool(x) 151 | e1 = self.encoder1(x) 152 | e2 = self.encoder2(e1) 153 | e3 = self.encoder3(e2) 154 | 155 | # Decoder with Skip Connections 156 | d3 = self.decoder3(e3) + e2 157 | d2 = self.decoder2(d3) + e1 158 | d1 = self.decoder1(d2) 159 | 160 | # Final Classification 161 | f1 = self.finaldeconv1(d1) 162 | f2 = self.finalrelu1(f1) 163 | f3 = self.finalconv2(f2) 164 | f4 = self.finalrelu2(f3) 165 | f5 = self.finalconv3(f4) 166 | 167 | return F.sigmoid(f5) 168 | 169 | class LinkNet50_full(nn.Module): 170 | def __init__(self, num_classes, num_channels=3): 171 | super().__init__() 172 | 173 | filters = [256, 512, 1024, 2048] 174 | resnet = models.resnet50(pretrained=True) 175 | 176 | # self.firstconv = resnet.conv1 177 | # assert num_channels == 3, "num channels not used now. to use changle first conv layer to support num channels other then 3" 178 | # try to use 8-channels as first input 179 | if num_channels==3: 180 | self.firstconv = resnet.conv1 181 | else: 182 | self.firstconv = nn.Conv2d(num_channels, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3)) 183 | 184 | self.firstbn = resnet.bn1 185 | self.firstrelu = resnet.relu 186 | self.firstmaxpool = resnet.maxpool 187 | self.encoder1 = resnet.layer1 188 | self.encoder2 = resnet.layer2 189 | self.encoder3 = resnet.layer3 190 | self.encoder4 = resnet.layer4 191 | 192 | # Decoder 193 | self.decoder4 = DecoderBlock(filters[3], filters[2]) 194 | self.decoder3 = DecoderBlock(filters[2], filters[1]) 195 | self.decoder2 = DecoderBlock(filters[1], filters[0]) 196 | self.decoder1 = DecoderBlock(filters[0], filters[0]) 197 | 198 | # Final Classifier 199 | self.finaldeconv1 = nn.ConvTranspose2d(filters[0], 32, 3, stride=2) 200 | self.finalrelu1 = nonlinearity(inplace=True) 201 | self.finalconv2 = nn.Conv2d(32, 32, 3) 202 | self.finalrelu2 = nonlinearity(inplace=True) 203 | self.finalconv3 = nn.Conv2d(32, num_classes, 2, padding=1) 204 | 205 | # noinspection PyCallingNonCallable 206 | def forward(self, x): 207 | # Encoder 208 | x = self.firstconv(x) 209 | x = self.firstbn(x) 210 | x = self.firstrelu(x) 211 | x = self.firstmaxpool(x) 212 | e1 = self.encoder1(x) 213 | e2 = self.encoder2(e1) 214 | e3 = self.encoder3(e2) 215 | e4 = self.encoder4(e3) 216 | 217 | # Decoder with Skip Connections 218 | d4 = self.decoder4(e4) + e3 219 | d3 = self.decoder3(d4) + e2 220 | d2 = self.decoder2(d3) + e1 221 | d1 = self.decoder1(d2) 222 | 223 | # Final Classification 224 | f1 = self.finaldeconv1(d1) 225 | f2 = self.finalrelu1(f1) 226 | f3 = self.finalconv2(f2) 227 | f4 = self.finalrelu2(f3) 228 | f5 = self.finalconv3(f4) 229 | 230 | return F.sigmoid(f5) 231 | 232 | class LinkNeXt(nn.Module): 233 | def __init__(self, num_classes, num_channels=3): 234 | super().__init__() 235 | 236 | filters = [256, 512, 1024, 2048] 237 | resnet = resnext101_32x4d(num_classes=1000, pretrained='imagenet') 238 | 239 | self.stem = resnet.stem 240 | self.encoder1 = resnet.layer1 241 | self.encoder2 = resnet.layer2 242 | self.encoder3 = resnet.layer3 243 | self.encoder4 = resnet.layer4 244 | 245 | # Decoder 246 | self.decoder4 = DecoderBlock(filters[3], filters[2]) 247 | self.decoder3 = DecoderBlock(filters[2], filters[1]) 248 | self.decoder2 = DecoderBlock(filters[1], filters[0]) 249 | self.decoder1 = DecoderBlock(filters[0], filters[0]) 250 | 251 | # Final Classifier 252 | self.finaldeconv1 = nn.ConvTranspose2d(filters[0], 32, 3, stride=2) 253 | self.finalrelu1 = nonlinearity(inplace=True) 254 | self.finalconv2 = nn.Conv2d(32, 32, 3) 255 | self.finalrelu2 = nonlinearity(inplace=True) 256 | self.finalconv3 = nn.Conv2d(32, num_classes, 2, padding=1) 257 | 258 | # noinspection PyCallingNonCallable 259 | def forward(self, x): 260 | # Encoder 261 | x = self.stem(x) 262 | e1 = self.encoder1(x) 263 | e2 = self.encoder2(e1) 264 | e3 = self.encoder3(e2) 265 | e4 = self.encoder4(e3) 266 | 267 | # Decoder with Skip Connections 268 | d4 = self.decoder4(e4) + e3 269 | # d4 = e3 270 | d3 = self.decoder3(d4) + e2 271 | d2 = self.decoder2(d3) + e1 272 | d1 = self.decoder1(d2) 273 | 274 | # Final Classification 275 | f1 = self.finaldeconv1(d1) 276 | f2 = self.finalrelu1(f1) 277 | f3 = self.finalconv2(f2) 278 | f4 = self.finalrelu2(f3) 279 | f5 = self.finalconv3(f4) 280 | 281 | # return f5 282 | return F.sigmoid(f5) -------------------------------------------------------------------------------- /src/Loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | class TDiceLoss(nn.Module): 6 | def __init__(self, dice_weight=1): 7 | super().__init__() 8 | self.nll_loss = nn.BCELoss() 9 | self.dice_weight = dice_weight 10 | 11 | def forward(self, outputs, targets): 12 | loss = self.nll_loss(outputs, targets) 13 | if self.dice_weight: 14 | eps = 1e-15 15 | dice_target = (targets == 1).float() 16 | dice_output = outputs 17 | intersection = (dice_output * dice_target).sum() 18 | union = dice_output.sum() + dice_target.sum() + eps 19 | loss += 1 - torch.log(2 * intersection / union) 20 | 21 | return loss 22 | 23 | def dice_loss(preds, trues, weight=None, is_average=True): 24 | num = preds.size(0) 25 | preds = preds.view(num, -1) 26 | trues = trues.view(num, -1) 27 | if weight is not None: 28 | w = torch.autograd.Variable(weight).view(num, -1) 29 | preds = preds * w 30 | trues = trues * w 31 | intersection = (preds * trues).sum(1) 32 | scores = 2. * (intersection + 1) / (preds.sum(1) + trues.sum(1) + 1) 33 | 34 | if is_average: 35 | score = scores.sum()/num 36 | return torch.clamp(score, 0., 1.) 37 | else: 38 | return scores 39 | 40 | def dice_clamp(preds, 41 | trues, 42 | is_average=True): 43 | preds = torch.round(preds) 44 | return dice_loss(preds, trues, is_average=is_average) 45 | 46 | class DiceLoss(nn.Module): 47 | def __init__(self, size_average=True): 48 | super().__init__() 49 | self.size_average = size_average 50 | 51 | def forward(self, input, target, weight=None): 52 | return 1-dice_loss(F.sigmoid(input), target, weight=weight, is_average=self.size_average) 53 | 54 | class BCEDiceLoss(nn.Module): 55 | def __init__(self, size_average=True): 56 | super().__init__() 57 | self.size_average = size_average 58 | self.dice = DiceLoss(size_average=size_average) 59 | 60 | def forward(self, input, target, weight=None): 61 | return nn.modules.loss.BCEWithLogitsLoss(size_average=self.size_average, weight=weight)(input, target) + self.dice(input, target, weight=weight) 62 | -------------------------------------------------------------------------------- /src/MaskUtils.py: -------------------------------------------------------------------------------- 1 | import rasterio 2 | import ast 3 | import numpy as np 4 | from skimage.io import imread 5 | from skimage import img_as_ubyte,img_as_float 6 | from skimage import exposure 7 | from skimage.draw import circle 8 | import cv2 9 | from collections import Counter 10 | import matplotlib.pyplot as plt 11 | from collections import Sequence 12 | from itertools import chain, count 13 | import os 14 | import pandas as pd 15 | 16 | data_prefix = '../data' 17 | geojson_df = pd.read_csv('geojson_df_full.csv') 18 | meta_df = pd.read_csv('../metadata.csv') 19 | 20 | # manual curation of the dataset 21 | # all unpaved roads with 2+ lanes set to paved 22 | geojson_df.loc[(geojson_df.lane_number>2)&(geojson_df.paved == 2),'paved'] = 1 23 | 24 | # a function to understand list depth 25 | def depth(seq): 26 | for level in count(): 27 | if not seq: 28 | return level 29 | seq = list(chain.from_iterable(s for s in seq if isinstance(s, Sequence))) 30 | 31 | def read_image(preset,path): 32 | img = imread(path) 33 | target_channels = np.zeros(shape=(preset['width'],preset['width'],len(preset['channels']))) 34 | 35 | # expand grayscale images to 3 dimensions 36 | if len(img.shape)<3: 37 | img = np.expand_dims(img, 2) 38 | 39 | for i,channel in enumerate(preset['channels']): 40 | target_channels[:,:,i] = img[:,:,channel-1] 41 | 42 | # target_channels = img_as_ubyte(target_channels) 43 | # target_channels = exposure.rescale_intensity(target_channels, in_range='uint8') 44 | 45 | return target_channels 46 | 47 | def draw_mask(circle_size, 48 | line_width, 49 | ls_list, 50 | mask_size 51 | ): 52 | 53 | mask = np.zeros((mask_size, mask_size), dtype=np.uint8) 54 | all_points = [] 55 | 56 | for line in ls_list: 57 | 58 | points_xy = line 59 | 60 | all_points.extend(points_xy) 61 | 62 | for i,[x,y] in enumerate(points_xy): 63 | if i-1>-1: 64 | prev_x = int(float(points_xy[i-1][0])) 65 | prev_y = int(float(points_xy[i-1][1])) 66 | mask = cv2.line(mask,(prev_x,prev_y),(int(float(x)),int(float(y))),(150),line_width) 67 | 68 | all_points_text = [(str(point[0])+' '+str(point[1])) for point in all_points] 69 | count_dict = Counter(all_points_text) 70 | 71 | for key, value in count_dict.items(): 72 | if(value>1): 73 | x,y = key.split() 74 | rr, cc = circle(int(float(y)), int(float(x)), circle_size) 75 | mask[rr.clip(min=0,max=mask_size-1), cc.clip(min=0,max=mask_size-1)] = 255 76 | return mask 77 | 78 | def draw_mask_width( 79 | ls_list, 80 | mask_size 81 | ): 82 | 83 | width_per_lane = 10 84 | mask = np.zeros((mask_size, mask_size), dtype=np.uint8) 85 | all_points = [] 86 | 87 | for line in ls_list: 88 | 89 | points_xy = line 90 | all_points.extend(points_xy) 91 | 92 | for i,[x,y,lanes] in enumerate(points_xy): 93 | if i-1>-1: 94 | prev_x = int(float(points_xy[i-1][0])) 95 | prev_y = int(float(points_xy[i-1][1])) 96 | mask = cv2.line(mask,(prev_x,prev_y), 97 | (int(float(x)),int(float(y))), 98 | (255), 99 | lanes*width_per_lane) 100 | return mask 101 | 102 | def draw_intersections(circle_size, 103 | ls_list, 104 | mask_size 105 | ): 106 | 107 | mask = np.zeros((mask_size, mask_size), dtype=np.uint8) 108 | all_points = [] 109 | 110 | for line in ls_list: 111 | points_xy = line 112 | all_points.extend(points_xy) 113 | all_points_text = [(str(point[0])+' '+str(point[1])) for point in all_points] 114 | count_dict = Counter(all_points_text) 115 | 116 | for key, value in count_dict.items(): 117 | if(value>1): 118 | x,y = key.split() 119 | rr, cc = circle(int(float(y)), int(float(x)), circle_size) 120 | mask[rr.clip(min=0,max=mask_size-1), cc.clip(min=0,max=mask_size-1)] = 255 121 | return mask 122 | 123 | def draw_masks(paved = 2, 124 | road_type = 6, 125 | lane_number = 6): 126 | 127 | random_image = geojson_df[(geojson_df.paved == paved) 128 | # &(geojson_df.road_type == road_type) 129 | &(geojson_df.lane_number == lane_number)].sample(n=1).img_id.values[0] 130 | 131 | sample_df = geojson_df[(geojson_df.paved == paved) 132 | # &(geojson_df.road_type == road_type) 133 | &(geojson_df.lane_number == lane_number) 134 | &(geojson_df.img_id == random_image)] 135 | 136 | rgb_ps_image = os.path.join(data_prefix, meta_df[(meta_df.img_subfolders.str.contains(random_image)) 137 | &(meta_df.img_folders=='RGB-PanSharpen')].img_files.values[0], 138 | 'RGB-PanSharpen', 139 | meta_df[(meta_df.img_subfolders.str.contains(random_image)) 140 | &(meta_df.img_folders=='RGB-PanSharpen')].img_subfolders.values[0]) 141 | 142 | mask2_path = os.path.join(data_prefix, meta_df[(meta_df.img_subfolders.str.contains(random_image)) 143 | &(meta_df.img_folders=='RGB-PanSharpen')].img_files.values[0], 144 | 'RGB-PanSharpen_mask', 145 | meta_df[(meta_df.img_subfolders.str.contains(random_image)) 146 | &(meta_df.img_folders=='RGB-PanSharpen')].img_subfolders.values[0])[:-3]+'jpg' 147 | 148 | 149 | src = rasterio.open(rgb_ps_image) 150 | ls_list = [ast.literal_eval(ls) for ls in sample_df.linestring.values] 151 | ls_list_image = [] 152 | 153 | for line in ls_list: 154 | points = [] 155 | for point in line: 156 | points.append(~src.affine * (point[0],point[1])) 157 | ls_list_image.append(points) 158 | 159 | img = read_image(preset_dict['rgb_ps'],rgb_ps_image) 160 | 161 | mask = draw_mask(circle_size=15, 162 | line_width=15, 163 | ls_list=ls_list_image, 164 | mask_size=1300 165 | ) 166 | 167 | mask2 = imread(mask2_path) 168 | 169 | fig=plt.figure(figsize=(20, 6)) 170 | 171 | fig.add_subplot(1, 3, 1) 172 | img += -img.min() 173 | img *= (1/img.max()) 174 | plt.imshow(img) 175 | 176 | fig.add_subplot(1, 3, 2) 177 | plt.imshow(mask) 178 | 179 | fig.add_subplot(1, 3, 3) 180 | plt.imshow(mask2) 181 | 182 | plt.show() 183 | 184 | def draw_masks_analyze_road_types(): 185 | road_types = [3,5,6,7] 186 | 187 | # select a random image and the gt 188 | random_image = geojson_df.sample(n=1).img_id.values[0] 189 | 190 | mask2_path = os.path.join(data_prefix, meta_df[(meta_df.img_subfolders.str.contains(random_image)) 191 | &(meta_df.img_folders=='RGB-PanSharpen')].img_files.values[0], 192 | 'RGB-PanSharpen_mask', 193 | meta_df[(meta_df.img_subfolders.str.contains(random_image)) 194 | &(meta_df.img_folders=='RGB-PanSharpen')].img_subfolders.values[0])[:-3]+'jpg' 195 | 196 | mask2 = imread(mask2_path) 197 | 198 | rgb_ps_image = os.path.join(data_prefix, meta_df[(meta_df.img_subfolders.str.contains(random_image)) 199 | &(meta_df.img_folders=='RGB-PanSharpen')].img_files.values[0], 200 | 'RGB-PanSharpen', 201 | meta_df[(meta_df.img_subfolders.str.contains(random_image)) 202 | &(meta_df.img_folders=='RGB-PanSharpen')].img_subfolders.values[0]) 203 | 204 | fig=plt.figure(figsize=(25, 10)) 205 | 206 | img = read_image(preset_dict['rgb_ps'],rgb_ps_image) 207 | fig.add_subplot(2, 5, 1) 208 | img += -img.min() 209 | img *= (1/img.max()) 210 | plt.imshow(img) 211 | 212 | fig.add_subplot(2, 5, 6) 213 | plt.imshow(mask2) 214 | 215 | for i,road_type in enumerate(road_types): 216 | 217 | sample_df = geojson_df[((geojson_df.road_type == road_type) 218 | &(geojson_df.img_id == random_image))] 219 | 220 | src = rasterio.open(rgb_ps_image) 221 | ls_list = [ast.literal_eval(ls) for ls in sample_df.linestring.values] 222 | ls_list_image = [] 223 | 224 | for line in ls_list: 225 | points = [] 226 | for point in line: 227 | points.append(~src.affine * (point[0],point[1])) 228 | ls_list_image.append(points) 229 | 230 | mask = draw_mask(circle_size=15, 231 | line_width=15, 232 | ls_list=ls_list_image, 233 | mask_size=1300 234 | ) 235 | 236 | fig.add_subplot(2, 5, i+2) 237 | plt.imshow(mask) 238 | 239 | plt.show() 240 | 241 | def process_ls(sample_df,lane_numbers): 242 | if lane_numbers == None: 243 | ls_list = [] 244 | for ls in sample_df.linestring.values: 245 | if depth(ast.literal_eval(ls)) == 2: 246 | ls_list.append(ast.literal_eval(ls)) 247 | elif depth(ast.literal_eval(ls)) == 3: 248 | for item in ast.literal_eval(ls): 249 | ls_list.append(item) 250 | else: 251 | raise ValueError('Wrong linestring format') 252 | return ls_list,None 253 | else: 254 | ls_list = [] 255 | lane_numbers_new = [] 256 | 257 | i = 0 258 | for ls in sample_df.linestring.values: 259 | if depth(ast.literal_eval(ls)) == 2: 260 | ls_list.append(ast.literal_eval(ls)) 261 | lane_numbers_new.append(lane_numbers[i]) 262 | elif depth(ast.literal_eval(ls)) == 3: 263 | for item in ast.literal_eval(ls): 264 | ls_list.append(item) 265 | lane_numbers_new.append(lane_numbers[i]) 266 | else: 267 | raise ValueError('Wrong linestring format') 268 | i+=1 269 | return ls_list,lane_numbers_new 270 | 271 | def draw_masks_new(): 272 | paved_types = [1,2] 273 | 274 | # select a random image and the gt 275 | random_image = geojson_df.sample(n=1).img_id.values[0] 276 | # random_image = 'AOI_4_Shanghai_img823' 277 | 278 | mask2_path = os.path.join(data_prefix, meta_df[(meta_df.img_subfolders.str.contains(random_image)) 279 | &(meta_df.img_folders=='RGB-PanSharpen')].img_files.values[0], 280 | 'RGB-PanSharpen_mask', 281 | meta_df[(meta_df.img_subfolders.str.contains(random_image)) 282 | &(meta_df.img_folders=='RGB-PanSharpen')].img_subfolders.values[0])[:-3]+'jpg' 283 | 284 | mask2 = imread(mask2_path) 285 | 286 | rgb_ps_image = os.path.join(data_prefix, meta_df[(meta_df.img_subfolders.str.contains(random_image)) 287 | &(meta_df.img_folders=='RGB-PanSharpen')].img_files.values[0], 288 | 'RGB-PanSharpen', 289 | meta_df[(meta_df.img_subfolders.str.contains(random_image)) 290 | &(meta_df.img_folders=='RGB-PanSharpen')].img_subfolders.values[0]) 291 | 292 | fig=plt.figure(figsize=(20, 15)) 293 | 294 | img = read_image(preset_dict['rgb_ps'],rgb_ps_image) 295 | fig.add_subplot(2, 3, 1) 296 | img += -img.min() 297 | img *= (1/img.max()) 298 | plt.imshow(img) 299 | 300 | fig.add_subplot(2, 3, 5) 301 | plt.imshow(mask2) 302 | 303 | # draw intersections only 304 | sample_df = geojson_df[((geojson_df.img_id == random_image))] 305 | 306 | src = rasterio.open(rgb_ps_image) 307 | lane_numbers = None 308 | ls_list,lane_numbers = process_ls(sample_df,lane_numbers) 309 | ls_list_image = [] 310 | 311 | for line in ls_list: 312 | points = [] 313 | for point in line: 314 | points.append(~src.affine * (point[0],point[1])) 315 | ls_list_image.append(points) 316 | 317 | intersections_mask = draw_intersections(circle_size=15, 318 | ls_list=ls_list_image, 319 | mask_size=1300) 320 | 321 | 322 | fig.add_subplot(2, 3, 6) 323 | plt.imshow(intersections_mask) 324 | 325 | for i,paved_type in enumerate(paved_types): 326 | 327 | sample_df = geojson_df[((geojson_df.paved == paved_type) 328 | &(geojson_df.img_id == random_image))] 329 | 330 | lane_numbers = list(sample_df.lane_number.values) 331 | 332 | src = rasterio.open(rgb_ps_image) 333 | ls_list,lane_numbers = process_ls(sample_df,lane_numbers) 334 | ls_list_image = [] 335 | 336 | for j,line in enumerate(ls_list): 337 | points = [] 338 | for point in line: 339 | # create pixel coordinates 340 | # add lane width to the tuple 341 | pixel_coordinates = ~src.affine * (point[0],point[1]) 342 | points.append(pixel_coordinates + (lane_numbers[j],)) 343 | ls_list_image.append(points) 344 | 345 | mask = draw_mask_width( 346 | ls_list=ls_list_image, 347 | mask_size=1300 348 | ) 349 | 350 | fig.add_subplot(2, 3, i+2) 351 | plt.imshow(mask) 352 | 353 | plt.show() 354 | 355 | def create_new_masks(random_image): 356 | paved_types = [1,2] 357 | 358 | rgb_ps_image = os.path.join(data_prefix, meta_df[(meta_df.img_subfolders.str.contains(random_image)) 359 | &(meta_df.img_folders=='PAN')].img_files.values[0], 360 | 'PAN', 361 | meta_df[(meta_df.img_subfolders.str.contains(random_image)) 362 | &(meta_df.img_folders=='PAN')].img_subfolders.values[0]) 363 | 364 | # draw intersections only 365 | sample_df = geojson_df[((geojson_df.img_id == random_image))] 366 | 367 | 368 | src = rasterio.open(rgb_ps_image) 369 | lane_numbers = None 370 | ls_list,lane_numbers = process_ls(sample_df,lane_numbers) 371 | ls_list_image = [] 372 | 373 | for line in ls_list: 374 | points = [] 375 | for point in line: 376 | points.append(~src.affine * (point[0],point[1])) 377 | ls_list_image.append(points) 378 | 379 | intersections_mask = draw_intersections(circle_size=15, 380 | ls_list=ls_list_image, 381 | mask_size=1300) 382 | 383 | road_masks = [] 384 | for i,paved_type in enumerate(paved_types): 385 | 386 | sample_df = geojson_df[((geojson_df.paved == paved_type) 387 | &(geojson_df.img_id == random_image))] 388 | 389 | lane_numbers = list(sample_df.lane_number.values) 390 | 391 | src = rasterio.open(rgb_ps_image) 392 | ls_list,lane_numbers = process_ls(sample_df,lane_numbers) 393 | ls_list_image = [] 394 | 395 | for j,line in enumerate(ls_list): 396 | points = [] 397 | for point in line: 398 | # create pixel coordinates 399 | # add lane width to the tuple 400 | pixel_coordinates = ~src.affine * (point[0],point[1]) 401 | points.append(pixel_coordinates + (lane_numbers[j],)) 402 | ls_list_image.append(points) 403 | 404 | road_masks.append(draw_mask_width( 405 | ls_list=ls_list_image, 406 | mask_size=1300 407 | )) 408 | return intersections_mask,road_masks[0],road_masks[1] -------------------------------------------------------------------------------- /src/ResNeXt.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | import torch.nn as nn 4 | import torch.utils.model_zoo as model_zoo 5 | from resnext_features.resnext101_32x4d_features import resnext101_32x4d_features,resnext101_32x4d_features_blob 6 | from resnext_features import resnext101_64x4d_features 7 | 8 | __all__ = ['ResNeXt101_32x4d', 'resnext101_32x4d', 9 | 'ResNeXt101_64x4d', 'resnext101_64x4d'] 10 | 11 | pretrained_settings = { 12 | 'resnext101_32x4d': { 13 | 'imagenet': { 14 | 'url': 'http://data.lip6.fr/cadene/pretrainedmodels/resnext101_32x4d-29e315fa.pth', 15 | 'input_space': 'RGB', 16 | 'input_size': [3, 224, 224], 17 | 'input_range': [0, 1], 18 | 'mean': [0.485, 0.456, 0.406], 19 | 'std': [0.229, 0.224, 0.225], 20 | 'num_classes': 1000 21 | } 22 | }, 23 | 'resnext101_64x4d': { 24 | 'imagenet': { 25 | 'url': 'http://data.lip6.fr/cadene/pretrainedmodels/resnext101_64x4d-e77a0586.pth', 26 | 'input_space': 'RGB', 27 | 'input_size': [3, 224, 224], 28 | 'input_range': [0, 1], 29 | 'mean': [0.485, 0.456, 0.406], 30 | 'std': [0.229, 0.224, 0.225], 31 | 'num_classes': 1000 32 | } 33 | } 34 | } 35 | 36 | class ResNeXt101_32x4d_blob(nn.Module): 37 | 38 | def __init__(self, num_classes=1000): 39 | super(ResNeXt101_32x4d_blob, self).__init__() 40 | self.num_classes = num_classes 41 | 42 | resnext = resnext101_32x4d_features_blob() 43 | 44 | self.features = resnext.resnext101_32x4d_features 45 | self.avg_pool = nn.AvgPool2d((7, 7), (1, 1)) 46 | self.last_linear = nn.Linear(2048, num_classes) 47 | 48 | def logits(self, input): 49 | x = self.avg_pool(input) 50 | x = x.view(x.size(0), -1) 51 | x = self.last_linear(x) 52 | return x 53 | 54 | def forward(self, input): 55 | x = self.features(input) 56 | x = self.logits(x) 57 | return x 58 | 59 | class ResNeXt101_32x4d(nn.Module): 60 | 61 | def __init__(self, num_classes=1000): 62 | super(ResNeXt101_32x4d, self).__init__() 63 | self.num_classes = num_classes 64 | 65 | resnext = resnext101_32x4d_features() 66 | 67 | self.stem = resnext.resnext101_32x4d_stem 68 | self.layer1 = resnext.resnext101_32x4d_layer1 69 | self.layer2 = resnext.resnext101_32x4d_layer2 70 | self.layer3 = resnext.resnext101_32x4d_layer3 71 | self.layer4 = resnext.resnext101_32x4d_layer4 72 | 73 | self.avg_pool = nn.AvgPool2d((7, 7), (1, 1)) 74 | self.last_linear = nn.Linear(2048, num_classes) 75 | 76 | def logits(self, input): 77 | x = self.avg_pool(input) 78 | x = x.view(x.size(0), -1) 79 | x = self.last_linear(x) 80 | return x 81 | 82 | def forward(self, input): 83 | x = self.stem(input) 84 | x = self.layer1(x) 85 | x = self.layer2(x) 86 | x = self.layer3(x) 87 | x = self.layer4(x) 88 | x = self.logits(x) 89 | return x 90 | 91 | class ResNeXt101_64x4d(nn.Module): 92 | 93 | def __init__(self, num_classes=1000): 94 | super(ResNeXt101_64x4d, self).__init__() 95 | self.num_classes = num_classes 96 | self.features = resnext101_64x4d_features 97 | self.avg_pool = nn.AvgPool2d((7, 7), (1, 1)) 98 | self.last_linear = nn.Linear(2048, num_classes) 99 | 100 | def logits(self, input): 101 | x = self.avg_pool(input) 102 | x = x.view(x.size(0), -1) 103 | x = self.last_linear(x) 104 | return x 105 | 106 | def forward(self, input): 107 | x = self.features(input) 108 | x = self.logits(x) 109 | return x 110 | 111 | def resnext101_32x4d(num_classes=1000, pretrained='imagenet'): 112 | model = ResNeXt101_32x4d(num_classes=num_classes) 113 | model_blob = ResNeXt101_32x4d_blob(num_classes=num_classes) 114 | if pretrained is not None: 115 | settings = pretrained_settings['resnext101_32x4d'][pretrained] 116 | assert num_classes == settings['num_classes'], \ 117 | "num_classes should be {}, but is {}".format(settings['num_classes'], num_classes) 118 | model_blob.load_state_dict(model_zoo.load_url(settings['url'])) 119 | 120 | model.stem = nn.Sequential( 121 | model_blob.features[0], 122 | model_blob.features[1], 123 | model_blob.features[2], 124 | model_blob.features[3], 125 | ) 126 | 127 | model.layer1 = nn.Sequential( 128 | model_blob.features[4], 129 | ) 130 | model.layer2 = nn.Sequential( 131 | model_blob.features[5], 132 | ) 133 | model.layer3 = nn.Sequential( 134 | model_blob.features[6], 135 | ) 136 | model.layer4 = nn.Sequential( 137 | model_blob.features[7], 138 | ) 139 | # finish here 140 | 141 | model.input_space = settings['input_space'] 142 | model.input_size = settings['input_size'] 143 | model.input_range = settings['input_range'] 144 | model.mean = settings['mean'] 145 | model.std = settings['std'] 146 | return model 147 | 148 | def resnext101_64x4d(num_classes=1000, pretrained='imagenet'): 149 | model = ResNeXt101_64x4d(num_classes=num_classes) 150 | if pretrained is not None: 151 | settings = pretrained_settings['resnext101_64x4d'][pretrained] 152 | assert num_classes == settings['num_classes'], \ 153 | "num_classes should be {}, but is {}".format(settings['num_classes'], num_classes) 154 | model.load_state_dict(model_zoo.load_url(settings['url'])) 155 | model.input_space = settings['input_space'] 156 | model.input_size = settings['input_size'] 157 | model.input_range = settings['input_range'] 158 | model.mean = settings['mean'] 159 | model.std = settings['std'] 160 | return model -------------------------------------------------------------------------------- /src/SatellitesAugs.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import time 3 | import types 4 | import numpy as np 5 | from numpy import random 6 | import random 7 | import imgaug as ia 8 | from imgaug import augmenters as iaa 9 | from torchvision import transforms 10 | from PIL import Image 11 | import torch 12 | from skimage.transform import rotate 13 | 14 | seed = 43 15 | is_mask = False 16 | 17 | # resnet and resnext means 18 | # 'mean': [0.485, 0.456, 0.406], 19 | # 'std': [0.229, 0.224, 0.225] 20 | 21 | # normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5], 22 | # std=[1, 1, 1, 1, 1, 1, 1, 1]) 23 | 24 | normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], 25 | std=[0.229, 0.224, 0.225]) 26 | 27 | class SatellitesTrainAugmentation(object): 28 | def __init__(self, 29 | shape=1280, 30 | aug_scheme=None): 31 | 32 | if aug_scheme == True: 33 | print('Augmentations are enabled for train') 34 | self.augment_img = Compose([ 35 | # RandomCrop(800), 36 | # NpyToPil(), 37 | # transforms.Scale(shape), 38 | # PilToNpy(), 39 | ImgAugAugs(), 40 | RandomCrop(shape), 41 | ToTensor(), 42 | normalize 43 | ]) 44 | self.augment_mask = Compose([ 45 | # RandomCrop(800), 46 | # NpyToPil(), 47 | # transforms.Scale(shape), 48 | # PilToNpy(), 49 | ImgAugAugs(), 50 | RandomCrop(shape), 51 | ToTensor() 52 | ]) 53 | else: 54 | print('Augmentations are NOT enabled for train') 55 | self.augment_img = Compose([ 56 | RandomCrop(shape), 57 | # NpyToPil(), 58 | # transforms.Scale(shape), 59 | # PilToNpy(), 60 | ToTensor(), 61 | normalize 62 | ]) 63 | self.augment_mask = Compose([ 64 | RandomCrop(shape), 65 | # NpyToPil(), 66 | # transforms.Scale(shape), 67 | # PilToNpy(), 68 | ToTensor() 69 | ]) 70 | 71 | def __call__(self, img, mask): 72 | global seed 73 | global is_mask 74 | seed = random.randint(0,100) 75 | # process image 76 | is_mask = False 77 | # naive solution to working with 8-channel images 78 | if img.shape[2]>3: 79 | img1 = self.augment_img(img[:,:,0:3]) 80 | img2 = self.augment_img(img[:,:,3:6]) 81 | img3 = self.augment_img(img[:,:,5:8]) 82 | img = torch.cat((img1[0:3,:,:],img2[0:3,:,:],img3[1:3,:,:])) 83 | else: 84 | img = self.augment_img(img) 85 | # process mask 86 | is_mask = True 87 | # quick hack to evaluate paved only or non-paved only roads 88 | # mask = self.augment(mask[:,:,1:3]) 89 | mask = self.augment_mask(mask) 90 | return img,mask 91 | class SatellitesTestAugmentation(object): 92 | def __init__(self,shape=1280,padding=6): 93 | self.augment_img = Compose([ 94 | # NpyToPil(), 95 | # transforms.Pad(padding=padding, fill=0), 96 | RandomCrop(shape), # most likely causing low score on test test! 97 | # NumpyPad(padding), 98 | # NpyToPil(), 99 | # transforms.Scale(shape), 100 | # PilToNpy(), 101 | ToTensor(), 102 | normalize 103 | ]) 104 | self.augment_mask = Compose([ 105 | # NpyToPil(), 106 | # transforms.Pad(padding=padding, fill=0), 107 | RandomCrop(shape), # most likely causing low score on test test! 108 | # NumpyPad(padding), 109 | # NpyToPil(), 110 | # transforms.Scale(shape), 111 | # PilToNpy(), 112 | ToTensor() 113 | ]) 114 | def __call__(self, img, mask,seed_param=None): 115 | global is_mask 116 | global seed 117 | 118 | if seed_param is None: 119 | seed = random.randint(0,100) 120 | else: 121 | seed = seed_param 122 | 123 | # naive solution to working with 8-channel images 124 | is_mask = False 125 | if img.shape[2]>3: 126 | img1 = self.augment_img(img[:,:,0:3]) 127 | img2 = self.augment_img(img[:,:,3:6]) 128 | img3 = self.augment_img(img[:,:,5:8]) 129 | img = torch.cat((img1[0:3,:,:],img2[0:3,:,:],img3[1:3,:,:])) 130 | else: 131 | img = self.augment_img(img) 132 | if mask is not None: 133 | is_mask = True 134 | # quick hack to evaluate paved only or non-paved only roads 135 | # mask = self.augment(mask[:,:,1:3]) 136 | mask = self.augment_mask(mask) 137 | return img,mask 138 | class SatellitesTestAugmentationPredict(object): 139 | def __init__(self,shape=1280,padding=6): 140 | self.augment_img = Compose([ 141 | NumpyPad(padding), 142 | ToTensor(), 143 | normalize 144 | ]) 145 | self.augment_mask = Compose([ 146 | NumpyPad(padding), 147 | ToTensor() 148 | ]) 149 | def __call__(self, img, mask,seed_param=None): 150 | global is_mask 151 | global seed 152 | 153 | if seed_param is None: 154 | seed = random.randint(0,100) 155 | else: 156 | seed = seed_param 157 | 158 | # naive solution to working with 8-channel images 159 | is_mask = False 160 | if img.shape[2]>3: 161 | img1 = self.augment_img(img[:,:,0:3]) 162 | img2 = self.augment_img(img[:,:,3:6]) 163 | img3 = self.augment_img(img[:,:,5:8]) 164 | img = torch.cat((img1[0:3,:,:],img2[0:3,:,:],img3[1:3,:,:])) 165 | else: 166 | img = self.augment_img(img) 167 | if mask is not None: 168 | is_mask = True 169 | # quick hack to evaluate paved only or non-paved only roads 170 | # mask = self.augment(mask[:,:,1:3]) 171 | mask = self.augment_mask(mask) 172 | return img,mask 173 | class SatellitesTestAugmentationTTA(object): 174 | def __init__(self, 175 | padding=6, 176 | hflip=False, 177 | vflip=False): 178 | if (hflip == True and vflip == False): 179 | self.augment_img = Compose([ 180 | NpyToPil(), 181 | transforms.Pad(padding=padding, fill=0), 182 | PilToNpy(), 183 | HFlip(), 184 | ToTensor(), 185 | normalize 186 | ]) 187 | self.augment_mask = Compose([ 188 | NpyToPil(), 189 | transforms.Pad(padding=padding, fill=0), 190 | PilToNpy(), 191 | HFlip(), 192 | ToTensor() 193 | ]) 194 | elif (hflip == False and vflip == True): 195 | self.augment_img = Compose([ 196 | NpyToPil(), 197 | transforms.Pad(padding=padding, fill=0), 198 | PilToNpy(), 199 | VFlip(), 200 | ToTensor(), 201 | normalize 202 | ]) 203 | self.augment_mask = Compose([ 204 | NpyToPil(), 205 | transforms.Pad(padding=padding, fill=0), 206 | PilToNpy(), 207 | VFlip(), 208 | ToTensor() 209 | ]) 210 | elif (hflip == True and vflip == True): 211 | self.augment_img = Compose([ 212 | NpyToPil(), 213 | transforms.Pad(padding=padding, fill=0), 214 | PilToNpy(), 215 | HFlip(), 216 | VFlip(), 217 | ToTensor(), 218 | normalize 219 | ]) 220 | self.augment_mask = Compose([ 221 | NpyToPil(), 222 | transforms.Pad(padding=padding, fill=0), 223 | PilToNpy(), 224 | HFlip(), 225 | VFlip(), 226 | ToTensor() 227 | ]) 228 | else: 229 | self.augment_img = Compose([ 230 | NpyToPil(), 231 | transforms.Pad(padding=padding, fill=0), 232 | PilToNpy(), 233 | ToTensor(), 234 | normalize 235 | ]) 236 | self.augment_mask = Compose([ 237 | NpyToPil(), 238 | transforms.Pad(padding=padding, fill=0), 239 | PilToNpy(), 240 | ToTensor() 241 | ]) 242 | 243 | def __call__(self, img, mask,seed_param=None): 244 | global is_mask 245 | global seed 246 | 247 | if seed_param is None: 248 | seed = random.randint(0,100) 249 | else: 250 | seed = seed_param 251 | 252 | # naive solution to working with 8-channel images 253 | is_mask = False 254 | if img.shape[2]>3: 255 | img1 = self.augment_img(img[:,:,0:3]) 256 | img2 = self.augment_img(img[:,:,3:6]) 257 | img3 = self.augment_img(img[:,:,5:8]) 258 | img = torch.cat((img1[0:3,:,:],img2[0:3,:,:],img3[1:3,:,:])) 259 | else: 260 | img = self.augment_img(img) 261 | if mask is not None: 262 | is_mask = True 263 | # quick hack to evaluate paved only or non-paved only roads 264 | # mask = self.augment(mask[:,:,1:3]) 265 | mask = self.augment_mask(mask) 266 | return img,mask 267 | class NumpyPad(object): 268 | def __init__(self, 269 | padding = 6): 270 | self.padding = padding 271 | def __call__(self, img): 272 | return np.pad(array=img,pad_width=((self.padding,self.padding), (self.padding,self.padding),(0,0)),mode='constant',constant_values=0) 273 | class HFlip(object): 274 | def __call__(self, 275 | image): 276 | seq = iaa.Sequential([ 277 | iaa.Fliplr(1.0), # horizontally flip 100% of all images 278 | ], random_order=False) 279 | return seq.augment_image(image) 280 | class VFlip(object): 281 | def __call__(self, 282 | image): 283 | seq = iaa.Sequential([ 284 | iaa.Flipud(1.0), # vertically flip 100% of all images 285 | ], random_order=False) 286 | return seq.augment_image(image) 287 | # version compatible with 16-bit images 288 | """ 289 | class ImgAugAugs(object): 290 | def __call__(self, 291 | image): 292 | global seed 293 | 294 | # poor man's flipping 295 | if seed%2==0: 296 | image = np.fliplr(image) 297 | elif seed%4==0: 298 | image = np.fliplr(image) 299 | image = np.flipud(image) 300 | 301 | # poor man's affine transformations 302 | image = rotate(image, 303 | angle=seed, 304 | resize=False, 305 | clip=True, 306 | preserve_range=True) 307 | 308 | return image 309 | """ 310 | class ImgAugAugs(object): 311 | def __call__(self, 312 | image): 313 | global seed 314 | ia.seed(seed) 315 | seq = iaa.Sequential([ 316 | # execute 0 to 1 of the following (less important) augmenters per image 317 | # don't execute all of them, as that would often be way too strong 318 | iaa.Fliplr(0.25), # horizontally flip 25% of all images 319 | iaa.Flipud(0.25), # vertically flip 25% of all images 320 | # Make some images brighter and some darker. 321 | # In 20% of all cases, we sample the multiplier once per channel, 322 | # which can end up changing the color of the images. 323 | iaa.Sometimes(0.25, 324 | iaa.Affine( 325 | scale={"x": (0.9, 1.1), "y": (0.9, 1.1)}, 326 | translate_percent={"x": (-0.1, 0.1), "y": (-0.1, 0.1)}, 327 | rotate=(-90, 90), 328 | shear=(-5, 5) 329 | ), 330 | ), 331 | ], random_order=True) # apply augmenters in random order 332 | 333 | return seq.augment_image(image) 334 | 335 | class RandomCrop(object): 336 | def __init__(self, 337 | shape = 512): 338 | self.shape = shape 339 | def __call__(self,img): 340 | global seed 341 | random.seed(seed) 342 | x_shift = random.randint(0, img.shape[0] - self.shape - 1) 343 | random.seed(seed-1) 344 | y_shift = random.randint(0, img.shape[1] - self.shape - 1) 345 | if len(img.shape)==3: 346 | return img[x_shift:x_shift+self.shape,x_shift:x_shift+self.shape,:] 347 | elif len(img.shape)==2: 348 | return img[x_shift:x_shift+self.shape,x_shift:x_shift+self.shape] 349 | else: 350 | raise NotImplementedError 351 | class Normalize(object): 352 | def __init__(self,mean,std): 353 | self.mean = np.array(mean, dtype=np.float32) 354 | self.std = np.array(std, dtype=np.float32) 355 | def __call__(self, image): 356 | image = image.astype(np.float32) 357 | image[:,:,0:3] -= self.mean 358 | image[:,:,0:3] *= (1/self.std) 359 | return image.astype(np.float32) 360 | class Compose(object): 361 | """Composes several augmentations together. 362 | Args: 363 | transforms (List[Transform]): list of transforms to compose. 364 | Example: 365 | >>> augmentations.Compose([ 366 | >>> transforms.CenterCrop(10), 367 | >>> transforms.ToTensor(), 368 | >>> ]) 369 | """ 370 | 371 | def __init__(self, transforms): 372 | self.transforms = transforms 373 | 374 | def __call__(self, img): 375 | for t in self.transforms: 376 | img = t(img) 377 | return img 378 | class ToCV2Image(object): 379 | def __call__(self, tensor): 380 | return tensor.cpu().numpy().astype(np.float32).transpose((1, 2, 0)) 381 | class PilToNpy(object): 382 | def __call__(self, pil_image): 383 | return np.array(pil_image) 384 | class NpyToPil(object): 385 | def __call__(self, cv2_image): 386 | return Image.fromarray(cv2_image) 387 | class ToTensor(object): 388 | def __call__(self, cvimage): 389 | global is_mask 390 | # process masks 391 | if (is_mask==True) and (len(cvimage.shape)==2): 392 | cvimage = np.expand_dims(cvimage, 2) 393 | cvimage = (cvimage > 255 * 0.5).astype(np.uint8) 394 | return torch.from_numpy(cvimage).permute(2, 0, 1).float() 395 | elif (is_mask==True) and (len(cvimage.shape)==3): 396 | cvimage = (cvimage > 255 * 0.5).astype(np.uint8) 397 | return torch.from_numpy(cvimage).permute(2, 0, 1).float() 398 | else: 399 | # process images 400 | try: 401 | return torch.from_numpy(cvimage).permute(2, 0, 1).float().div(255) 402 | except: 403 | return torch.from_numpy((cvimage.transpose((2, 0, 1))).copy()).float().div(255) 404 | class CannyEdges(object): 405 | def __init__(self,threshold1=100,threshold2=200): 406 | self.threshold1 = threshold1 407 | self.threshold2 = threshold2 408 | def __call__(self, image): 409 | canny_region = np.uint8(image[:,:,0:3]) 410 | edges = cv2.Canny(canny_region,self.threshold1,self.threshold2) 411 | return np.dstack( ( image,edges) ) 412 | class SaliencyMap(object): 413 | def __call__(self, image): 414 | sm = pySaliencyMap(image.shape[0], image.shape[1]) 415 | return np.dstack( ( image,sm.SMGetSM(image)) ) 416 | class PhotometricDistort(object): 417 | def __init__(self): 418 | self.pd = [ 419 | RandomContrast(), 420 | ConvertColor(transform='HSV'), 421 | RandomSaturation(), 422 | RandomHue(), 423 | ConvertColor(current='HSV', transform='BGR'), 424 | RandomContrast() 425 | ] 426 | self.rand_brightness = RandomBrightness() 427 | self.rand_light_noise = RandomLightingNoise() 428 | 429 | def __call__(self, image): 430 | im = image.copy() 431 | im = self.rand_brightness(im) 432 | if random.randint(1,2): 433 | distort = Compose(self.pd[:-1]) 434 | else: 435 | distort = Compose(self.pd[1:]) 436 | im = distort(im) 437 | return self.rand_light_noise(im) 438 | class RandomSaturation(object): 439 | def __init__(self, lower=0.5, upper=1.5): 440 | self.lower = lower 441 | self.upper = upper 442 | assert self.upper >= self.lower, "contrast upper must be >= lower." 443 | assert self.lower >= 0, "contrast lower must be non-negative." 444 | 445 | def __call__(self, image): 446 | if random.randint(1,2): 447 | image[:, :, 1] *= random.uniform(self.lower, self.upper) 448 | 449 | return image 450 | class RandomHue(object): 451 | def __init__(self, delta=18.0): 452 | assert delta >= 0.0 and delta <= 360.0 453 | self.delta = delta 454 | 455 | def __call__(self, image): 456 | if random.randint(1,2): 457 | image[:, :, 0] += random.uniform(-self.delta, self.delta) 458 | image[:, :, 0][image[:, :, 0] > 360.0] -= 360.0 459 | image[:, :, 0][image[:, :, 0] < 0.0] += 360.0 460 | return image 461 | class RandomLightingNoise(object): 462 | def __init__(self): 463 | self.perms = ((0, 1, 2), (0, 2, 1), 464 | (1, 0, 2), (1, 2, 0), 465 | (2, 0, 1), (2, 1, 0)) 466 | 467 | def __call__(self, image): 468 | if random.randint(1,2): 469 | swap = self.perms[random.randint(0,len(self.perms)-1)] 470 | shuffle = SwapChannels(swap) # shuffle channels 471 | image = shuffle(image) 472 | return image 473 | class ConvertColor(object): 474 | def __init__(self, current='BGR', transform='HSV'): 475 | self.transform = transform 476 | self.current = current 477 | 478 | def __call__(self, image, boxes=None, labels=None): 479 | if self.current == 'BGR' and self.transform == 'HSV': 480 | image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) 481 | elif self.current == 'HSV' and self.transform == 'BGR': 482 | image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR) 483 | else: 484 | raise NotImplementedError 485 | return image 486 | class RandomContrast(object): 487 | def __init__(self, lower=0.5, upper=1.5): 488 | self.lower = lower 489 | self.upper = upper 490 | assert self.upper >= self.lower, "contrast upper must be >= lower." 491 | assert self.lower >= 0, "contrast lower must be non-negative." 492 | 493 | # expects float image 494 | def __call__(self, image): 495 | if random.randint(1,2): 496 | alpha = random.uniform(self.lower, self.upper) 497 | image *= alpha 498 | return image 499 | class RandomBrightness(object): 500 | def __init__(self, delta=32): 501 | assert delta >= 0.0 502 | assert delta <= 255.0 503 | self.delta = delta 504 | 505 | def __call__(self, image): 506 | if random.randint(1,2): 507 | delta = random.uniform(-self.delta, self.delta) 508 | image += delta 509 | return image 510 | class ConvertFromInts(object): 511 | def __call__(self, image, boxes=None, labels=None): 512 | return image.astype(np.float32) 513 | class SwapChannels(object): 514 | """Transforms a tensorized image by swapping the channels in the order 515 | specified in the swap tuple. 516 | Args: 517 | swaps (int triple): final order of channels 518 | eg: (2, 1, 0) 519 | """ 520 | 521 | def __init__(self, swaps): 522 | self.swaps = swaps 523 | 524 | def __call__(self, image): 525 | """ 526 | Args: 527 | image (Tensor): image tensor to be transformed 528 | Return: 529 | a tensor with channels swapped according to swap 530 | """ 531 | # if torch.is_tensor(image): 532 | # image = image.data.cpu().numpy() 533 | # else: 534 | # image = np.array(image) 535 | image = image[:, :, self.swaps] 536 | return image 537 | -------------------------------------------------------------------------------- /src/SatellitesDataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | import numpy as np 4 | import pandas as pd 5 | from skimage.io import imread 6 | import torch.utils.data as data 7 | 8 | import warnings 9 | warnings.filterwarnings('ignore') 10 | 11 | # relative dataset storage path for test dataset 12 | prefix = '../data' 13 | # relative dataset storage path for train dataset 14 | prefix_train = '../' 15 | # metadata and mask_df files 16 | meta_prefix = '../' 17 | meta_data_file = os.path.join(meta_prefix,'metadata.csv') 18 | mask_df_file = os.path.join(meta_prefix,'mask_df.csv') 19 | wide_mask_df_file = 'new_masks.csv' 20 | layered_mask_df_file = 'new_masks_layered.csv' 21 | 22 | # high level function that return list of images and cities under presets 23 | def get_test_dataset(preset, 24 | preset_dict, 25 | city='all'): 26 | meta_df = pd.read_csv(meta_data_file) 27 | 28 | test_folders = ['AOI_2_Vegas_Roads_Test_Public','AOI_5_Khartoum_Roads_Test_Public', 29 | 'AOI_3_Paris_Roads_Test_Public','AOI_4_Shanghai_Roads_Test_Public'] 30 | 31 | cities_dict = {'all': ['AOI_2_Vegas_Roads_Test_Public', 'AOI_5_Khartoum_Roads_Test_Public','AOI_3_Paris_Roads_Test_Public', 'AOI_4_Shanghai_Roads_Test_Public'], 32 | 'vegas':['AOI_2_Vegas_Roads_Test_Public'], 33 | 'paris':['AOI_3_Paris_Roads_Test_Public'], 34 | 'shanghai':['AOI_4_Shanghai_Roads_Test_Public'], 35 | 'khartoum': ['AOI_5_Khartoum_Roads_Test_Public']} 36 | 37 | # select the images 38 | sample_df = meta_df[(meta_df.img_files.isin(test_folders)) 39 | &(meta_df.width == preset_dict[preset]['width']) 40 | &(meta_df.channels == preset_dict[preset]['channel_count']) 41 | &(meta_df.img_folders == preset_dict[preset]['subfolder']) 42 | &(meta_df.img_files.isin(cities_dict[city]))] 43 | 44 | # get the data as lists for simplicity 45 | or_imgs = list(sample_df[['img_subfolders','img_files','img_folders']] 46 | .apply(lambda row: os.path.join(prefix,row['img_files'],row['img_folders']+'_8bit',row['img_subfolders']), axis=1).values) 47 | 48 | le, u = sample_df['img_folders'].factorize() 49 | sample_df.loc[:,'city_no'] = le 50 | cty_no = list(sample_df.city_no.values) 51 | 52 | city_folders = list(sample_df.img_files.values) 53 | img_names = list(sample_df.img_subfolders.values) 54 | 55 | return or_imgs,city_folders,img_names,cty_no,prefix 56 | 57 | # high level function that return list of images and cities under presets 58 | def get_train_dataset(preset, 59 | preset_dict, 60 | city='all'): 61 | mask_df = pd.read_csv(mask_df_file) 62 | meta_df = pd.read_csv(meta_data_file) 63 | data_df = mask_df.merge(meta_df[['img_subfolders','width','channels']], how = 'left', left_on = 'img_file', right_on = 'img_subfolders') 64 | 65 | cities_dict = {'all': ['AOI_2_Vegas_Roads_Train', 'AOI_5_Khartoum_Roads_Train','AOI_3_Paris_Roads_Train', 'AOI_4_Shanghai_Roads_Train'], 66 | 'vegas':['AOI_2_Vegas_Roads_Train'], 67 | 'paris':['AOI_3_Paris_Roads_Train'], 68 | 'shanghai':['AOI_4_Shanghai_Roads_Train'], 69 | 'khartoum': ['AOI_5_Khartoum_Roads_Train']} 70 | 71 | # select the images 72 | sample_df = data_df[(data_df.width == preset_dict[preset]['width']) 73 | &(data_df.mask_max > 0) 74 | &(data_df.channels == preset_dict[preset]['channel_count']) 75 | &(data_df.img_subfolder == preset_dict[preset]['subfolder']) 76 | &(data_df.img_folder.isin(cities_dict[city]))] 77 | 78 | # get the data as lists for simplicity 79 | bit8_imgs = list(sample_df.bit8_path.values) 80 | bit8_masks = list(sample_df.mask_path.values) 81 | bit8_imgs = [(os.path.join(prefix_train,path)) for path in bit8_imgs] 82 | bit8_masks = [(os.path.join(prefix_train,path)) for path in bit8_masks] 83 | le, u = sample_df['img_folder'].factorize() 84 | sample_df.loc[:,'city_no'] = le 85 | cty_no = list(sample_df.city_no.values) 86 | 87 | return bit8_imgs,bit8_masks,cty_no 88 | 89 | # dataset class 90 | class SatellitesDataset(data.Dataset): 91 | def __init__(self, 92 | preset, 93 | image_paths = [], 94 | mask_paths = None, 95 | transforms = None, 96 | ): 97 | 98 | self.mask_paths = mask_paths 99 | self.preset = preset 100 | self.transforms = transforms 101 | 102 | if mask_paths is not None: 103 | self.image_paths = sorted(image_paths) 104 | self.mask_paths = sorted(mask_paths) 105 | 106 | if len(self.image_paths) != len(mask_paths): 107 | raise ValueError('Mask list length <> image list lenth') 108 | if [path.split('/')[4].split('img')[1].split('.')[0] for path in self.image_paths] != [path.split('/')[4].split('img')[1].split('.')[0] for path in self.mask_paths]: 109 | raise ValueError('Mask list sorting <> image list sorting') 110 | else: 111 | self.image_paths = image_paths 112 | # self.image_paths = sorted(image_paths) 113 | 114 | def __len__(self): 115 | return len(self.image_paths) 116 | 117 | def __getitem__(self, idx): 118 | if self.mask_paths is not None: 119 | 120 | img = imread(self.image_paths[idx]) 121 | target_channels = np.zeros(shape=(self.preset['width'],self.preset['width'],len(self.preset['channels']))) 122 | 123 | # expand grayscale images to 3 dimensions 124 | if len(img.shape)<3: 125 | img = np.expand_dims(img, 2) 126 | 127 | for i,channel in enumerate(self.preset['channels']): 128 | target_channels[:,:,i] = img[:,:,channel-1] 129 | 130 | target_channels = target_channels.astype('uint8') 131 | 132 | mask = imread(self.mask_paths[idx]) 133 | mask = mask.astype('uint8') 134 | 135 | if self.transforms is not None: 136 | target_channels, mask = self.transforms(target_channels, mask) 137 | 138 | return target_channels,mask 139 | 140 | else: 141 | img = imread(self.image_paths[idx]) 142 | target_channels = np.zeros(shape=(self.preset['width'],self.preset['width'],len(self.preset['channels']))) 143 | 144 | for i,channel in enumerate(self.preset['channels']): 145 | target_channels[:,:,i] = img[:,:,channel-1] 146 | 147 | target_channels = target_channels.astype('uint8') 148 | 149 | if self.transforms is not None: 150 | target_channels, _ = self.transforms(target_channels, None) 151 | return target_channels 152 | 153 | def get_train_dataset_for_predict(preset, 154 | preset_dict, 155 | city='all'): 156 | mask_df = pd.read_csv(mask_df_file) 157 | meta_df = pd.read_csv(meta_data_file) 158 | data_df = mask_df.merge(meta_df[['img_subfolders','width','channels']], how = 'left', left_on = 'img_file', right_on = 'img_subfolders') 159 | 160 | cities_dict = {'all': ['AOI_2_Vegas_Roads_Train', 'AOI_5_Khartoum_Roads_Train','AOI_3_Paris_Roads_Train', 'AOI_4_Shanghai_Roads_Train'], 161 | 'vegas':['AOI_2_Vegas_Roads_Train'], 162 | 'paris':['AOI_3_Paris_Roads_Train'], 163 | 'shanghai':['AOI_4_Shanghai_Roads_Train'], 164 | 'khartoum': ['AOI_5_Khartoum_Roads_Train']} 165 | 166 | # select the images 167 | sample_df = data_df[(data_df.width == preset_dict[preset]['width']) 168 | &(data_df.mask_max > 0) 169 | &(data_df.channels == preset_dict[preset]['channel_count']) 170 | &(data_df.img_subfolder == preset_dict[preset]['subfolder']) 171 | &(data_df.img_folder.isin(cities_dict[city]))] 172 | 173 | # get the data as lists for simplicity 174 | bit8_imgs = list(sample_df.bit8_path.values) 175 | bit8_imgs = [(os.path.join(prefix_train,path)) for path in bit8_imgs] 176 | 177 | le, u = sample_df['img_folder'].factorize() 178 | sample_df.loc[:,'city_no'] = le 179 | cty_no = list(sample_df.city_no.values) 180 | 181 | city_folders = list(sample_df.img_folder.values) 182 | img_names = list(sample_df.img_file.values) 183 | 184 | return bit8_imgs,city_folders,img_names,cty_no,prefix 185 | 186 | def get_train_dataset_wide_masks(preset, 187 | preset_dict): 188 | mask_df = pd.read_csv(mask_df_file) 189 | meta_df = pd.read_csv(meta_data_file) 190 | data_df = mask_df.merge(meta_df[['img_subfolders','width','channels']], how = 'left', left_on = 'img_file', right_on = 'img_subfolders') 191 | 192 | # filter out bad new masks 193 | new_mask_df = pd.read_csv(wide_mask_df_file) 194 | good_new_masks = list(new_mask_df[new_mask_df.correct == 1].img_names.values) 195 | 196 | # filter only roads with mostly non-paved roads 197 | 198 | df = pd.read_csv('geojson_df_full.csv') 199 | table = pd.pivot_table(df, 200 | index=["img_id"], 201 | columns = ['paved'], 202 | values=["linestring"], 203 | aggfunc={len},fill_value=0) 204 | table.columns = ['count_paved','count_non_paved'] 205 | mostly_non_paved_imgs = list(table[table.count_paved < table.count_non_paved].index.values) 206 | 207 | 208 | # select the images 209 | sample_df = data_df[(data_df.width == preset_dict[preset]['width']) 210 | &(data_df.mask_max > 0) # filter broken masks 211 | &(data_df.channels == preset_dict[preset]['channel_count']) # preset filter 212 | &(data_df.img_subfolder == preset_dict[preset]['subfolder']) # preset filter 213 | &(data_df.img_file.isin(good_new_masks) # filter new good masks 214 | &(data_df['img_subfolders'] 215 | .apply(lambda x: 'AOI'+x.split('AOI')[1][:-4]) 216 | .isin(mostly_non_paved_imgs)) # filter out mostly paved roads 217 | ) 218 | ] 219 | 220 | # get the data as lists for simplicity 221 | bit8_imgs = list(sample_df.bit8_path.values) 222 | bit8_masks = list(sample_df.mask_path.values) 223 | 224 | # replace masks with width masks 225 | bit8_masks = [(path.replace("_mask","_width_mask")) for path in bit8_masks] 226 | 227 | bit8_imgs = [(os.path.join(prefix_train,path)) for path in bit8_imgs] 228 | bit8_masks = [(os.path.join(prefix_train,path)) for path in bit8_masks] 229 | le, u = sample_df['img_folder'].factorize() 230 | sample_df.loc[:,'city_no'] = le 231 | cty_no = list(sample_df.city_no.values) 232 | 233 | return bit8_imgs,bit8_masks,cty_no 234 | 235 | def get_train_dataset_layered_masks(preset, 236 | preset_dict): 237 | mask_df = pd.read_csv(mask_df_file) 238 | meta_df = pd.read_csv(meta_data_file) 239 | data_df = mask_df.merge(meta_df[['img_subfolders','width','channels']], how = 'left', left_on = 'img_file', right_on = 'img_subfolders') 240 | 241 | # note that I am removing wide masks - because layered masks with zero pixels somehow passed my filter 242 | new_mask_df = pd.read_csv(wide_mask_df_file) 243 | good_new_masks = list(new_mask_df[new_mask_df.correct == 1].img_names.values) 244 | 245 | # select the images 246 | sample_df = data_df[(data_df.width == preset_dict[preset]['width']) 247 | &(data_df.mask_max > 0) 248 | &(data_df.channels == preset_dict[preset]['channel_count']) 249 | &(data_df.img_subfolder == preset_dict[preset]['subfolder']) 250 | &(data_df.img_file.isin(good_new_masks)) 251 | ] 252 | 253 | # get the data as lists for simplicity 254 | bit8_imgs = list(sample_df.bit8_path.values) 255 | bit8_masks = list(sample_df.mask_path.values) 256 | 257 | # replace masks with width masks 258 | bit8_masks = [(path.replace("_mask","_layered_mask")) for path in bit8_masks] 259 | 260 | bit8_imgs = [(os.path.join(prefix_train,path)) for path in bit8_imgs] 261 | bit8_masks = [(os.path.join(prefix_train,path)) for path in bit8_masks] 262 | le, u = sample_df['img_folder'].factorize() 263 | sample_df.loc[:,'city_no'] = le 264 | cty_no = list(sample_df.city_no.values) 265 | 266 | return bit8_imgs,bit8_masks,cty_no 267 | 268 | # high level function that return list of images and cities under presets 269 | def get_train_dataset_all(preset, 270 | preset_dict, 271 | city='all'): 272 | mask_df = pd.read_csv(mask_df_file) 273 | meta_df = pd.read_csv(meta_data_file) 274 | data_df = mask_df.merge(meta_df[['img_subfolders','width','channels']], how = 'left', left_on = 'img_file', right_on = 'img_subfolders') 275 | 276 | cities_dict = {'all': ['AOI_2_Vegas_Roads_Train', 'AOI_5_Khartoum_Roads_Train','AOI_3_Paris_Roads_Train', 'AOI_4_Shanghai_Roads_Train'], 277 | 'vegas':['AOI_2_Vegas_Roads_Train'], 278 | 'paris':['AOI_3_Paris_Roads_Train'], 279 | 'shanghai':['AOI_4_Shanghai_Roads_Train'], 280 | 'khartoum': ['AOI_5_Khartoum_Roads_Train']} 281 | 282 | # select the images 283 | sample_df = data_df[(data_df.width == preset_dict[preset]['width']) 284 | # &(data_df.mask_max > 0) 285 | &(data_df.channels == preset_dict[preset]['channel_count']) 286 | &(data_df.img_subfolder == preset_dict[preset]['subfolder']) 287 | &(data_df.img_folder.isin(cities_dict[city]))] 288 | 289 | # get the data as lists for simplicity 290 | bit8_imgs = list(sample_df.bit8_path.values) 291 | bit8_masks = list(sample_df.mask_path.values) 292 | 293 | bit8_masks = [(path.replace("_mask","_all_mask")) for path in bit8_masks] 294 | bit8_masks = [(path.replace("RGB-PanSharpen_all_mask/RGB-PanSharpen","MUL-PanSharpen_all_mask/MUL-PanSharpen")) for path in bit8_masks] 295 | 296 | 297 | bit8_imgs = [(os.path.join(prefix_train,path)) for path in bit8_imgs] 298 | bit8_masks = [(os.path.join(prefix_train,path)) for path in bit8_masks] 299 | le, u = sample_df['img_folder'].factorize() 300 | sample_df.loc[:,'city_no'] = le 301 | cty_no = list(sample_df.city_no.values) 302 | 303 | return bit8_imgs,bit8_masks,cty_no 304 | 305 | def get_train_dataset_for_predict_all(preset, 306 | preset_dict, 307 | city='all'): 308 | mask_df = pd.read_csv(mask_df_file) 309 | meta_df = pd.read_csv(meta_data_file) 310 | data_df = mask_df.merge(meta_df[['img_subfolders','width','channels']], how = 'left', left_on = 'img_file', right_on = 'img_subfolders') 311 | 312 | cities_dict = {'all': ['AOI_2_Vegas_Roads_Train', 'AOI_5_Khartoum_Roads_Train','AOI_3_Paris_Roads_Train', 'AOI_4_Shanghai_Roads_Train'], 313 | 'vegas':['AOI_2_Vegas_Roads_Train'], 314 | 'paris':['AOI_3_Paris_Roads_Train'], 315 | 'shanghai':['AOI_4_Shanghai_Roads_Train'], 316 | 'khartoum': ['AOI_5_Khartoum_Roads_Train']} 317 | 318 | # select the images 319 | sample_df = data_df[(data_df.width == preset_dict[preset]['width']) 320 | # &(data_df.mask_max > 0) 321 | &(data_df.channels == preset_dict[preset]['channel_count']) 322 | &(data_df.img_subfolder == preset_dict[preset]['subfolder']) 323 | &(data_df.img_folder.isin(cities_dict[city]))] 324 | 325 | # get the data as lists for simplicity 326 | bit8_imgs = list(sample_df.bit8_path.values) 327 | bit8_imgs = [(os.path.join(prefix_train,path)) for path in bit8_imgs] 328 | 329 | le, u = sample_df['img_folder'].factorize() 330 | sample_df.loc[:,'city_no'] = le 331 | cty_no = list(sample_df.city_no.values) 332 | 333 | city_folders = list(sample_df.img_folder.values) 334 | img_names = list(sample_df.img_file.values) 335 | 336 | return bit8_imgs,city_folders,img_names,cty_no,prefix 337 | def get_train_dataset_all_16bit(preset, 338 | preset_dict, 339 | city='all'): 340 | mask_df = pd.read_csv(mask_df_file) 341 | meta_df = pd.read_csv(meta_data_file) 342 | data_df = mask_df.merge(meta_df[['img_subfolders','width','channels']], how = 'left', left_on = 'img_file', right_on = 'img_subfolders') 343 | 344 | cities_dict = {'all': ['AOI_2_Vegas_Roads_Train', 'AOI_5_Khartoum_Roads_Train','AOI_3_Paris_Roads_Train', 'AOI_4_Shanghai_Roads_Train'], 345 | 'vegas':['AOI_2_Vegas_Roads_Train'], 346 | 'paris':['AOI_3_Paris_Roads_Train'], 347 | 'shanghai':['AOI_4_Shanghai_Roads_Train'], 348 | 'khartoum': ['AOI_5_Khartoum_Roads_Train']} 349 | 350 | # select the images 351 | sample_df = data_df[(data_df.width == preset_dict[preset]['width']) 352 | # &(data_df.mask_max > 0) 353 | &(data_df.channels == preset_dict[preset]['channel_count']) 354 | &(data_df.img_subfolder == preset_dict[preset]['subfolder']) 355 | &(data_df.img_folder.isin(cities_dict[city]))] 356 | 357 | # get the data as lists for simplicity 358 | bit16_imgs = list(sample_df.img_path.values) 359 | bit8_masks = list(sample_df.mask_path.values) 360 | 361 | bit8_masks = [(path.replace("_mask","_all_mask")) for path in bit8_masks] 362 | 363 | bit16_imgs = [(os.path.join(prefix_train,path)) for path in bit16_imgs] 364 | bit8_masks = [(os.path.join(prefix_train,path)) for path in bit8_masks] 365 | le, u = sample_df['img_folder'].factorize() 366 | sample_df.loc[:,'city_no'] = le 367 | cty_no = list(sample_df.city_no.values) 368 | 369 | return bit16_imgs,bit8_masks,cty_no 370 | def get_train_dataset_for_predict_all_16bit(preset, 371 | preset_dict, 372 | city='all'): 373 | mask_df = pd.read_csv(mask_df_file) 374 | meta_df = pd.read_csv(meta_data_file) 375 | data_df = mask_df.merge(meta_df[['img_subfolders','width','channels']], how = 'left', left_on = 'img_file', right_on = 'img_subfolders') 376 | 377 | cities_dict = {'all': ['AOI_2_Vegas_Roads_Train', 'AOI_5_Khartoum_Roads_Train','AOI_3_Paris_Roads_Train', 'AOI_4_Shanghai_Roads_Train'], 378 | 'vegas':['AOI_2_Vegas_Roads_Train'], 379 | 'paris':['AOI_3_Paris_Roads_Train'], 380 | 'shanghai':['AOI_4_Shanghai_Roads_Train'], 381 | 'khartoum': ['AOI_5_Khartoum_Roads_Train']} 382 | 383 | # select the images 384 | sample_df = data_df[(data_df.width == preset_dict[preset]['width']) 385 | # &(data_df.mask_max > 0) 386 | &(data_df.channels == preset_dict[preset]['channel_count']) 387 | &(data_df.img_subfolder == preset_dict[preset]['subfolder']) 388 | &(data_df.img_folder.isin(cities_dict[city]))] 389 | 390 | # get the data as lists for simplicity 391 | bit16_imgs = list(sample_df.img_path.values) 392 | bit16_imgs = [(os.path.join(prefix_train,path)) for path in bit16_imgs] 393 | 394 | le, u = sample_df['img_folder'].factorize() 395 | sample_df.loc[:,'city_no'] = le 396 | cty_no = list(sample_df.city_no.values) 397 | 398 | city_folders = list(sample_df.img_folder.values) 399 | img_names = list(sample_df.img_file.values) 400 | 401 | return bit16_imgs,city_folders,img_names,cty_no,prefix 402 | def get_test_dataset_16bit(preset, 403 | preset_dict, 404 | city='all'): 405 | meta_df = pd.read_csv(meta_data_file) 406 | 407 | test_folders = ['AOI_2_Vegas_Roads_Test_Public','AOI_5_Khartoum_Roads_Test_Public', 408 | 'AOI_3_Paris_Roads_Test_Public','AOI_4_Shanghai_Roads_Test_Public'] 409 | 410 | cities_dict = {'all': ['AOI_2_Vegas_Roads_Test_Public', 'AOI_5_Khartoum_Roads_Test_Public','AOI_3_Paris_Roads_Test_Public', 'AOI_4_Shanghai_Roads_Test_Public'], 411 | 'vegas':['AOI_2_Vegas_Roads_Test_Public'], 412 | 'paris':['AOI_3_Paris_Roads_Test_Public'], 413 | 'shanghai':['AOI_4_Shanghai_Roads_Test_Public'], 414 | 'khartoum': ['AOI_5_Khartoum_Roads_Test_Public']} 415 | 416 | # select the images 417 | sample_df = meta_df[(meta_df.img_files.isin(test_folders)) 418 | &(meta_df.width == preset_dict[preset]['width']) 419 | &(meta_df.channels == preset_dict[preset]['channel_count']) 420 | &(meta_df.img_folders == preset_dict[preset]['subfolder']) 421 | &(meta_df.img_files.isin(cities_dict[city]))] 422 | 423 | # get the data as lists for simplicity 424 | or_imgs = list(sample_df[['img_subfolders','img_files','img_folders']] 425 | .apply(lambda row: os.path.join(prefix,row['img_files'],row['img_folders'],row['img_subfolders']), axis=1).values) 426 | 427 | le, u = sample_df['img_folders'].factorize() 428 | sample_df.loc[:,'city_no'] = le 429 | cty_no = list(sample_df.city_no.values) 430 | 431 | city_folders = list(sample_df.img_files.values) 432 | img_names = list(sample_df.img_subfolders.values) 433 | 434 | return or_imgs,city_folders,img_names,cty_no,prefix -------------------------------------------------------------------------------- /src/TbLogger.py: -------------------------------------------------------------------------------- 1 | # Code referenced from https://gist.github.com/gyglim/1f8dfb1b5c82627ae3efcfbbadb9f514 2 | import tensorflow as tf 3 | import numpy as np 4 | import scipy.misc 5 | try: 6 | from StringIO import StringIO # Python 2.7 7 | except ImportError: 8 | from io import BytesIO # Python 3.x 9 | 10 | 11 | class Logger(object): 12 | 13 | def __init__(self, log_dir): 14 | """Create a summary writer logging to log_dir.""" 15 | self.writer = tf.summary.FileWriter(log_dir) 16 | 17 | def scalar_summary(self, tag, value, step): 18 | """Log a scalar variable.""" 19 | summary = tf.Summary(value=[tf.Summary.Value(tag=tag, simple_value=value)]) 20 | self.writer.add_summary(summary, step) 21 | 22 | def image_summary(self, tag, images, step): 23 | """Log a list of images.""" 24 | 25 | img_summaries = [] 26 | for i, img in enumerate(images): 27 | # Write the image to a string 28 | try: 29 | s = StringIO() 30 | except: 31 | s = BytesIO() 32 | scipy.misc.toimage(img).save(s, format="png") 33 | 34 | # Create an Image object 35 | img_sum = tf.Summary.Image(encoded_image_string=s.getvalue(), 36 | height=img.shape[0], 37 | width=img.shape[1]) 38 | # Create a Summary value 39 | img_summaries.append(tf.Summary.Value(tag='%s/%d' % (tag, i), image=img_sum)) 40 | 41 | # Create and write Summary 42 | summary = tf.Summary(value=img_summaries) 43 | self.writer.add_summary(summary, step) 44 | 45 | def histo_summary(self, tag, values, step, bins=1000): 46 | """Log a histogram of the tensor of values.""" 47 | 48 | # Create a histogram using numpy 49 | counts, bin_edges = np.histogram(values, bins=bins) 50 | 51 | # Fill the fields of the histogram proto 52 | hist = tf.HistogramProto() 53 | hist.min = float(np.min(values)) 54 | hist.max = float(np.max(values)) 55 | hist.num = int(np.prod(values.shape)) 56 | hist.sum = float(np.sum(values)) 57 | hist.sum_squares = float(np.sum(values**2)) 58 | 59 | # Drop the start of the first bin 60 | bin_edges = bin_edges[1:] 61 | 62 | # Add bin edges and counts 63 | for edge in bin_edges: 64 | hist.bucket_limit.append(edge) 65 | for c in counts: 66 | hist.bucket.append(c) 67 | 68 | # Create and write Summary 69 | summary = tf.Summary(value=[tf.Summary.Value(tag=tag, histo=hist)]) 70 | self.writer.add_summary(summary, step) 71 | self.writer.flush() 72 | -------------------------------------------------------------------------------- /src/UNet.py: -------------------------------------------------------------------------------- 1 | """ 2 | all credits to ternaus and albu 3 | """ 4 | import torch 5 | import torch.nn as nn 6 | from torch.autograd import Variable 7 | from torchvision import models 8 | from torch.nn import functional as F 9 | 10 | def conv3x3(in_, out): 11 | return nn.Conv2d(in_, out, 3, padding=1) 12 | 13 | def concat(xs): 14 | return torch.cat(xs, 1) 15 | 16 | class Conv3BN(nn.Module): 17 | def __init__(self, in_: int, out: int, bn=False): 18 | super().__init__() 19 | self.conv = conv3x3(in_, out) 20 | self.bn = nn.BatchNorm2d(out) if bn else None 21 | self.activation = nn.SELU(inplace=True) 22 | 23 | def forward(self, x): 24 | x = self.conv(x) 25 | if self.bn is not None: 26 | x = self.bn(x) 27 | x = self.activation(x) 28 | return x 29 | 30 | class UNetModule(nn.Module): 31 | def __init__(self, in_: int, out: int): 32 | super().__init__() 33 | self.l1 = Conv3BN(in_, out) 34 | self.l2 = Conv3BN(out, out) 35 | 36 | def forward(self, x): 37 | x = self.l1(x) 38 | x = self.l2(x) 39 | return x 40 | 41 | class ConvRelu(nn.Module): 42 | def __init__(self, in_: int, out: int): 43 | super().__init__() 44 | self.conv = conv3x3(in_, out) 45 | self.activation = nn.ReLU(inplace=True) 46 | 47 | def forward(self, x): 48 | x = self.conv(x) 49 | x = self.activation(x) 50 | return x 51 | 52 | class DecoderBlock(nn.Module): 53 | def __init__(self, in_channels, middle_channels, out_channels): 54 | super().__init__() 55 | 56 | self.block = nn.Sequential( 57 | ConvRelu(in_channels, middle_channels), 58 | nn.ConvTranspose2d(middle_channels, out_channels, kernel_size=3, stride=2, padding=1, output_padding=1), 59 | nn.ReLU(inplace=True) 60 | ) 61 | 62 | def forward(self, x): 63 | return self.block(x) 64 | 65 | class UNet11(nn.Module): 66 | def __init__(self, num_classes=1, num_filters=32, num_channels=3): 67 | super().__init__() 68 | self.pool = nn.MaxPool2d(2, 2) 69 | encoder = models.vgg11(pretrained=True).features 70 | self.relu = encoder[1] 71 | 72 | # try to use 8-channels as first input 73 | if num_channels==3: 74 | self.conv1 = encoder[0] 75 | else: 76 | self.conv1 = nn.Conv2d(num_channels, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) 77 | 78 | self.conv2 = encoder[3] 79 | self.conv3s = encoder[6] 80 | self.conv3 = encoder[8] 81 | self.conv4s = encoder[11] 82 | self.conv4 = encoder[13] 83 | self.conv5s = encoder[16] 84 | self.conv5 = encoder[18] 85 | 86 | self.center = DecoderBlock(num_filters * 8 * 2, num_filters * 8 * 2, num_filters * 8) 87 | self.dec5 = DecoderBlock(num_filters * (16 + 8), num_filters * 8 * 2, num_filters * 8) 88 | self.dec4 = DecoderBlock(num_filters * (16 + 8), num_filters * 8 * 2, num_filters * 4) 89 | self.dec3 = DecoderBlock(num_filters * (8 + 4), num_filters * 4 * 2, num_filters * 2) 90 | self.dec2 = DecoderBlock(num_filters * (4 + 2), num_filters * 2 * 2, num_filters) 91 | self.dec1 = ConvRelu(num_filters * (2 + 1), num_filters) 92 | 93 | self.final = nn.Conv2d(num_filters, num_classes, kernel_size=1) 94 | 95 | def forward(self, x): 96 | conv1 = self.relu(self.conv1(x)) 97 | conv2 = self.relu(self.conv2(self.pool(conv1))) 98 | conv3s = self.relu(self.conv3s(self.pool(conv2))) 99 | conv3 = self.relu(self.conv3(conv3s)) 100 | conv4s = self.relu(self.conv4s(self.pool(conv3))) 101 | conv4 = self.relu(self.conv4(conv4s)) 102 | conv5s = self.relu(self.conv5s(self.pool(conv4))) 103 | conv5 = self.relu(self.conv5(conv5s)) 104 | 105 | center = self.center(self.pool(conv5)) 106 | 107 | dec5 = self.dec5(torch.cat([center, conv5], 1)) 108 | dec4 = self.dec4(torch.cat([dec5, conv4], 1)) 109 | dec3 = self.dec3(torch.cat([dec4, conv3], 1)) 110 | dec2 = self.dec2(torch.cat([dec3, conv2], 1)) 111 | dec1 = self.dec1(torch.cat([dec2, conv1], 1)) 112 | #return F.sigmoid(self.final(dec1)) 113 | return self.final(dec1) 114 | -------------------------------------------------------------------------------- /src/__pycache__/DilatedResnet.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snakers4/spacenet-three/5e1b202b824c6d976298488225078d850c908a8e/src/__pycache__/DilatedResnet.cpython-35.pyc -------------------------------------------------------------------------------- /src/__pycache__/InceptionResnetv2.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snakers4/spacenet-three/5e1b202b824c6d976298488225078d850c908a8e/src/__pycache__/InceptionResnetv2.cpython-35.pyc -------------------------------------------------------------------------------- /src/__pycache__/LinkNet.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snakers4/spacenet-three/5e1b202b824c6d976298488225078d850c908a8e/src/__pycache__/LinkNet.cpython-35.pyc -------------------------------------------------------------------------------- /src/__pycache__/Loss.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snakers4/spacenet-three/5e1b202b824c6d976298488225078d850c908a8e/src/__pycache__/Loss.cpython-35.pyc -------------------------------------------------------------------------------- /src/__pycache__/LossSemSeg.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snakers4/spacenet-three/5e1b202b824c6d976298488225078d850c908a8e/src/__pycache__/LossSemSeg.cpython-35.pyc -------------------------------------------------------------------------------- /src/__pycache__/MaskUtils.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snakers4/spacenet-three/5e1b202b824c6d976298488225078d850c908a8e/src/__pycache__/MaskUtils.cpython-35.pyc -------------------------------------------------------------------------------- /src/__pycache__/ResNeXt.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snakers4/spacenet-three/5e1b202b824c6d976298488225078d850c908a8e/src/__pycache__/ResNeXt.cpython-35.pyc -------------------------------------------------------------------------------- /src/__pycache__/SatellitesAugs.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snakers4/spacenet-three/5e1b202b824c6d976298488225078d850c908a8e/src/__pycache__/SatellitesAugs.cpython-35.pyc -------------------------------------------------------------------------------- /src/__pycache__/SatellitesDataset.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snakers4/spacenet-three/5e1b202b824c6d976298488225078d850c908a8e/src/__pycache__/SatellitesDataset.cpython-35.pyc -------------------------------------------------------------------------------- /src/__pycache__/TbLogger.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snakers4/spacenet-three/5e1b202b824c6d976298488225078d850c908a8e/src/__pycache__/TbLogger.cpython-35.pyc -------------------------------------------------------------------------------- /src/__pycache__/UNet.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snakers4/spacenet-three/5e1b202b824c6d976298488225078d850c908a8e/src/__pycache__/UNet.cpython-35.pyc -------------------------------------------------------------------------------- /src/__pycache__/presets.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snakers4/spacenet-three/5e1b202b824c6d976298488225078d850c908a8e/src/__pycache__/presets.cpython-35.pyc -------------------------------------------------------------------------------- /src/__pycache__/sknw.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snakers4/spacenet-three/5e1b202b824c6d976298488225078d850c908a8e/src/__pycache__/sknw.cpython-35.pyc -------------------------------------------------------------------------------- /src/create_8bit_test_images.py: -------------------------------------------------------------------------------- 1 | import pandas as pd 2 | from multiprocessing import Pool 3 | import tqdm 4 | import numpy as np 5 | import os 6 | import glob as glob 7 | from skimage.io import imread, imsave 8 | import osmnx as ox 9 | import numpy as np 10 | import matplotlib 11 | import matplotlib.pyplot as plt 12 | import geopandas as gpd 13 | from osgeo import gdal, ogr, osr 14 | import cv2 15 | import subprocess 16 | import shapely 17 | from shapely.geometry import MultiLineString 18 | from matplotlib.patches import PathPatch 19 | import matplotlib.path 20 | 21 | imgs = [] 22 | 23 | # change this to your data prefix 24 | path_prefix = '../data' 25 | 26 | # default variables from the hosts of the challenge 27 | buffer_meters = 2 28 | burnValue = 150 29 | 30 | # only test folders 31 | folders = ['AOI_2_Vegas_Roads_Test_Public', 32 | 'AOI_3_Paris_Roads_Test_Public', 33 | 'AOI_4_Shanghai_Roads_Test_Public', 34 | 'AOI_5_Khartoum_Roads_Test_Public'] 35 | 36 | # image types 37 | prefix_dict = { 38 | 'mul': 'MUL', 39 | 'muls': 'MUL-PanSharpen', 40 | 'pan': 'PAN', 41 | 'rgbps': 'RGB-PanSharpen', 42 | } 43 | 44 | for folder in folders: 45 | for prefix in prefix_dict.items(): 46 | g = glob.glob(path_prefix+'/{}/{}/*.tif'.format(folder,prefix[1])) 47 | imgs.extend(g) 48 | 49 | img_folders = [(img.split('/')[2]) for img in imgs] 50 | img_subfolders = [(img.split('/')[3]) for img in imgs] 51 | img_files = [(img.split('/')[4]) for img in imgs] 52 | 53 | def create_8bit_test_images(input_data): 54 | img_path = input_data[0] 55 | img_folder = input_data[1] 56 | img_subfolder = input_data[2] 57 | img_file = input_data[3] 58 | 59 | # create paths for masks and 8bit images 60 | bit8_folder = os.path.join(path_prefix,img_folder,img_subfolder+'_8bit') 61 | bit8_path = os.path.join(bit8_folder,img_file) 62 | 63 | if not os.path.exists(bit8_folder): 64 | os.mkdir(bit8_folder) 65 | if os.path.isfile(bit8_path): 66 | os.remove(bit8_path) 67 | 68 | try: 69 | # convert images to 8-bit 70 | convert_to_8Bit(img_path, 71 | bit8_path, 72 | outputPixType='Byte', 73 | outputFormat='GTiff', 74 | rescale_type='rescale', 75 | percentiles=[2,98]) 76 | 77 | except BaseException as e: 78 | print(str(e)) 79 | 80 | return [bit8_folder,bit8_path,img_path,img_folder,img_subfolder,img_file] 81 | 82 | def convert_to_8Bit(inputRaster, outputRaster, 83 | outputPixType='Byte', 84 | outputFormat='GTiff', 85 | rescale_type='rescale', 86 | percentiles=[2, 98]): 87 | ''' 88 | Convert 16bit image to 8bit 89 | rescale_type = [clip, rescale] 90 | if clip, scaling is done strictly between 0 65535 91 | if rescale, each band is rescaled to a min and max 92 | set by percentiles 93 | ''' 94 | 95 | srcRaster = gdal.Open(inputRaster) 96 | cmd = ['gdal_translate', '-ot', outputPixType, '-of', 97 | outputFormat] 98 | 99 | # iterate through bands 100 | for bandId in range(srcRaster.RasterCount): 101 | bandId = bandId+1 102 | band = srcRaster.GetRasterBand(bandId) 103 | if rescale_type == 'rescale': 104 | bmin = band.GetMinimum() 105 | bmax = band.GetMaximum() 106 | # if not exist minimum and maximum values 107 | if bmin is None or bmax is None: 108 | (bmin, bmax) = band.ComputeRasterMinMax(1) 109 | # else, rescale 110 | band_arr_tmp = band.ReadAsArray() 111 | bmin = np.percentile(band_arr_tmp.flatten(), 112 | percentiles[0]) 113 | bmax= np.percentile(band_arr_tmp.flatten(), 114 | percentiles[1]) 115 | 116 | else: 117 | bmin, bmax = 0, 65535 118 | 119 | cmd.append('-scale_{}'.format(bandId)) 120 | cmd.append('{}'.format(bmin)) 121 | cmd.append('{}'.format(bmax)) 122 | cmd.append('{}'.format(0)) 123 | cmd.append('{}'.format(255)) 124 | 125 | cmd.append(inputRaster) 126 | cmd.append(outputRaster) 127 | # print("Conversin command:", cmd) 128 | subprocess.call(cmd) 129 | 130 | return 131 | 132 | input_data = zip(imgs,img_folders,img_subfolders,img_files) 133 | input_data = [item for item in input_data] 134 | 135 | with Pool(10) as p: 136 | bit8_data = list(tqdm.tqdm(p.imap(create_8bit_test_images, input_data), 137 | total=len(input_data))) 138 | # transpose the list 139 | bit8_data = list(map(list, zip(*bit8_data))) 140 | 141 | bit8_df = pd.DataFrame() 142 | 143 | for i,key in enumerate(['bit8_folder','bit8_path','img_path','img_folder','img_subfolder','img_file']): 144 | bit8_df[key] = bit8_data[i] 145 | 146 | bit8_df.to_csv('bit8_test_run.csv') 147 | 148 | -------------------------------------------------------------------------------- /src/final_model_lstrs.py: -------------------------------------------------------------------------------- 1 | import pandas as pd 2 | import numpy as np 3 | import glob as glob 4 | import matplotlib.pyplot as plt 5 | from skimage.io import imread 6 | from tqdm import tqdm 7 | import skimage.transform 8 | import math 9 | from skimage.morphology import skeletonize 10 | import sknw 11 | import cv2 12 | from collections import Counter 13 | from skimage.draw import circle 14 | import os 15 | import argparse 16 | 17 | parser = argparse.ArgumentParser(description='Masks into linestrings') 18 | 19 | parser.add_argument('--folder', '-fld', default='norm_ln34_mul_ps_vegetation_aug_dice_predict', type=str, 20 | metavar='FLD', help='masks containing folder name') 21 | 22 | args = parser.parse_args() 23 | 24 | #globbing 25 | root = '../data' 26 | bit8_folders2 = ['RGB-PanSharpen_8bit'] 27 | test_folders1 = ['AOI_2_Vegas_Roads_Test_Public', 28 | 'AOI_3_Paris_Roads_Test_Public', 29 | 'AOI_4_Shanghai_Roads_Test_Public', 30 | 'AOI_5_Khartoum_Roads_Test_Public'] 31 | mask_folders2_test_pad = [args.folder] 32 | 33 | globdf_test = pd.DataFrame() 34 | for test_folder in test_folders1: 35 | 36 | for bit8_folder in bit8_folders2: 37 | local = pd.DataFrame() 38 | bit8_folder_path = '{}/{}/{}'.format(root, test_folder, bit8_folder) 39 | local['bit8_img'] = glob.glob('{}/*.tif'.format(bit8_folder_path)) 40 | local['bit8_folder'] = bit8_folder 41 | local['test_folder'] = test_folder 42 | globdf_test = pd.concat([globdf_test, local],ignore_index = True) 43 | del local 44 | 45 | globdf_masks_test_pad = pd.DataFrame() 46 | for test_folder in test_folders1: 47 | 48 | 49 | for mask_folder_test in mask_folders2_test_pad: 50 | local = pd.DataFrame() 51 | mask_folder_test_path = '{}/{}/{}'.format(root, test_folder, mask_folder_test) 52 | local['mask_img'] = glob.glob('{}/*.jpg'.format(mask_folder_test_path)) 53 | local['mask_folder_test'] = mask_folder_test 54 | local['test_folder'] = test_folder 55 | globdf_masks_test_pad = pd.concat([globdf_masks_test_pad, local],ignore_index = True) 56 | del local 57 | 58 | globdf_test['bit8_name'] = globdf_test['bit8_img'].apply(lambda x: x.split('/')[len(x.split('/')) - 1]) 59 | globdf_test['img_id'] = globdf_test['bit8_name'].apply(lambda x: x[x.find('AOI'):x.find('.')]) 60 | 61 | globdf_masks_test_pad['mask_name'] = globdf_masks_test_pad['mask_img'].apply(lambda x: x.split('/')[len(x.split('/')) - 1]) 62 | globdf_masks_test_pad['img_id'] = globdf_masks_test_pad['mask_name'].apply(lambda x: x[x.find('AOI'):x.find('.')]) 63 | 64 | globdf_test_pad = globdf_test.merge(globdf_masks_test_pad, how = 'left', on = ['test_folder', 'img_id']) 65 | 66 | # functions 67 | def simplify_edge(ps: np.ndarray, max_distance=3): 68 | """ 69 | Combine multiple points of graph edges to line segments 70 | so distance from points to segments <= max_distance 71 | :param ps: array of points in the edge, including node coordinates 72 | :param max_distance: maximum distance, if exceeded new segment started 73 | :return: ndarray of new nodes coordinates 74 | """ 75 | res_points = [] 76 | cur_idx = 0 77 | # combine points to the single line while distance from the line to any point < max_distance 78 | for i in range(1, len(ps) - 1): 79 | segment = ps[cur_idx:i + 1, :] - ps[cur_idx, :] 80 | angle = -math.atan2(segment[-1, 1], segment[-1, 0]) 81 | ca = math.cos(angle) 82 | sa = math.sin(angle) 83 | # rotate all the points so line is alongside first column coordinate 84 | # and the second col coordinate means the distance to the line 85 | segment_rotated = np.array([[ca, -sa], [sa, ca]]).dot(segment.T) 86 | distance = np.max(np.abs(segment_rotated[1, :])) 87 | if distance > max_distance: 88 | res_points.append(ps[cur_idx, :]) 89 | cur_idx = i 90 | if len(res_points) == 0: 91 | res_points.append(ps[0, :]) 92 | res_points.append(ps[-1, :]) 93 | 94 | return np.array(res_points) 95 | 96 | def simplify_graph(graph, max_distance=2): 97 | """ 98 | :type graph: MultiGraph 99 | """ 100 | all_segments = [] 101 | for (s, e) in graph.edges(): 102 | for _, val in graph[s][e].items(): 103 | ps = val['pts'] 104 | full_segments = np.row_stack([ 105 | graph.node[s]['o'], 106 | ps, 107 | graph.node[e]['o'] 108 | ]) 109 | 110 | segments = simplify_edge(full_segments, max_distance=max_distance) 111 | all_segments.append(segments) 112 | 113 | return all_segments 114 | 115 | def segment_to_linestring(segment): 116 | 117 | if len(segment) < 2: 118 | return [] 119 | 120 | linestring = 'LINESTRING ({})' 121 | sublinestring = '' 122 | 123 | for i, node in enumerate(segment): 124 | 125 | if i == 0: 126 | sublinestring = sublinestring + '{:.1f} {:.1f}'.format(node[1], node[0]) 127 | else: 128 | if node[0] == segment[i - 1][0] and node[1] == segment[i - 1][1]: 129 | if len(segment) == 2: 130 | return [] 131 | continue 132 | if i > 1 and node[0] == segment[i - 2][0] and node[1] == segment[i - 2][1]: 133 | continue 134 | sublinestring = sublinestring + ', {:.1f} {:.1f}'.format(node[1], node[0]) 135 | linestring = linestring.format(sublinestring) 136 | return linestring 137 | 138 | def segmets_to_linestrings(segments): 139 | linestrings = [] 140 | for segment in segments: 141 | linestring = segment_to_linestring(segment) 142 | if len(linestring) > 0: 143 | linestrings.append(linestring) 144 | if len(linestrings) == 0: 145 | linestrings = ['LINESTRING EMPTY'] 146 | return linestrings 147 | 148 | def process_masks(mask_paths): 149 | lnstr_df = pd.DataFrame() 150 | with tqdm(total=len(mask_paths)) as pbar: 151 | for msk_pth in mask_paths: 152 | #print(msk_pth) 153 | msk = imread(msk_pth) 154 | msk = msk[6:1306, 6:1306] 155 | msk_nme = msk_pth.split('/')[3] 156 | img_id = msk_nme[msk_nme.find('AOI'):msk_nme.find('.')] 157 | 158 | # open and skeletonize 159 | thresh = 30 160 | binary = (msk > thresh)*1 161 | 162 | ske = skeletonize(binary).astype(np.uint16) 163 | 164 | # build graph from skeleton 165 | graph = sknw.build_sknw(ske, multi=True) 166 | segments = simplify_graph(graph) 167 | 168 | linestrings = segmets_to_linestrings(segments) 169 | local = pd.DataFrame() 170 | local['WKT_Pix'] = linestrings 171 | local['ImageId'] = img_id 172 | 173 | lnstr_df = pd.concat([lnstr_df, local], ignore_index = True) 174 | pbar.update(1) 175 | return lnstr_df 176 | 177 | # calculating result 178 | print('Processing masks into linestrings...') 179 | 180 | globdf_test_narrow_vegetation = globdf_test_pad[globdf_test_pad['mask_folder_test'] == args.folder].copy() 181 | lstrs_test = process_masks(globdf_test_narrow_vegetation.mask_img) 182 | 183 | lstrs_test = lstrs_test[['ImageId', 'WKT_Pix']].copy() 184 | lstrs_test = lstrs_test.drop_duplicates() 185 | os.makedirs('../solutions', exist_ok=True) 186 | lstrs_test.to_csv('../solutions/norm_test.csv', index = False) 187 | print('Resulting norm_test.csv file is saved under ../solutions/ directory') -------------------------------------------------------------------------------- /src/presets.py: -------------------------------------------------------------------------------- 1 | # select some repsentative data 2 | # following the presets we are going to use 3 | preset_dict = { 4 | 'mul_vegetation': {'width':325,'channel_count':8,'channels':[7,5,3],'subfolder':'MUL'}, 5 | 'mul_urban': {'width':325,'channel_count':8,'channels':[8,7,5],'subfolder':'MUL'}, 6 | 'mul_blackwater': {'width':325,'channel_count':8,'channels':[7,8,1],'subfolder':'MUL'}, 7 | 'mul_ir1': {'width':325,'channel_count':8,'channels':[7,7,7],'subfolder':'MUL'}, 8 | 'mul_ir2': {'width':325,'channel_count':8,'channels':[8,8,8],'subfolder':'MUL'}, 9 | 'mul_naive1': {'width':325,'channel_count':8,'channels':[1,2,3],'subfolder':'MUL'}, 10 | 'mul_naive2': {'width':325,'channel_count':8,'channels':[4,5,6],'subfolder':'MUL'}, 11 | 'mul_naive3': {'width':325,'channel_count':8,'channels':[6,7,8],'subfolder':'MUL'}, 12 | 13 | 'pan': {'width':1300,'channel_count':1,'channels':[1,1,1],'subfolder':'PAN'}, 14 | 15 | 'rgb_ps': {'width':1300,'channel_count':3,'channels':[1,2,3],'subfolder':'RGB-PanSharpen'}, 16 | 17 | 'mul_ps_vegetation': {'width':1300,'channel_count':8,'channels':[7,5,3],'subfolder':'MUL-PanSharpen'}, 18 | 'mul_ps_urban': {'width':1300,'channel_count':8,'channels':[8,7,5],'subfolder':'MUL-PanSharpen'}, 19 | 'mul_ps_blackwater': {'width':1300,'channel_count':8,'channels':[7,8,1],'subfolder':'MUL-PanSharpen'}, 20 | 'mul_ps_ir1': {'width':1300,'channel_count':8,'channels':[7,7,7],'subfolder':'MUL-PanSharpen'}, 21 | 'mul_ps_ir2': {'width':1300,'channel_count':8,'channels':[8,8,8],'subfolder':'MUL-PanSharpen'}, 22 | 'mul_ps_naive1': {'width':1300,'channel_count':8,'channels':[1,2,3],'subfolder':'MUL-PanSharpen'}, 23 | 'mul_ps_naive2': {'width':1300,'channel_count':8,'channels':[4,5,6],'subfolder':'MUL-PanSharpen'}, 24 | 'mul_ps_naive3': {'width':1300,'channel_count':8,'channels':[6,7,8],'subfolder':'MUL-PanSharpen'}, 25 | 26 | 'mul_ps_8channel': {'width':1300,'channel_count':8,'channels':[1,2,3,4,5,6,7,8],'subfolder':'MUL-PanSharpen'}, 27 | 'mul_8channel': {'width':325,'channel_count':8,'channels':[1,2,3,4,5,6,7,8],'subfolder':'MUL'}, 28 | } 29 | -------------------------------------------------------------------------------- /src/readme.md.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "metadata": { 6 | "ExecuteTime": { 7 | "end_time": "2018-02-04T07:16:18.892637Z", 8 | "start_time": "2018-02-04T07:16:18.882802Z" 9 | } 10 | }, 11 | "source": [ 12 | "# 1 Hardware requirements\n", 13 | "\n", 14 | "**Training**\n", 15 | "\n", 16 | "- 6+ core modern CPU (Xeon, i7) for fast image pre-processing;\n", 17 | "- The models were trained on 2 * GeForce 1080 Ti;\n", 18 | "- Training time on my setup ~ **3 hours** for models with 8-bit images as inputs;\n", 19 | "- Disk space - 40GB should be more than enough;\n", 20 | "\n", 21 | "**Inference**\n", 22 | "\n", 23 | "- 6+ core modern CPU (Xeon, i7) for fast image pre-processing;\n", 24 | "- On 2 * GeForce 1080 Ti inference takes **3-5 minutes**;\n", 25 | "- Graph creation takes **5-10 minutes**;\n", 26 | "\n", 27 | "# 2 Preparing and launching the Docker environment\n", 28 | "\n", 29 | "**Clone the repository**\n", 30 | "\n", 31 | "`git clone https://github.com/snakers4/spacenet-three .`\n", 32 | "\n", 33 | "\n", 34 | "**This repository contains 2 Dockerfiles**\n", 35 | "- `/dockerfiles/Dockerfile` - this is the main Dockerfile which was used as environment to run the training and inference scripts\n", 36 | "- `/dockerfiles/Dockerfile2`- this is an additional backup Docker file with newer versions of the drivers and PyTorch, just in case\n", 37 | "\n", 38 | "**Build a Docker image**\n", 39 | "\n", 40 | "`\n", 41 | "cd dockerfiles\n", 42 | "docker build -t aveysov .\n", 43 | "`\n", 44 | "\n", 45 | "**Install the latest nvidia docker**\n", 46 | "\n", 47 | "Follow instructions from [here](https://github.com/NVIDIA/nvidia-docker).\n", 48 | "Please prefer nvidia-docker2 for more stable performance.\n", 49 | "\n", 50 | "\n", 51 | "To test all works fine run:\n", 52 | "\n", 53 | "\n", 54 | "`docker run --runtime=nvidia --rm nvidia/cuda nvidia-smi`\n", 55 | "\n", 56 | "**(IMPORTANT) Run docker container (IMPORTANT)**\n", 57 | "\n", 58 | "Unless you use this exact command (with --shm-size flag) (you can change ports and mounted volumes, of course), then the PyTorch generators **WILL NOT WORK**. \n", 59 | "\n", 60 | "\n", 61 | "- nvidia-docker 2: `docker run --runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=all -it -v /path/to/cloned/repository:/home/keras/notebook -p 8888:8888 -p 6006:6006 --shm-size 8G aveysov`\n", 62 | "- nvidia-docker: `nvidia-docker -it -v /path/to/cloned/repository:/home/keras/notebook -p 8888:8888 -p 6006:6006 --shm-size 8G aveysov`\n", 63 | "\n", 64 | "**Installing project specific software**\n", 65 | "\n", 66 | "1. Exec into the docker machine via `docker exec -it --user root YOUR_CONTAINER_ID /bin/bash`\n", 67 | "2. Run these scripts one after another\n", 68 | "```\n", 69 | "conda install -y -c conda-forge cython\n", 70 | "conda install -y -c conda-forge rasterio\n", 71 | "conda install -y -c conda-forge libgdal\n", 72 | "conda install -y -c conda-forge gdal\n", 73 | "conda install -y -c conda-forge scikit-image\n", 74 | "conda install -y -c conda-forge pyproj\n", 75 | "conda install -y -c conda-forge geopandas\n", 76 | "conda install -y -c conda-forge tqdm\n", 77 | "conda install -y -c conda-forge shapely=1.5.16\n", 78 | "conda install -y -c conda-forge scipy\n", 79 | "conda install -y -c conda-forge networkx=1.11\n", 80 | "conda install -y -c conda-forge fiona\n", 81 | "pip3 install utm\n", 82 | "pip3 install osmnx==0.5.1\n", 83 | "```\n", 84 | "3. Run these scripts one after another\n", 85 | "```\n", 86 | "pip3 install numba\n", 87 | "conda install -y -c conda-forge scikit-image\n", 88 | "```\n", 89 | "\n", 90 | "Steps 2-3 are to ensure compatibility with legacy software from APLS [repository](https://github.com/CosmiQ/apls).\n", 91 | "An alternative to that - is to use pip's requirements.txt in the same order.\n", 92 | "These steps are required to run 8-bit mask creation step from APLS repository and mask creation step.\n", 93 | "If you will be trying to re-do this step - reserve 5-6 hours for experiments.\n", 94 | "\n", 95 | "**To start the stopped container**\n", 96 | "\n", 97 | "\n", 98 | "`docker start -i YOUR_CONTAINER_ID`\n", 99 | "\n", 100 | "\n", 101 | "# 3 Preparing the data and the machine for running scripts\n", 102 | "\n", 103 | "- Ssh into the docker container via `docker exec -it YOUR_CONTAINER_ID`\n", 104 | "- Cd to the root folder of thre repo\n", 105 | "- Dowload the data into `data/`\n", 106 | "- Run these commands:\n", 107 | " - `mkdir src/weights`\n", 108 | " - `mkdir src/tb_logs`\n", 109 | " - `cd scripts` \n", 110 | " - `python3 create_binary_masks.py` \n", 111 | " - `python3 create_8bit_test_images.py` \n", 112 | " \n", 113 | " \n", 114 | "After all of your manipulations your directory should look like:\n", 115 | "\n", 116 | "```\n", 117 | "├── README.md <- The top-level README for developers using this project.\n", 118 | "├── data\n", 119 | "│ ├── AOI_2_Vegas_Roads_Test_Public <- Test set for each city\n", 120 | "│ └── AOI_2_Vegas_Roads_Train <- Train set for each city\n", 121 | "│ ├─ geojson\n", 122 | "│ ├─ summaryData\n", 123 | "│ ├─ MUL\n", 124 | "│ ├─ RGB-PanSharpen\n", 125 | "│ ├─ PAN\n", 126 | "│ ├─ MUL-PanSharpen\n", 127 | "│ ├─ MUL-PanSharpen_8bit\n", 128 | "│ ├─ RGB-PanSharpen_8bit\n", 129 | "│ ├─ PAN_8bit\n", 130 | "│ ├─ MUL_8bit\n", 131 | "│ ├─ MUL-PanSharpen_mask\n", 132 | "│ ├─ RGB-PanSharpen_mask\n", 133 | "│ ├─ PAN_mask\n", 134 | "│ ├─ MUL_mask\n", 135 | "│ └─ RGB-PanSharpen_mask\n", 136 | "│ │\n", 137 | "│ ...\n", 138 | "│ │\n", 139 | "│ ├── AOI_5_Khartoum_Roads_Test_Public <- Test set for each city\n", 140 | "│ └── AOI_5_Khartoum_Roads_Train <- Train set for each city\n", 141 | "│\n", 142 | "├── dockerfiles <- A folder with Dockerfiles\n", 143 | "│\n", 144 | "├── src <- Source code\n", 145 | "│\n", 146 | "└── scripts <- One-off preparation scripts\n", 147 | "```\n", 148 | "\n", 149 | "# 4 Training the model\n", 150 | "\n", 151 | "If all is ok, then use the following command to train the model\n", 152 | "\n", 153 | "- Ssh into the docker container via `docker exec -it YOUR_CONTAINER_ID`\n", 154 | "- Cd to the root folder of thre repo\n", 155 | "- `cd src`\n", 156 | "- optional - turn on tensorboard for monitoring progress `tensorboard --logdir='satellites_roads/src/tb_logs' --port=6006` via jupyter notebook console or via tmux + docker exec (model converges in 30-40 epochs)\n", 157 | "- then\n", 158 | "```\n", 159 | "echo 'python3 train_satellites.py \\\n", 160 | "\t--arch linknet34 --batch-size 6 \\\n", 161 | "\t--imsize 1280 --preset mul_ps_vegetation --augs True \\\n", 162 | "\t--workers 6 --epochs 40 --start-epoch 0 \\\n", 163 | "\t--seed 42 --print-freq 20 \\\n", 164 | "\t--lr 1e-3 --optimizer adam \\\n", 165 | "\t--tensorboard True --lognumber ln34_mul_ps_vegetation_aug_dice' > train.sh\n", 166 | "```\n", 167 | "- `sh train.sh`\n", 168 | "\n", 169 | "# 5 Predicting masks\n", 170 | "\n", 171 | "- Ssh into the docker container via `docker exec -it YOUR_CONTAINER_ID`\n", 172 | "- Cd to the root folder of thre repo\n", 173 | "- `cd src`\n", 174 | "- then\n", 175 | "``` \n", 176 | "echo 'python3 train_satellites.py\\\n", 177 | "\t--arch linknet34 --batch-size 12\\\n", 178 | "\t--imsize 1312 --preset mul_ps_vegetation --augs True\\\n", 179 | "\t--workers 6 --epochs 50 --start-epoch 0\\\n", 180 | "\t--seed 42 --print-freq 10\\\n", 181 | "\t--lr 1e-3 --optimizer adam\\\n", 182 | "\t--lognumber norm_ln34_mul_ps_vegetation_aug_dice_predict\\\n", 183 | "\t--predict --resume weights/norm_ln34_mul_ps_vegetation_aug_dice_best.pth.tar\\' > predict.sh\n", 184 | "```\n", 185 | "- `sh predict.sh`\n", 186 | "\n", 187 | "\n", 188 | "# 6 Creating graphs and submission files\n", 189 | "`cd` into `src` directory and execute `final_model_lstrs.py` script as follows:\n", 190 | "```\n", 191 | "docker exec -it YOUR_CONTAINER_ID sh -c \"cd path/to/src && python3 final_model_lstrs.py --folder norm_ln34_mul_ps_vegetation_aug_dice_predict\"\n", 192 | "```\n", 193 | "`folder` argument is for masks containing folder name, default is `norm_ln34_mul_ps_vegetation_aug_dice_predict`.\n", 194 | "Scipt saves a file called `norm_test.csv` into `../solutions` directory. The resulting file is used then as a submission file.\n", 195 | "\n", 196 | "# 7 Additional notes\n", 197 | "\n", 198 | "- You can run training and inference on the presets from `/src/presets.py`;\n", 199 | "- So the model can be evaluated on RGB-PS images and / or 8-channel images as well;\n", 200 | "- This script, for example will train an 8-channel model:\n", 201 | "```\n", 202 | "python3 train_satellites.py \\\n", 203 | "\t--arch linknet34 --batch-size 6 \\\n", 204 | "\t--imsize 1280 --preset mul_ps_vegetation --augs True \\\n", 205 | "\t--workers 6 --epochs 40 --start-epoch 0 \\\n", 206 | "\t--seed 42 --print-freq 20 \\\n", 207 | "\t--lr 1e-3 --optimizer adam \\\n", 208 | "\t--tensorboard True --lognumber ln34_mul_ps_vegetation_aug_dice\n", 209 | "```\n", 210 | "\n", 211 | "To train an 8-channel model you should also replace mean and std settings in the `src/SatellitesAugs.py`\n", 212 | "\n", 213 | "```\n", 214 | "# 8-channel settings\n", 215 | "# normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5],\n", 216 | "# std=[1, 1, 1, 1, 1, 1, 1, 1])\n", 217 | "\n", 218 | "# 3 channel settings\n", 219 | "normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n", 220 | " std=[0.229, 0.224, 0.225])\n", 221 | "```\n", 222 | "\n", 223 | "- 16-bit images are also supported:\n", 224 | "\n", 225 | "This snippet is commented in `src/SatellitesAugs.py`\n", 226 | "\n", 227 | "```\n", 228 | "# version compatible with 16-bit images\n", 229 | "\"\"\" \n", 230 | "class ImgAugAugs(object):\n", 231 | " def __call__(self,\n", 232 | " image):\n", 233 | " global seed \n", 234 | " \n", 235 | " # poor man's flipping\n", 236 | " if seed%2==0:\n", 237 | " image = np.fliplr(image)\n", 238 | " elif seed%4==0:\n", 239 | " image = np.fliplr(image)\n", 240 | " image = np.flipud(image)\n", 241 | " \n", 242 | " # poor man's affine transformations\n", 243 | " image = rotate(image,\n", 244 | " angle=seed,\n", 245 | " resize=False,\n", 246 | " clip=True,\n", 247 | " preserve_range=True) \n", 248 | "\n", 249 | " return image\n", 250 | "\"\"\"\n", 251 | "```\n", 252 | "\n", 253 | "- Also the following models are supported\n", 254 | " - `unet11` (VGG11 + Unet)\n", 255 | " - `linknet50` (ResNet50 + LinkNet, 3 layers)\n", 256 | " - `linknet50_full` (ResNet50 + LinkNet, 4 layers)\n", 257 | " - `linknext` (ResNext-101-32 + LinkNet, 4 layers)\n", 258 | " \n", 259 | "- Also in the repo you can find scripts to generate wide masks (i.e. wide roads have varying width) and layered masks (paved / non-paved). There are scripts in the `src/SatellitesDataset.py` that support that. They basically just replace some paths; \n", 260 | " \n", 261 | "# 8 Juputer notebooks\n", 262 | "\n", 263 | "Use these notebooks on your own risk!\n", 264 | "\n", 265 | "- `src/experiments.ipynb` - general debugging notebook with new models / generators / etc\n", 266 | "- `src/play_w_stuff.ipynb` - visualizing the solutions\n", 267 | "- `src/pipeline_experiments.ipynb`- some minor experiments with the graph creation script" 268 | ] 269 | }, 270 | { 271 | "cell_type": "code", 272 | "execution_count": null, 273 | "metadata": {}, 274 | "outputs": [], 275 | "source": [ 276 | "*" 277 | ] 278 | } 279 | ], 280 | "metadata": { 281 | "kernelspec": { 282 | "display_name": "Python 3", 283 | "language": "python", 284 | "name": "python3" 285 | }, 286 | "language_info": { 287 | "codemirror_mode": { 288 | "name": "ipython", 289 | "version": 3 290 | }, 291 | "file_extension": ".py", 292 | "mimetype": "text/x-python", 293 | "name": "python", 294 | "nbconvert_exporter": "python", 295 | "pygments_lexer": "ipython3", 296 | "version": "3.5.4" 297 | }, 298 | "toc": { 299 | "nav_menu": {}, 300 | "number_sections": true, 301 | "sideBar": true, 302 | "skip_h1_title": false, 303 | "toc_cell": false, 304 | "toc_position": {}, 305 | "toc_section_display": "block", 306 | "toc_window_display": false 307 | } 308 | }, 309 | "nbformat": 4, 310 | "nbformat_minor": 2 311 | } 312 | -------------------------------------------------------------------------------- /src/resnext_features/__pycache__/resnext101_32x4d_features.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snakers4/spacenet-three/5e1b202b824c6d976298488225078d850c908a8e/src/resnext_features/__pycache__/resnext101_32x4d_features.cpython-35.pyc -------------------------------------------------------------------------------- /src/resnext_features/__pycache__/resnext101_64x4d_features.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snakers4/spacenet-three/5e1b202b824c6d976298488225078d850c908a8e/src/resnext_features/__pycache__/resnext101_64x4d_features.cpython-35.pyc -------------------------------------------------------------------------------- /src/script.sh: -------------------------------------------------------------------------------- 1 | python3 train_gapnet.py \ 2 | --arch gapnet18 --batch-size 4 --dilation 40 \ 3 | --imsize 640 --preset mul_ps_vegetation --augs True \ 4 | --workers 6 --epochs 50 --start-epoch 0 \ 5 | --seed 42 --print-freq 20 \ 6 | --lr 1e-4 --optimizer adam \ 7 | --tensorboard True --tensorboard_images True \ 8 | --lognumber gapnet_img18_vegetation_dilation_40 9 | -------------------------------------------------------------------------------- /src/sknw.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from numba import jit 3 | import networkx as nx 4 | 5 | # author https://github.com/yxdragon/sknw/blob/master/sknw/sknw.py 6 | 7 | # get neighbors d index 8 | def neighbors(shape): 9 | dim = len(shape) 10 | block = np.ones([3]*dim) 11 | block[tuple([1]*dim)] = 0 12 | idx = np.where(block>0) 13 | idx = np.array(idx, dtype=np.uint8).T 14 | idx = np.array(idx-[1]*dim) 15 | acc = np.cumprod((1,)+shape[::-1][:-1]) 16 | return np.dot(idx, acc[::-1]) 17 | 18 | @jit # my mark 19 | def mark(img): # mark the array use (0, 1, 2) 20 | nbs = neighbors(img.shape) 21 | img = img.ravel() 22 | for p in range(len(img)): 23 | if img[p]==0:continue 24 | s = 0 25 | for dp in nbs: 26 | if img[p+dp]!=0:s+=1 27 | if s==2:img[p]=1 28 | else:img[p]=2 29 | 30 | @jit # trans index to r, c... 31 | def idx2rc(idx, acc): 32 | rst = np.zeros((len(idx), len(acc)), dtype=np.int16) 33 | for i in range(len(idx)): 34 | for j in range(len(acc)): 35 | rst[i,j] = idx[i]//acc[j] 36 | idx[i] -= rst[i,j]*acc[j] 37 | rst -= 1 38 | return rst 39 | 40 | @jit # fill a node (may be two or more points) 41 | def fill(img, p, num, nbs, acc, buf): 42 | back = img[p] 43 | img[p] = num 44 | buf[0] = p 45 | cur = 0; s = 1; 46 | 47 | while True: 48 | p = buf[cur] 49 | for dp in nbs: 50 | cp = p+dp 51 | if img[cp]==back: 52 | img[cp] = num 53 | buf[s] = cp 54 | s+=1 55 | cur += 1 56 | if cur==s:break 57 | return idx2rc(buf[:s], acc) 58 | 59 | @jit # trace the edge and use a buffer, then buf.copy, if use [] numba not works 60 | def trace(img, p, nbs, acc, buf): 61 | c1 = 0; c2 = 0; 62 | newp = 0 63 | cur = 0 64 | 65 | while True: 66 | buf[cur] = p 67 | img[p] = 0 68 | cur += 1 69 | for dp in nbs: 70 | cp = p + dp 71 | if img[cp] >= 10: 72 | if c1==0:c1=img[cp] 73 | else: c2 = img[cp] 74 | if img[cp] == 1: 75 | newp = cp 76 | p = newp 77 | if c2!=0:break 78 | return (c1-10, c2-10, idx2rc(buf[:cur], acc)) 79 | 80 | @jit # parse the image then get the nodes and edges 81 | def parse_struc(img): 82 | nbs = neighbors(img.shape) 83 | acc = np.cumprod((1,)+img.shape[::-1][:-1])[::-1] 84 | img = img.ravel() 85 | pts = np.array(np.where(img==2))[0] 86 | buf = np.zeros(131072, dtype=np.int64) 87 | num = 10 88 | nodes = [] 89 | for p in pts: 90 | if img[p] == 2: 91 | nds = fill(img, p, num, nbs, acc, buf) 92 | num += 1 93 | nodes.append(nds) 94 | 95 | edges = [] 96 | for p in pts: 97 | for dp in nbs: 98 | if img[p+dp]==1: 99 | edge = trace(img, p+dp, nbs, acc, buf) 100 | edges.append(edge) 101 | return nodes, edges 102 | 103 | # use nodes and edges build a networkx graph 104 | def build_graph(nodes, edges, multi=False): 105 | graph = nx.MultiGraph() if multi else nx.Graph() 106 | for i in range(len(nodes)): 107 | graph.add_node(i, pts=nodes[i], o=nodes[i].mean(axis=0)) 108 | for s,e,pts in edges: 109 | l = np.linalg.norm(pts[1:]-pts[:-1], axis=1).sum() 110 | graph.add_edge(s,e, pts=pts, weight=l) 111 | return graph 112 | 113 | def buffer(ske): 114 | buf = np.zeros(tuple(np.array(ske.shape)+2), dtype=np.uint16) 115 | buf[tuple([slice(1,-1)]*buf.ndim)] = ske 116 | return buf 117 | 118 | def build_sknw(ske, multi=False): 119 | buf = buffer(ske) 120 | mark(buf) 121 | nodes, edges = parse_struc(buf) 122 | return build_graph(nodes, edges, multi) 123 | 124 | # draw the graph 125 | def draw_graph(img, graph, cn=255, ce=128): 126 | acc = np.cumprod((1,)+img.shape[::-1][:-1])[::-1] 127 | img = img.ravel() 128 | for idx in graph.nodes(): 129 | pts = graph.node[idx]['pts'] 130 | img[np.dot(pts, acc)] = cn 131 | for (s, e) in graph.edges(): 132 | eds = graph[s][e] 133 | for i in eds: 134 | pts = eds[i]['pts'] 135 | img[np.dot(pts, acc)] = ce 136 | 137 | if __name__ == '__main__': 138 | g = nx.MultiGraph() 139 | g.add_nodes_from([1,2,3,4,5]) 140 | g.add_edges_from([(1,2),(1,3),(2,3),(4,5),(5,4)]) 141 | print(g.nodes()) 142 | print(g.edges()) 143 | a = g.subgraph(1) 144 | print('d') 145 | print(a) 146 | print('d') -------------------------------------------------------------------------------- /src/test.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snakers4/spacenet-three/5e1b202b824c6d976298488225078d850c908a8e/src/test.jpg -------------------------------------------------------------------------------- /src/untitled.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snakers4/spacenet-three/5e1b202b824c6d976298488225078d850c908a8e/src/untitled.txt --------------------------------------------------------------------------------