visualize.py

#!/usr/bin/env python3
 

 
# Author: Ryan Luna
 
import sys
import yaml
from math import sin, cos
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import matplotlib.patches as patches
 
 
# Reads a YAML world specification from the worldFile.
# Returns the set of obstacles and the bounds of the world.
def readWorldFile(worldFile):
    handle = open(worldFile, "r")
    # use safe_load instead load
    dataMap = yaml.safe_load(handle)
    handle.close()
 
    # Read in bounds
    xlower = float(dataMap["bounds"]["x"]["lower"])
    xupper = float(dataMap["bounds"]["x"]["upper"])
    ylower = float(dataMap["bounds"]["y"]["lower"])
    yupper = float(dataMap["bounds"]["y"]["upper"])
 
    bounds = [[xlower, xupper], [ylower, yupper]]
 
    obs = {}
 
    if "obstacles" in dataMap:
        for r in dataMap["obstacles"]:
            for k in r.keys():
                if k not in obs:
                    obs[k] = []
                obs[k].append(r[k])
    return obs, bounds
 
 
def readPathFile(pathFile):
    # Read in file line by line
    lines = [line.rstrip() for line in open(pathFile) if len(line.rstrip()) > 0]
 
    # first line is meta data
    metadata = lines[0]
    metadata = metadata.split(" ")
    if len(metadata) != 5:
        raise RuntimeError(
            "Malformed path file.  Expected first line with # links, link length, originX, originY, xySlices"
        )
 
    numLinks = int(metadata[0])
    linkLength = float(metadata[1])
    origin = (float(metadata[2]), float(metadata[3]))
    slices = int(metadata[4])
 
    path = []
 
    for l in lines[1:]:
        entries = l.split(" ")
        if len(entries) != numLinks:
            raise RuntimeError(
                "Malformed path file.  Path entries must have length = # links"
            )
        config = [float(e) for e in entries]
        path.append(config)
 
    return numLinks, linkLength, origin, path, slices
 
 
def plotPolygon(axes, data, color, alpha=1.0, edgecolor=None):
    points = []
    for d in data:
        for k in d.keys():
            if k == "vertex":
                coords = d[k].strip()
                if coords.startswith("("):
                    coords = coords[1:-1]
                coords = coords.strip().split(",")
                pt = [float(coords[0]), float(coords[1])]
                points.append(pt)
 
            else:
                raise RuntimeError('Expected "vertex", but got ', k)
 
    if len(points) > 0:
        arr = np.array(points)
        axes.add_patch(
            patches.Polygon(
                arr, facecolor=color, alpha=alpha, fill=True, edgecolor=edgecolor
            )
        )
 
 
def plotObstacles(axes, obstacles, bounds):
    for k, v in obstacles.items():
        for o in v:
            if k == "polygon":
                plotPolygon(axes, o, "0.4", edgecolor="0.4")
            else:
                raise RuntimeError("Unknown geometry type: ", k)
 
    plt.axis([bounds[0][0], bounds[0][1], bounds[1][0], bounds[1][1]])
 
 
def plotChain(axes, angles, origin, color="blue"):
    x = origin[0]
    y = origin[1]
    angle = 0.0
 
    linkLength = 1.0 / len(angles)
    for theta in angles:
        angle += theta
 
        xN = x + (cos(angle) * linkLength)
        yN = y + (sin(angle) * linkLength)
 
        axes.plot([x, xN], [y, yN], "-", color=color)
 
        x = xN
        y = yN
 
 
def plotChainFrame(frame_num, ax, path, obstacles, bounds, origin, rate, slices):
    ax.clear()
    plotObstacles(ax, obstacles, bounds)
 
    # Hold for 1 second at the beginning and end
    if frame_num < rate:
        configuration = path[0]
    elif frame_num >= len(path) + rate:
        configuration = path[-1]
    else:
        configuration = path[frame_num - rate]
 
    plotChain(ax, configuration, origin)
 
 
def AnimatePath(obstacles, bounds, numLinks, linkLength, origin, path, slices):
    fig = plt.figure()
    axes = plt.axes()
    framerate = 30
 
    numFrames = len(path)
    anim = animation.FuncAnimation(
        fig,
        plotChainFrame,
        fargs=[axes, path, obstacles, bounds, origin, slices, framerate],
        frames=numFrames + (framerate * 2),
        interval=(1.0 / framerate) * 1000,
        blit=False,
    )
 
    filename = "animation.mp4"
    anim.save(filename, fps=framerate)
 
 
if __name__ == "__main__":
    worldFile = "./world.yaml"
    pathFile = "./manipulator_path.txt"
 
    obs, bounds = readWorldFile(worldFile)
    numLinks, linkLength, origin, path, slices = readPathFile(pathFile)
 
    print("Creating animation for planar kinematic chain...  ")
    AnimatePath(obs, bounds, numLinks, linkLength, origin, path, slices)
    print("Done")