生成mask之间空隙权重

import random
from collections import defaultdict


import sys
import logging


import numpy as np
import os
import cv2
import time
from scipy import ndimage


TRAIN_PATH = 'D:/2018 Data Science Bowl/stage1_train/'
TEST_PATH = 'D:/2018 Data Science Bowl/stage1_test'


def calculate_unet_background_weight(merged_mask, masks, w0=10, q=5,):
    weight = np.zeros(merged_mask.shape)
    # calculate weight for important pixels
    distances = np.array([ndimage.distance_transform_edt(m==0) for m in masks])
    shortest_dist = np.sort(distances, axis=0)
    # distance to the border of the nearest cell 
    d1 = shortest_dist[0]
    # distance to the border of the second nearest cell
    d2 = shortest_dist[1] if len(shortest_dist) > 1 else np.zeros(d1.shape)


    weight = w0 * np.exp(-(d1+d2)**2/(2*q**2)).astype(np.float32)
    weight = 1 + (merged_mask == 0) * weight
    return weight




masks=[]
for root,dirs,files in os.walk(TRAIN_PATH):
    for dir_ in dirs:
        masks=[]
        dir__=TRAIN_PATH+dir_+'/'
        msk_dir=dir__+'/'+'masks'+'/'
        for mask_file_root,mask_folder,mask_files in os.walk(msk_dir):
            for mask_file in mask_files:
                mask=cv2.imread(msk_dir+'/'+mask_file,0)
                masks.append(mask/255)
        msk_shape=masks[0].shape
        multi_mask=np.zeros([msk_shape[0],msk_shape[1],len(masks)])
        merge_mask=np.zeros(msk_shape)
        for i in range(len(masks)):
            multi_mask[:,:,i]=masks[i]
            merge_mask=merge_mask+masks[i]
        weight=calculate_unet_background_weight(merge_mask, masks, 10, 5,)
        imgpath=dir__+'/'+'images'+'/'+dir_+'.png'
        img=cv2.imread(imgpath)
        dir__+'/'+'masks'+'/'
cv2.imshow('gap',np.uint8(weight*25)),cv2.imshow('img',img), cv2.imshow('margedMask',merge_mask),cv2.waitKey(0)