def numpy_to_image(array: ndarray, uri: str) -> Image:
"""Convert a numpy array to image, and upload to external storage.
Parameters
----------
array : :py:class:`numpy.ndarray`
Image data.
uri : str
The base directory to copy the image to.
Return
------
Image
Return a new image pointed to the new URI.
Example
-------
>>> spark.createDataFrame(..).registerTempTable("df")
>>>
>>> spark.sql(\"\"\"SELECT numpy_to_image(
... resize(grayscale(image)),
... lit('s3://asset')
... ) AS new_image FROM df\"\"\")
See Also
--------
:py:meth:`rikai.types.vision.Image.from_array`
"""
return Image.from_array(array, uri)
def video_to_images(
video: Union[VideoStream, YouTubeVideo],
output_uri: str,
segment: Segment = Segment(0, -1),
sample_rate: int = 1,
max_samples: int = 15000,
quality: str = "worst",
) -> list:
"""Extract video frames into a list of images.
Parameters
----------
video : Video
An video object, either YouTubeVideo or VideoStream.
output_uri: str
Frames will be written as <output_uri>/<fno>.jpg
segment: Segment, default Segment(0, -1)
A Segment object, localizing video in time to (start_fno, end_fno)
sample_rate : int, default 1
Keep 1 out of every sample_rate frames.
max_samples : int, default 15000
Return at most this many frames (-1 means no max)
quality: str, default 'worst'
Either 'worst' (lowest bitrate) or 'best' (highest bitrate)
See: https://pythonhosted.org/Pafy/index.html#Pafy.Pafy.getbest
Return
------
List
Return a list of images from video indexed by frame number.
"""
assert isinstance(
video, (YouTubeVideo,
VideoStream)), "Input type must be YouTubeVideo or VideoStream"
assert isinstance(segment, Segment), "Second input type must be Segment"
start_frame = segment.start_fno
if segment.end_fno > 0:
max_samples = min((segment.end_fno - start_frame), max_samples)
if isinstance(video, YouTubeVideo):
video_iterator = SingleFrameSampler(
video.get_stream(quality=quality),
sample_rate,
start_frame,
max_samples,
)
else:
video_iterator = SingleFrameSampler(video, sample_rate, start_frame,
max_samples)
return [
Image.from_array(
img,
os.path.join(output_uri, "{}.jpg".format(
(start_frame + idx) * sample_rate)),
) for idx, img in enumerate(video_iterator)
]
def video_to_images(
video,
segment: Segment = Segment(0, -1),
sample_rate: int = 1,
max_samples: int = 15000,
quality: str = "worst",
) -> list:
"""Extract video frames into a list of images.
Parameters
----------
video : Video
An video object, either YouTubeVideo or VideoStream.
segment: Segment
A Segment object, localizing video in time to (start_fno, end_fno)
sample_rate : Int
The sampling rate in number of frames
max_samples : Int
Yield at most this many frames (-1 means no max)
quality: str, default 'worst'
Either 'worst' (lowest bitrate) or 'best' (highest bitrate)
See: https://pythonhosted.org/Pafy/index.html#Pafy.Pafy.getbest
Return
------
List
Return a list of images from video indexed by frame number.
"""
assert isinstance(
video, (YouTubeVideo,
VideoStream)), "Input type must be YouTubeVideo or VideoStream"
assert isinstance(segment, Segment), "Second input type must be Segment"
base_path = video.uri
start_frame = segment.start_fno
if segment.end_fno > 0:
max_samples = min((segment.end_fno - start_frame), max_samples)
if isinstance(video, YouTubeVideo):
base_path = video.vid
video_iterator = SingleFrameSampler(
video.get_stream(quality=quality),
sample_rate,
start_frame,
max_samples,
)
else:
video_iterator = SingleFrameSampler(video, sample_rate, start_frame,
max_samples)
return [
Image.from_array(
img,
"{}_{}.jpg".format(base_path, (start_frame + idx) * sample_rate),
) for idx, img in enumerate(video_iterator)
]
def spectrogram_image(
video: Union[VideoStream, YouTubeVideo],
output_uri: str,
segment: Segment = Segment(0, -1),
size: int = 224,
max_samples: int = 15000,
) -> Image:
"""Applies ffmpeg filter to generate spectrogram image.
Parameters
----------
video : VideoStream or YouTubeVideo
A video object whose audio track will be converted to spectrogram
output_uri: str
The uri to which the spectrogram image will be written to
segment: Segment
A Segment object, localizing video in time to (start_fno, end_fno)
max_samples : Int
Yield at most this many frames (-1 means no max)
size : Int
Sets resolution of frequency, time spectrogram image.
Return
------
Image
Return an Image of the audio spectrogram.
"""
try:
import ffmpeg
except ImportError:
raise ValueError("Couldn't import ffmpeg. Please make sure to "
"`pip install ffmpeg-python` explicitly or install "
"the correct extras like `pip install rikai[all]`")
assert isinstance(
video, (YouTubeVideo,
VideoStream)), "Input type must be YouTubeVideo or VideoStream"
assert isinstance(segment, Segment), "Second input type must be Segment"
start_frame = segment.start_fno
if segment.end_fno > 0:
max_samples = min((segment.end_fno - start_frame), max_samples)
video_uri = (video.get_stream().uri
if isinstance(video, YouTubeVideo) else video.uri)
output, _ = (ffmpeg.input(video_uri).filter("showspectrumpic",
"{}x{}".format(size, size),
legend=0).output(
"pipe:",
format="rawvideo",
pix_fmt="rgb24",
start_number=start_frame,
vframes=max_samples,
).run(capture_stdout=True))
return Image.from_array(
np.frombuffer(output, np.uint8).reshape([size, size, 3]), output_uri)
def test_format_kwargs(tmp_path):
data = np.random.random((100, 100))
rescaled = (255.0 / data.max() * (data - data.min())).astype(np.uint8)
result_uri = tmp_path / "result.jpg"
Image.from_array(rescaled, result_uri, format="jpeg", optimize=True)
expected_uri = tmp_path / "expected.jpg"
PILImage.fromarray(rescaled).save(expected_uri,
format="jpeg",
optimize=True)
assert filecmp.cmp(result_uri, expected_uri)
result_uri = tmp_path / "result.png"
Image.from_array(rescaled, result_uri, format="png", compress_level=1)
expected_uri = tmp_path / "expected.png"
PILImage.fromarray(rescaled).save(expected_uri,
format="png",
compress_level=1)
assert filecmp.cmp(result_uri, expected_uri)
def spectrogram_image(
video,
segment: Segment = Segment(0, -1),
size: int = 224,
max_samples: int = 15000,
) -> Image:
"""Applies ffmpeg filter to generate spectrogram image.
Parameters
----------
video : Video
A video object, either YouTubeVideo or VideoStream.
segment: Segment
A Segment object, localizing video in time to (start_fno, end_fno)
max_samples : Int
Yield at most this many frames (-1 means no max)
size : Int
Sets resolution of frequency, time spectrogram image.
Return
------
Image
Return an Image of the audio spectrogram.
"""
import ffmpeg
assert isinstance(
video, (YouTubeVideo,
VideoStream)), "Input type must be YouTubeVideo or VideoStream"
assert isinstance(segment, Segment), "Second input type must be Segment"
base_path = video.vid if isinstance(video, YouTubeVideo) else video.uri
start_frame = segment.start_fno
if segment.end_fno > 0:
max_samples = min((segment.end_fno - start_frame), max_samples)
video_uri = (video.get_stream().uri
if isinstance(video, YouTubeVideo) else video.uri)
output, _ = (ffmpeg.input(video_uri).filter("showspectrumpic",
"{}x{}".format(size, size),
legend=0).output(
"pipe:",
format="rawvideo",
pix_fmt="rgb24",
start_number=start_frame,
vframes=max_samples,
).run(capture_stdout=True))
return Image.from_array(
np.frombuffer(output, np.uint8).reshape([size, size, 3]),
"{}_spectrogram.jpg".format(base_path),
)
def numpy_to_image(array: ndarray,
uri: str,
format: str = None,
**kwargs) -> Image:
"""Convert a numpy array to image, and upload to external storage.
Parameters
----------
array : :py:class:`numpy.ndarray`
Image data.
uri : str
The base directory to copy the image to.
format : str, optional
The image format to save as. See
`supported formats <https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.save>`_ for details.
kwargs : dict, optional
Optional arguments to pass to `PIL.Image.save <https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.save>`_.
Return
------
Image
Return a new image pointed to the new URI.
Example
-------
>>> spark.createDataFrame(..).registerTempTable("df")
>>>
>>> spark.sql(\"\"\"SELECT numpy_to_image(
... resize(grayscale(image)),
... lit('s3://asset')
... ) AS new_image FROM df\"\"\")
See Also
--------
:py:meth:`rikai.types.vision.Image.from_array`
""" # noqa: E501
return Image.from_array(array, uri, format=format, **kwargs)
def video_to_images(
video: Union[VideoStream, YouTubeVideo],
output_uri: str,
segment: Segment = Segment(0, -1),
sample_rate: int = 1,
max_samples: int = 15000,
quality: str = "worst",
image_format: str = "png",
**image_kwargs,
) -> list:
"""Extract video frames into a list of images.
Parameters
----------
video : Video
An video object, either YouTubeVideo or VideoStream.
output_uri: str
Frames will be written as <output_uri>/<fno>.<img_format>
segment: Segment, default Segment(0, -1)
A Segment object, localizing video in time to (start_fno, end_fno)
sample_rate : int, default 1
Keep 1 out of every sample_rate frames.
max_samples : int, default 15000
Return at most this many frames (-1 means no max)
quality: str, default 'worst'
Either 'worst' (lowest bitrate) or 'best' (highest bitrate)
See: https://pythonhosted.org/Pafy/index.html#Pafy.Pafy.getbest
image_format : str, optional
The image format to save as. See
`supported formats <https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.save>`_ for details.
image_kwargs : dict, optional
Optional arguments to pass to `PIL.Image.save <https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.save>`_.
------
List
Return a list of images from video indexed by frame number.
""" # noqa: E501
assert isinstance(
video, (YouTubeVideo,
VideoStream)), "Input type must be YouTubeVideo or VideoStream"
assert isinstance(segment, Segment), "Second input type must be Segment"
start_frame = segment.start_fno
if segment.end_fno > 0:
max_samples = min((segment.end_fno - start_frame), max_samples)
if isinstance(video, YouTubeVideo):
video_iterator = SingleFrameSampler(
video.get_stream(quality=quality),
sample_rate,
start_frame,
max_samples,
)
else:
video_iterator = SingleFrameSampler(video, sample_rate, start_frame,
max_samples)
return [
Image.from_array(
img,
os.path.join(
output_uri,
"{}.{}".format((start_frame + idx) * sample_rate,
image_format),
),
format=image_format,
**image_kwargs,
) for idx, img in enumerate(video_iterator)
]
def test_crop_image():
data = np.random.randint(0, 255, size=(100, 100), dtype=np.uint8)
im = Image.from_array(data)
patch = im.crop(Box2d(10, 10, 30, 30))
cropped_data = patch.to_numpy()
assert np.array_equal(cropped_data, data[10:30, 10:30])
