raw_to_mask.py 文件源码

def mask_to_poly(image_id):
    preds = joblib.load('raw_preds/raw_blend5/{}.pkl'.format(image_id))
    size = preds.shape[1]
    if n_out == 10:
#        preds = (preds > 0.3).astype(np.uint8)

        thresholds = np.array([0.4, 0.4, 0.4, 0.4, 0.8,
                               0.4, 0.4, 0.4, 0.1, 0.1]).reshape((10, 1))
        preds = (preds.reshape((10, -1)) > thresholds).reshape((10, size, size))
        preds = preds.astype(np.uint8)
    else:
        preds = np.argmax(preds, axis=0)
        preds = unsoft(preds)

    rg = colorize_raster(preds.transpose((1, 2, 0)))
#    cv2.imwrite('1.png', rg)
    size = 900
    rg = cv2.resize(rg, (size, size))
#    cv2.imshow('mask', rg)
#    cv2.waitKey()
    im = get_rgb_image(image_id, size, size)
    rg = np.hstack([rg, im])
    cv2.imwrite('raw_temp5_1/{}.png'.format(image_id), rg)

    shs = []
    for i in range(10):
        mask = preds[i]

        y_sf, x_sf = get_scale_factor(image_id, mask.shape[0], mask.shape[1])
        y_sf = 1. / y_sf
        x_sf = 1. / x_sf

        sh = polygonize_cv(mask)
#        sh = polygonize_sk((mask>0)*255, 0)
#        sh = (sh1.buffer(0).intersection(sh2.buffer(0))).buffer(0)

#        if not sh.is_valid:
#            sh = sh.buffer(0)
        sh = affinity.scale(sh, xfact=x_sf, yfact=y_sf, origin=(0, 0, 0))
        shs.append(sh)
    return shs