Function postprocess_kmeans

postprocessing/postprocess_saliency_maps.py:103–150

(args)

        cv2.imwrite(output, segmented_image)

def postprocess_kmeans(args):
    
    files = os.listdir(args.sal_path)

    if (not os.path.exists(args.output_path)):
        os.makedirs(args.output_path)

    for file in tqdm(files):

        kmeans = KMeans(n_clusters=2,random_state=10)
        attn_weights = cv2.imread(args.sal_path+'/'+file, 0) / 255
        h, w = attn_weights.shape
        image = cv2.resize(attn_weights, (256, 256),interpolation=cv2.INTER_NEAREST)
        flat_image = image.reshape(-1, 1)

        labels = kmeans.fit_predict(flat_image)

        segmented_image = labels.reshape(256, 256)

        centroids = kmeans.cluster_centers_.flatten()

        # Identify the background cluster (assuming it has the lowest centroid value)
        background_cluster = np.argmin(centroids)

        # Mark background pixels as 0 and foreground pixels as 1
        segmented_image = np.where(segmented_image == background_cluster, 0, 1)

        segmented_image = cv2.resize(segmented_image, (w,h),interpolation=cv2.INTER_NEAREST)
        segmented_image = segmented_image.astype(np.uint8)*255

        nb_blobs, im_with_separated_blobs, stats, _ = cv2.connectedComponentsWithStats(segmented_image)
        sizes = stats[:, cv2.CC_STAT_AREA]
        
        # Sort sizes (ignoring the background at index 0)
        sorted_sizes = sorted(sizes[1:], reverse=True) 
        
        # Determine the top K sizes
        top_k_sizes = sorted_sizes[:args.num_contours]
        
        im_result = np.zeros_like(im_with_separated_blobs)
        
        for index_blob in range(1, nb_blobs):
            if sizes[index_blob] in top_k_sizes:
                im_result[im_with_separated_blobs == index_blob] = 255
        
        segmented_image = im_result

        cv2.imwrite(args.output_path+'/'+file, segmented_image)


def get_parser():