def test_youtubevideo(spark, tmpdir):
df = spark.createDataFrame(
[
Row(YouTubeVideo("video_id")),
Row(YouTubeVideo("other_video_id")),
]
)
_check_roundtrip(spark, df, tmpdir)
def test_youtube():
from IPython.display import YouTubeVideo as IYT
vid = "pD1gDSao1eA"
yt = YouTubeVideo(vid)
assert yt.uri == "https://www.youtube.com/watch?v={0}".format(vid)
assert yt.embed_url == "http://www.youtube.com/embed/{0}".format(vid)
assert YouTubeVideo(vid) == YouTubeVideo(vid)
assert YouTubeVideo(vid) != YouTubeVideo("othervid")
def test_video_show():
vid = "pD1gDSao1eA"
v = YouTubeVideo(vid).get_stream()
result = v._repr_html_()
assert result == v.display()._repr_html_()
# TODO actually parse the html and check kwargs
from IPython.display import Video as IV
expected = IV(v.uri)._repr_html_()
assert result == expected
def test_youtube_show():
vid = "pD1gDSao1eA"
yt = YouTubeVideo(vid)
result = yt._repr_html_()
assert result == yt.display()._repr_html_()
# TODO actually parse the html and check kwargs
from IPython.display import YouTubeVideo as IYT
expected = IYT(vid)._repr_html_()
assert result == expected
def test_spectrogram_image(
spark: SparkSession, tmp_path: Path, asset_path: Path
):
"""Test generate spectrogram image
from YouTubeVideo/VideoStream videos types."""
video = VideoStream(str(asset_path / "big_buck_bunny_short.mp4"))
s1 = (
spark.createDataFrame([(video,)], ["video"])
.withColumn(
"spectrogram",
spectrogram_image(col("video"), lit(str(tmp_path / "s1.jpg"))),
)
.first()["spectrogram"]
)
assert type(s1) == Image
yt = YouTubeVideo(vid="rUWxSEwctFU")
s2 = (
spark.createDataFrame([(yt,)], ["video"])
.withColumn(
"spectrogram",
spectrogram_image(col("video"), lit(str(tmp_path / "s2.jpg"))),
)
.first()["spectrogram"]
)
assert type(s2) == Image
def test_spectrogram_image(spark: SparkSession):
"""Test generate spectrogram image
from YouTubeVideo/VideoStream videos types."""
videostream_df = spark.createDataFrame(
[
(VideoStream(uri=os.path.abspath(
os.path.join(
os.path.dirname(__file__),
"..",
"assets",
"big_buck_bunny_short.mp4",
))), ),
],
["video"],
)
youtube_df = spark.createDataFrame(
[
(YouTubeVideo(vid="rUWxSEwctFU"), ),
],
["video"],
)
videostream_df = videostream_df.withColumn(
"spectrogram",
spectrogram_image(col("video")),
)
youtube_df = youtube_df.withColumn(
"spectrogram",
spectrogram_image(col("video")),
)
videostream_sample = videostream_df.first()["spectrogram"]
youtube_sample = youtube_df.first()["spectrogram"]
assert type(videostream_sample) == Image
assert type(youtube_sample) == Image
def test_video_to_images(
spark: SparkSession, tmp_path: Path, asset_path: Path
):
"""Test extract video frames from YouTubeVideo/VideoStream types
into list of Image assets.
"""
sample_rate = 2
max_samples = 10
video = VideoStream(str(asset_path / "big_buck_bunny_short.mp4"))
df1 = spark.createDataFrame(
[(video, Segment(0, 20))], ["video", "segment"]
)
output_dir = tmp_path / "videostream_test"
output_dir.mkdir(parents=True)
df1 = df1.withColumn(
"images",
video_to_images(
col("video"),
lit(str(output_dir)),
col("segment"),
lit(sample_rate),
lit(max_samples),
),
)
df2 = spark.createDataFrame(
[(YouTubeVideo(vid="rUWxSEwctFU"), Segment(0, 20))],
["video", "segment"],
)
output_dir = tmp_path / "youtube_test"
output_dir.mkdir(parents=True)
df2 = df2.withColumn(
"images",
video_to_images(
col("video"),
lit(str(output_dir)),
col("segment"),
lit(sample_rate),
lit(max_samples),
),
)
videostream_sample = df1.first()["images"]
youtube_sample = df2.first()["images"]
assert (
type(videostream_sample) == list
and type(videostream_sample[0]) == Image
and len(videostream_sample) == max_samples
)
assert (
type(youtube_sample) == list
and type(youtube_sample[0]) == Image
and len(youtube_sample) == max_samples
)
def test_video_to_images(spark: SparkSession):
"""Test extract video frames from YouTubeVideo/VideoStream types
into list of Image assets.
"""
sample_rate = 2
max_samples = 10
videostream_df = spark.createDataFrame(
[
(
VideoStream(uri=os.path.abspath(
os.path.join(
os.path.dirname(__file__),
"..",
"assets",
"big_buck_bunny_short.mp4",
))),
Segment(0, 20),
),
],
["video", "segment"],
)
youtube_df = spark.createDataFrame(
[
(YouTubeVideo(vid="rUWxSEwctFU"), Segment(0, 20)),
],
["video", "segment"],
)
videostream_df = videostream_df.withColumn(
"images",
video_to_images(col("video"), col("segment"), lit(sample_rate),
lit(max_samples)),
)
youtube_df = youtube_df.withColumn(
"images",
video_to_images(col("video"), col("segment"), lit(sample_rate),
lit(max_samples)),
)
videostream_sample = videostream_df.first()["images"]
youtube_sample = youtube_df.first()["images"]
assert (type(videostream_sample) == list
and type(videostream_sample[0]) == Image
and len(videostream_sample) == max_samples)
assert (type(youtube_sample) == list and type(youtube_sample[0]) == Image
and len(youtube_sample) == max_samples)
def test_youtube_sample_stream():
vid = "pD1gDSao1eA"
yt = YouTubeVideo(vid)
v = yt.get_stream()
assert isinstance(v, VideoStream)
isinstance(next(v.__iter__()), np.ndarray)
def deserialize(self, datum) -> "YouTubeVideo":
from rikai.types import YouTubeVideo # pylint: disable=import-outside-toplevel
return YouTubeVideo(datum[0])
def test_youtubevideo(self):
df = self.spark.createDataFrame(
[Row(YouTubeVideo("video_id")), Row(YouTubeVideo("other_video_id"))]
)
self._check_roundtrip(df)