File size: 2,380 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 |
import os
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"
os.environ["OMP_NUM_THREADS"] = "1"
from os import path, makedirs, listdir
import sys
import numpy as np
np.random.seed(1)
import random
random.seed(1)
import timeit
import cv2
from skimage.morphology import square, dilation
cv2.setNumThreads(0)
cv2.ocl.setUseOpenCL(False)
pred_folders = ['dpn92cls_0_tuned', 'dpn92cls_1_tuned', 'dpn92cls_2_tuned', 'res34cls2_0_tuned', 'res34cls2_1_tuned', 'res34cls2_2_tuned',
'res50cls_0_tuned', 'res50cls_1_tuned', 'res50cls_2_tuned', 'se154cls_0_tuned', 'se154cls_1_tuned', 'se154cls_2_tuned']
loc_folders = ['pred50_loc_tuned', 'pred92_loc_tuned', 'pred34_loc', 'pred154_loc']
_thr = [0.38, 0.13, 0.14]
if __name__ == '__main__':
t0 = timeit.default_timer()
pre_file = sys.argv[1]
post_file = sys.argv[2]
loc_pred_file = sys.argv[3]
cls_pred_file = sys.argv[4]
loc_fn = os.path.basename(loc_pred_file)
loc_fn = '{0}'.format(loc_fn + '_part1.png')
cls_fn = os.path.basename(cls_pred_file)
preds = []
for d in pred_folders:
msk1 = cv2.imread(path.join(d, '{0}'.format(cls_fn + '_part1.png')), cv2.IMREAD_UNCHANGED)
msk2 = cv2.imread(path.join(d, '{0}'.format(cls_fn + '_part2.png')), cv2.IMREAD_UNCHANGED)
msk = np.concatenate([msk1, msk2[..., 1:]], axis=2)
preds.append(msk)
preds = np.asarray(preds).astype('float').sum(axis=0) / len(pred_folders) / 255
loc_preds = []
for d in loc_folders:
msk = cv2.imread(path.join(d, loc_fn), cv2.IMREAD_UNCHANGED)
loc_preds.append(msk)
loc_preds = np.asarray(loc_preds).astype('float').sum(axis=0) / len(loc_folders) / 255
msk_dmg = preds[..., 1:].argmax(axis=2) + 1
msk_loc = (1 * ((loc_preds > _thr[0]) | ((loc_preds > _thr[1]) & (msk_dmg > 1) & (msk_dmg < 4)) | ((loc_preds > _thr[2]) & (msk_dmg > 1)))).astype('uint8')
msk_dmg = msk_dmg * msk_loc
_msk = (msk_dmg == 2)
if _msk.sum() > 0:
_msk = dilation(_msk, square(5))
msk_dmg[_msk & msk_dmg == 1] = 2
msk_dmg = msk_dmg.astype('uint8')
cv2.imwrite(loc_pred_file, msk_loc, [cv2.IMWRITE_PNG_COMPRESSION, 9])
cv2.imwrite(cls_pred_file, msk_dmg, [cv2.IMWRITE_PNG_COMPRESSION, 9])
elapsed = timeit.default_timer() - t0
print('Time: {:.3f} min'.format(elapsed / 60)) |