# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0.
'''
Amazon Kinesis Video Stream (KVS) Consumer Library for Python.
This class provides post-processing fiunctions for a MKV fragement that has been parsed
by the Amazon Kinesis Video Streams Cosumer Library for Python.
'''
__version__ = "0.0.1"
__status__ = "Development"
__copyright__ = "Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved."
__author__ = "Dean Colcott "
import io
import logging
import imageio.v3 as iio
import amazon_kinesis_video_consumer_library.ebmlite.util as emblite_utils
# Init the logger.
log = logging.getLogger(__name__)
class KvsFragementProcessor():
####################################################
# Fragment processing functions
def get_fragment_tags(self, fragment_dom):
'''
Parses a MKV Fragment Doc (of type ebmlite.core.MatroskaDocument) that is returned to the provided callback
from get_streaming_fragments() in this class and returns a dict of the SimpleTag elements found.
### Parameters:
**fragment_dom**: ebmlite.core.Document
The DOM like structure describing the fragment parsed by EBMLite.
### Returns:
simple_tags: dict
Dictionary of all SimpleTag elements with format - TagName : TagValue .
'''
# Get the Segment Element of the Fragment DOM - error if not found
segment_element = None
for element in fragment_dom:
if (element.id == 0x18538067): # MKV Segment Element ID
segment_element = element
break
if (not segment_element):
raise KeyError('Segment Element required but not found in fragment_doc' )
# Save all of the SimpleTag elements in the Segment element
simple_tag_elements = []
for element in segment_element:
if (element.id == 0x1254C367): # Tags element type ID
for tags in element:
if (tags.id == 0x7373): # Tag element type ID
for tag_type in tags:
if (tag_type.id == 0x67C8 ): # SimpleTag element type ID
simple_tag_elements.append(tag_type)
# For all SimpleTags types (ID: 0x67C8), save for TagName (ID: 0x7373) and values of TagString (ID:0x4487) or TagBinary (ID: 0x4485 )
simple_tags_dict = {}
for simple_tag in simple_tag_elements:
tag_name = None
tag_value = None
for element in simple_tag:
if (element.id == 0x45A3): # Tag Name element type ID
tag_name = element.value
elif (element.id == 0x4487 or element.id == 0x4485): # TagString and TagBinary element type IDs respectively
tag_value = element.value
# As long as tag name was found add the Tag to the return dict.
if (tag_name):
simple_tags_dict[tag_name] = tag_value
return simple_tags_dict
def get_fragement_dom_pretty_string(self, fragment_dom):
'''
Returns the Pretty Print parsing of the EBMLite fragment DOM as a string
### Parameters:
**fragment_dom**: ebmlite.core.Document
The DOM like structure describing the fragment parsed by EBMLite.
### Return:
**pretty_print_str**: str
Pretty print string of the Fragment DOM object
'''
pretty_print_str = io.StringIO()
emblite_utils.pprint(fragment_dom, out=pretty_print_str)
return pretty_print_str.getvalue()
def save_fragment_as_local_mkv(self, fragment_bytes, file_name_path):
'''
Save the provided fragment_bytes as stand-alone MKV file on local disk.
fragment_bytes as it arrives in is already a well formatted MKV fragment
so can just write the bytes straight to disk and it will be a playable MKV file.
### Parameters:
fragment_bytes: bytearray
A ByteArray with raw bytes from exactly one fragment.
file_name_path: Str
Local file path / name to save the MKV file to.
'''
f = open(file_name_path, "wb")
f.write(fragment_bytes)
f.close()
def get_frames_as_ndarray(self, fragment_bytes, one_in_frames_ratio):
'''
Parses fragment_bytes and returns a ratio of available frames in the MKV fragment as
a list of numpy.ndarray's.
e.g: Setting one_in_frames_ratio = 5 will return every 5th frame found in the fragment.
(Starting with the first)
To return all available frames just set one_in_frames_ratio = 1
### Parameters:
fragment_bytes: bytearray
A ByteArray with raw bytes from exactly one fragment.
one_in_frames_ratio: Str
Ratio of the available frames in the fragment to process and return.
### Return:
frames: List
A list of frames extracted from the fragment as numpy.ndarray
'''
# Parse all frames in the fragment to frames list
frames = iio.imread(io.BytesIO(fragment_bytes), plugin="pyav", index=...)
# Store and return frames in frame ratio of total available
ret_frames = []
for i in range(0, len(frames), one_in_frames_ratio):
ret_frames.append(frames[i])
return ret_frames
def save_frames_as_jpeg(self, fragment_bytes, one_in_frames_ratio, jpg_file_base_path):
'''
Parses fragment_bytes and saves a ratio of available frames in the MKV fragment as
JPEGs on the local disk.
e.g: Setting one_in_frames_ratio = 5 will return every 5th frame found in the fragment
(starting with the first).
To return all available frames just set one_in_frames_ratio = 1
### Parameters:
fragment_bytes: ByteArray
A ByteArray with raw bytes from exactly one fragment.
one_in_frames_ratio: Str
Ratio of the available frames in the fragment to process and save.
### Return
jpeg_paths : List
A list of file paths to the saved JPEN files.
'''
# Parse all frames in the fragment to frames list
ndarray_frames = self.get_frames_as_ndarray(fragment_bytes, one_in_frames_ratio)
# Write frames to disk as JPEG images
jpeg_paths = []
for i in range(len(ndarray_frames)):
frame = ndarray_frames[i]
image_file_path = '{}-{}.jpg'.format(jpg_file_base_path, i)
iio.imwrite(image_file_path, frame, format=None)
jpeg_paths.append(image_file_path)
return jpeg_paths