nreHieW/SoccerNet_Field_Segmentation

def show_item(item):
    fig, axs = plt.subplots(nrows = 1, ncols = 4, figsize = (20, 4))
    axs[0].imshow(item['image'])
    axs[0].set_title("Image")
    axs[0].axis('off')

    axs[1].imshow(overlay_mask(item['image'], item['outlines']))
    axs[1].set_title("Outlines")
    axs[1].axis('off')

    axs[2].imshow(show_segments(item['segments']))
    axs[2].set_title("Segments")
    axs[2].axis('off')
            # PART 3: GET MASK OUTLINES
    kernel = np.array([[0, 1, 0],
                        [1, -4, 1],
                        [0, 1, 0]])
    segments = np.array(item['segments']).astype(np.uint8)
    class_edges = np.zeros(segments.shape[1:], dtype=int)

    for i in range(segments.shape[0]):
        edge = convolve(segments[i], kernel, mode='constant', cval=0)
        edge_detected = edge != 0
        class_edges[edge_detected] = i

    axs[3].imshow(overlay_mask(item['image'], class_edges))
    axs[3].set_title("Segments Outlines")
    axs[3].axis('off')
    if item['is_bad']:
        s = f"Bad ID: {item['id']}"
    else:
        s = f"ID: {item['id']}"
    fig.suptitle(s, fontsize = 8)
    plt.subplots_adjust(hspace = -0.2, wspace = -0.05)
    plt.show()

show_item(dataset['train'][99])