예제 #1
0
def annotate_video(file_path, coordinates):
    """
    Annotates supplied video from predicted coordinates.
    
    Args:
        file_path: path
            System path of video to annotate
        coordinates: list
            Predicted body part coordinates for each frame in the video
    """

    # Load raw video
    from skvideo.io import vreader, ffprobe, FFmpegWriter
    videogen = vreader(file_path)
    video_metadata = ffprobe(file_path)['video']
    fps = video_metadata['@r_frame_rate']
    frame_height, frame_width = next(vreader(file_path)).shape[:2]
    frame_side = frame_width if frame_width >= frame_height else frame_height

    # Initialize annotated video
    vcodec = 'libvpx-vp9'  #'libx264'
    writer = FFmpegWriter(normpath(file_path.split('.')[0] + '_tracked.mp4'),
                          inputdict={'-r': fps},
                          outputdict={
                              '-r': fps,
                              '-bitrate': '-1',
                              '-vcodec': vcodec,
                              '-pix_fmt': 'yuv420p',
                              '-lossless': '1'
                          })  #'-lossless': '1'

    # Annotate video
    from PIL import Image, ImageDraw
    i = 0
    while True:
        try:
            frame = next(videogen)
            image = Image.fromarray(frame)
            image_draw = ImageDraw.Draw(image)
            image_coordinates = coordinates[i]
            image = helpers.display_body_parts(image,
                                               image_draw,
                                               image_coordinates,
                                               image_height=frame_height,
                                               image_width=frame_width,
                                               marker_radius=int(frame_side /
                                                                 150))
            image = helpers.display_segments(image,
                                             image_draw,
                                             image_coordinates,
                                             image_height=frame_height,
                                             image_width=frame_width,
                                             segment_width=int(frame_side /
                                                               100))
            writer.writeFrame(np.array(image))
            i += 1
        except:
            break

    writer.close()
예제 #2
0
def annotate_image(file_path, coordinates):
    """
    Annotates supplied image from predicted coordinates.
    
    Args:
        file_path: path
            System path of image to annotate
        coordinates: list
            Predicted body part coordinates for image
    """

    # Load raw image
    from PIL import Image, ImageDraw
    image = Image.open(file_path)
    image_width, image_height = image.size
    image_side = image_width if image_width >= image_height else image_height

    # Annotate image
    image_draw = ImageDraw.Draw(image)
    image_coordinates = coordinates[0]
    image = helpers.display_body_parts(image,
                                       image_draw,
                                       image_coordinates,
                                       image_height=image_height,
                                       image_width=image_width,
                                       marker_radius=int(image_side / 150))
    image = helpers.display_segments(image,
                                     image_draw,
                                     image_coordinates,
                                     image_height=image_height,
                                     image_width=image_width,
                                     segment_width=int(image_side / 100))

    # Save annotated image
    image.save(normpath(file_path.split('.')[0] + '_tracked.png'))