IdaVision/visualization_utils.py at main · njc1583/IdaVision · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
import matplotlib.pyplot as plt
import numpy as np

from frame_utils import *
from VidFrame import *

def view_vidframe_segmentations(vid_frame, N_background_segments, background_compactness):
    vid_frame.calculateMotionSuperpixelMasks()

    background_labels = vid_frame.calculateBackgroundLabels(N_background_segments,
        background_compactness)

    background_boundary = get_boundary_segments(vid_frame.vid_frame,
                                                background_labels)

    plt.figure()
    plt.imshow(vid_frame.label_frame)
    plt.show()

    plt.figure()
    plt.imshow(background_labels)
    plt.show()

    N_background_segments = np.unique(background_labels)

    plt.figure()
    plt.imshow(background_boundary)
    plt.show()

    _, obj_labels = vid_frame.calculateForegroundLabels()

    plt.figure()
    plt.imshow(obj_labels)
    plt.show()

def visualize_images(images, nrows, ncols):
    if len(images.shape) == 4:
        N, H, W, C = images.shape
    else:
        N, H, W = images.shape

    assert N == nrows * ncols

    fig, ax = plt.subplots(nrows, ncols)

    for i, a in enumerate(ax.ravel()):
        a.imshow(images[i])
        a.set_axis_off()

    return fig


def segment_frames(frame_imgs, frames, nrows, ncols, n_segments=500,
                   compactness=10, mask=None):

    N, H, W, C = frame_imgs.shape

    assert nrows * ncols == len(frames)

    segmented_frames = np.zeros((len(frames), H, W, C), dtype=frame_imgs.dtype)

    for fidx, frame in enumerate(frames):
        segment_labels = slic_segment_image(
            frame_imgs[frame], n_segments, compactness, mask)

        segmented_img = get_boundary_segments(
            frame_imgs[frame], segment_labels)

        segmented_frames[fidx] = segmented_img

    return visualize_images(segmented_frames, nrows, ncols)