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()
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'))
