File size: 2,665 Bytes
844f7c0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
import os
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"
os.environ["OMP_NUM_THREADS"] = "1"
import numpy as np
np.random.seed(1)
import random
random.seed(1)
import pandas as pd
import cv2
import timeit
from os import path, makedirs, listdir
import sys
sys.setrecursionlimit(10000)
from multiprocessing import Pool
from skimage.morphology import square, dilation, erosion
from skimage.segmentation import watershed
from skimage import io
from shapely.wkt import loads
from shapely.geometry import mapping, Polygon
# import matplotlib.pyplot as plt
# import seaborn as sns
import json
import os
os.chdir('xview2')
masks_dir = 'masks'
train_dirs = ['train', 'tier3']
def mask_for_polygon(poly, im_size=(1024, 1024)):
img_mask = np.zeros(im_size, np.uint8)
int_coords = lambda x: np.array(x).round().astype(np.int32)
exteriors = [int_coords(poly.exterior.coords)]
interiors = [int_coords(pi.coords) for pi in poly.interiors]
cv2.fillPoly(img_mask, exteriors, 1)
cv2.fillPoly(img_mask, interiors, 0)
return img_mask
damage_dict = {
"no-damage": 1,
"minor-damage": 2,
"major-damage": 3,
"destroyed": 4,
"un-classified": 1 # ?
}
def process_image(json_file):
js1 = json.load(open(json_file))
js2 = json.load(open(json_file.replace('_pre_disaster', '_post_disaster')))
msk = np.zeros((1024, 1024), dtype='uint8')
msk_damage = np.zeros((1024, 1024), dtype='uint8')
for feat in js1['features']['xy']:
poly = loads(feat['wkt'])
_msk = mask_for_polygon(poly)
msk[_msk > 0] = 255
for feat in js2['features']['xy']:
poly = loads(feat['wkt'])
subtype = feat['properties']['subtype']
_msk = mask_for_polygon(poly)
msk_damage[_msk > 0] = damage_dict[subtype]
cv2.imwrite(json_file.replace('/labels/', '/masks/').replace('_pre_disaster.json', '_pre_disaster.png'), msk, [cv2.IMWRITE_PNG_COMPRESSION, 9])
cv2.imwrite(json_file.replace('/labels/', '/masks/').replace('_pre_disaster.json', '_post_disaster.png'), msk_damage, [cv2.IMWRITE_PNG_COMPRESSION, 9])
if __name__ == '__main__':
t0 = timeit.default_timer()
all_files = []
for d in train_dirs:
makedirs(path.join(d, masks_dir), exist_ok=True)
for f in sorted(listdir(path.join(d, 'images'))):
if '_pre_disaster.png' in f:
all_files.append(path.join(d, 'labels', f.replace('_pre_disaster.png', '_pre_disaster.json')))
with Pool() as pool:
_ = pool.map(process_image, all_files)
elapsed = timeit.default_timer() - t0
print('Time: {:.3f} min'.format(elapsed / 60)) |