Блог

Я тренирую U-сеть для извлечения площади зданий из спутниковых снимков. Результаты неплохие, но я хочу, чтобы контуры фигур внутри изображения были четкими.

Чтобы улучшить его, я пытаюсь использовать весовую карту контуров или границ фигуры внутри изображения. Поэтому я пытаюсь построить карту весов с высокими значениями — например, 10 — на границах и значениями, уменьшающимися с обеих сторон. Но я еще не знал, как это сделать.

Я адаптировал из другого кода решение, которое работает для этого. Вот код:

 def unet_weight_map(y, wc=None, w0 = 10, sigma = 20):

"""
Generate weight maps as specified in the U-Net paper
for boolean mask.

"U-Net: Convolutional Networks for Biomedical Image Segmentation"
https://arxiv.org/pdf/1505.04597.pdf

Parameters
----------
mask: Numpy array
    2D array of shape (image_height, image_width) representing binary mask
    of objects.
wc: dict
    Dictionary of weight classes.
w0: int
    Border weight parameter.
sigma: int
    Border width parameter.

Returns
-------
Numpy array
    Training weights. A 2D array of shape (image_height, image_width).
"""
y = y.reshape(y.shape[0], y.shape[1])
labels = label(y)
no_labels = labels == 0
label_ids = sorted(np.unique(labels))
if len(label_ids) > 0:
    distances = np.zeros((y.shape[0], y.shape[1], len(label_ids)))
    for i, label_id in enumerate(label_ids):
        distances[:,:,i] = distance_transform_edt(labels != label_id)
    distances = np.sort(distances, axis=2)
    d1 = distances[:,:,0]
    d2 = distances[:,:,1]
    w = w0 * np.exp(-1/2*((d1   d2) / sigma)**2) * no_labels
else:
    w = np.zeros_like(y)
if wc:
    class_weights = np.zeros_like(y)
    for k, v in wc.items():
        class_weights[y == k] = v
    w = w   class_weights
return w


wc = {
    0: 0, # background
    1: 1  # objects
}

w = unet_weight_map(img, wc)
 

Ввод:

Выход:

Если у кого-то есть лучшее решение, пожалуйста!