Skip to content

Plots

Awive provides some basic plotting capabilities.

awive.plots.draw_velocities(velocities, image, max_velocity=5)

Draw velocities arrows in an image.

The color of those arrows depends on the velocity value. max is red and min is blue.

Parameters:

Name Type Description Default
velocities dict[str, dict[str, float]]

dictionary of velocities. keys are the velocity and values are dicts with three keys: "velocity", "position", and "count". velocity: velocity value in p/s. position: position of the velocity in pixels in y column count: number of occurrences of the velocity.

required
image NDArray[uint8]

image to draw the velocities on.

required
max_velocity float

maximum velocity value. if a velocity is higher than this value, it will be considered as this value.

5

Returns:

Type Description
NDArray[uint8]

an image with the velocities arrows.

Source code in awive/plots.py
 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
def draw_velocities(
    velocities: dict[str, dict[str, float]],
    image: npt.NDArray[np.uint8],
    max_velocity: float = 5,
) -> npt.NDArray[np.uint8]:
    """Draw velocities arrows in an image.

    The color of those arrows depends on the velocity value. max is red and
    min is blue.

    Args:
        velocities: dictionary of velocities. keys are the velocity and values
            are dicts with three keys: "velocity", "position", and "count".
            velocity: velocity value in p/s.
            position: position of the velocity in pixels in y column
            count: number of occurrences of the velocity.
        image: image to draw the velocities on.
        max_velocity: maximum velocity value. if a velocity is higher than this
            value, it will be considered as this value.

    Returns:
        an image with the velocities arrows.
    """
    # Extract all velocity values
    velocity_values = [v["velocity"] for v in velocities.values()]

    # Determine min and max velocities
    min_velocity = min(min(velocity_values), 0)
    max_velocity = max(max(velocity_values), max_velocity)
    max_length = image.shape[1] // 4
    thickness = image.shape[0] // 100

    # Iterate over the velocities and draw arrows
    for velocity_info in velocities.values():
        velocity = velocity_info["velocity"]
        position = (2 * image.shape[0] // 3, velocity_info["position"])

        # Normalize velocity between 0 and 1
        norm_velocity = (velocity - min_velocity) / (max_velocity - min_velocity)
        # Calculate the end position of the arrow based on the velocity
        end_position = (position[0] - int(norm_velocity * max_length), position[1])

        # Map normalized velocity to color
        color = cv2.applyColorMap(
            np.array([[int(255 - norm_velocity * 255)]], dtype=np.uint8),
            cv2.COLORMAP_JET,
        ).flatten().tolist()

        # Draw the arrow
        cv2.arrowedLine(  # pyright: ignore[reportCallIssue]
            image,
            position,  # pyright: ignore[reportArgumentType]
            end_position,  # pyright: ignore[reportArgumentType]
            color,
            thickness=thickness,
            tipLength=0.15,
        )

    return image