Update Python type annotations to modern syntax and types

Author: scotts

examples/decoding/approximate_mode.py

@@ -42,21 +42,25 @@
 
 temp_dir = tempfile.mkdtemp()
 short_video_path = Path(temp_dir) / "short_video.mp4"
-with open(short_video_path, 'wb') as f:
+with open(short_video_path, "wb") as f:
     for chunk in response.iter_content():
         f.write(chunk)
 
 long_video_path = Path(temp_dir) / "long_video.mp4"
 ffmpeg_command = [
     "ffmpeg",
-    "-stream_loop", "99",  # repeat video 100 times
-    "-i", f"{short_video_path}",
-    "-c", "copy",
-    f"{long_video_path}"
+    "-stream_loop",
+    "99",  # repeat video 100 times
+    "-i",
+    f"{short_video_path}",
+    "-c",
+    "copy",
+    f"{long_video_path}",
 ]
 subprocess.run(ffmpeg_command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 
 from torchcodec.decoders import VideoDecoder
+
 print(f"Short video duration: {VideoDecoder(short_video_path).metadata.duration_seconds} seconds")
 print(f"Long video duration: {VideoDecoder(long_video_path).metadata.duration_seconds / 60} minutes")
 
@@ -116,10 +120,7 @@ def bench(f, average_over=50, warmup=2, **f_kwargs):
 
 def sample_clips(seek_mode):
     return samplers.clips_at_random_indices(
-        decoder=VideoDecoder(
-            source=long_video_path,
-            seek_mode=seek_mode
-        ),
+        decoder=VideoDecoder(source=long_video_path, seek_mode=seek_mode),
         num_clips=5,
         num_frames_per_clip=2,
     )
@@ -154,7 +155,8 @@ def sample_clips(seek_mode):
     torch.testing.assert_close(
         exact_decoder.get_frame_at(i).data,
         approx_decoder.get_frame_at(i).data,
-        atol=0, rtol=0,
+        atol=0,
+        rtol=0,
     )
 print("Frame seeking is the same for this video!")
 

examples/decoding/basic_cuda_example.py

@@ -154,9 +154,7 @@ def plot_cpu_and_cuda_frames(cpu_frames: torch.Tensor, cuda_frames: torch.Tensor
 # differences because CUDA math is not bit-exact with respect to CPU math.
 #
 frames_equal = torch.equal(cpu_frames.to("cuda"), cuda_frames)
-mean_abs_diff = torch.mean(
-    torch.abs(cpu_frames.float().to("cuda") - cuda_frames.float())
-)
+mean_abs_diff = torch.mean(torch.abs(cpu_frames.float().to("cuda") - cuda_frames.float()))
 max_abs_diff = torch.max(torch.abs(cpu_frames.to("cuda").float() - cuda_frames.float()))
 print(f"{frames_equal=}")
 print(f"{mean_abs_diff=}")

examples/decoding/basic_example.py

@@ -13,12 +13,13 @@
 :class:`~torchcodec.decoders.VideoDecoder` class.
 """
 
+from __future__ import annotations
+
 # %%
 # First, a bit of boilerplate: we'll download a video from the web, and define a
 # plotting utility. You can ignore that part and jump right below to
 # :ref:`creating_decoder`.
 
-from typing import Optional
 import torch
 import requests
 
@@ -33,7 +34,7 @@
 raw_video_bytes = response.content
 
 
-def plot(frames: torch.Tensor, title : Optional[str] = None):
+def plot(frames: torch.Tensor, title: str | None = None):
     try:
         from torchvision.utils import make_grid
         from torchvision.transforms.v2.functional import to_pil_image
@@ -42,7 +43,7 @@ def plot(frames: torch.Tensor, title : Optional[str] = None):
         print("Cannot plot, please run `pip install torchvision matplotlib`")
         return
 
-    plt.rcParams["savefig.bbox"] = 'tight'
+    plt.rcParams["savefig.bbox"] = "tight"
     fig, ax = plt.subplots()
     ax.imshow(to_pil_image(make_grid(frames)))
     ax.set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])
@@ -76,7 +77,7 @@ def plot(frames: torch.Tensor, title : Optional[str] = None):
 # ---------------------------------------
 
 first_frame = decoder[0]  # using a single int index
-every_twenty_frame = decoder[0 : -1 : 20]  # using slices
+every_twenty_frame = decoder[0:-1:20]  # using slices
 
 print(f"{first_frame.shape = }")
 print(f"{first_frame.dtype = }")
@@ -116,10 +117,7 @@ def plot(frames: torch.Tensor, title : Optional[str] = None):
 # The decoder is a normal iterable object and can be iterated over like so:
 
 for frame in decoder:
-    assert (
-        isinstance(frame, torch.Tensor)
-        and frame.shape == (3, decoder.metadata.height, decoder.metadata.width)
-    )
+    assert isinstance(frame, torch.Tensor) and frame.shape == (3, decoder.metadata.height, decoder.metadata.width)
 
 # %%
 # Retrieving pts and duration of frames

examples/decoding/custom_frame_mappings.py

@@ -41,21 +41,25 @@
 
 temp_dir = tempfile.mkdtemp()
 short_video_path = Path(temp_dir) / "short_video.mp4"
-with open(short_video_path, 'wb') as f:
+with open(short_video_path, "wb") as f:
     for chunk in response.iter_content():
         f.write(chunk)
 
 long_video_path = Path(temp_dir) / "long_video.mp4"
 ffmpeg_command = [
     "ffmpeg",
-    "-stream_loop", "50",  # repeat video 50 times to get a ~12 min video
-    "-i", f"{short_video_path}",
-    "-c", "copy",
-    f"{long_video_path}"
+    "-stream_loop",
+    "50",  # repeat video 50 times to get a ~12 min video
+    "-i",
+    f"{short_video_path}",
+    "-c",
+    "copy",
+    f"{long_video_path}",
 ]
 subprocess.run(ffmpeg_command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 
 from torchcodec.decoders import VideoDecoder
+
 print(f"Short video duration: {VideoDecoder(short_video_path).metadata.duration_seconds} seconds")
 print(f"Long video duration: {VideoDecoder(long_video_path).metadata.duration_seconds / 60} minutes")
 
@@ -82,7 +86,18 @@
 
 # Lets define a simple function to run ffprobe on a video's first stream index, then writes the results in output_json_path.
 def generate_frame_mappings(video_path, output_json_path, stream_index):
-    ffprobe_cmd = ["ffprobe", "-i", f"{video_path}", "-select_streams", f"{stream_index}", "-show_frames", "-show_entries", "frame=pts,duration,key_frame", "-of", "json"]
+    ffprobe_cmd = [
+        "ffprobe",
+        "-i",
+        f"{video_path}",
+        "-select_streams",
+        f"{stream_index}",
+        "-show_frames",
+        "-show_entries",
+        "frame=pts,duration,key_frame",
+        "-of",
+        "json",
+    ]
     print(f"Running ffprobe:\n{' '.join(ffprobe_cmd)}\n")
     ffprobe_result = subprocess.run(ffprobe_cmd, check=True, capture_output=True, text=True)
     with open(output_json_path, "w") as f:
@@ -157,12 +172,8 @@ def bench(f, file_like=False, average_over=50, warmup=2, **f_kwargs):
 # so the performance benefits are realized.
 
 
-def decode_frames(video_path, seek_mode = "exact", custom_frame_mappings = None):
-    decoder = VideoDecoder(
-        source=video_path,
-        seek_mode=seek_mode,
-        custom_frame_mappings=custom_frame_mappings
-    )
+def decode_frames(video_path, seek_mode="exact", custom_frame_mappings=None):
+    decoder = VideoDecoder(source=video_path, seek_mode=seek_mode, custom_frame_mappings=custom_frame_mappings)
     decoder.get_frames_in_range(start=0, stop=10)
 
 
@@ -196,7 +207,8 @@ def decode_frames(video_path, seek_mode = "exact", custom_frame_mappings = None)
     torch.testing.assert_close(
         exact_decoder.get_frame_at(i).data,
         custom_frame_mappings_decoder.get_frame_at(i).data,
-        atol=0, rtol=0,
+        atol=0,
+        rtol=0,
     )
 print("Frame seeking is the same for this video!")
 

examples/decoding/file_like.py

@@ -158,7 +158,7 @@ def stream_while_decode():
     # session; we need to indicate that we need to trust the environment
     # settings for proxy configuration. Depending on your environment, you may
     # not need this setting.
-    with fsspec.open(nasa_url, client_kwargs={'trust_env': True}) as file_like:
+    with fsspec.open(nasa_url, client_kwargs={"trust_env": True}) as file_like:
         decoder = VideoDecoder(file_like, seek_mode="approximate")
         return decoder[0]
 
@@ -226,18 +226,20 @@ def seek(self, offset: int, whence: int) -> int:
 file_op_counter = FileOpCounter(open(nasa_video_path, "rb"))
 counter_decoder = VideoDecoder(file_op_counter, seek_mode="approximate")
 
-print("Decoder initialization required "
-      f"{file_op_counter.num_reads} reads and "
-      f"{file_op_counter.num_seeks} seeks.")
+print(
+    "Decoder initialization required " f"{file_op_counter.num_reads} reads and " f"{file_op_counter.num_seeks} seeks."
+)
 
 init_reads = file_op_counter.num_reads
 init_seeks = file_op_counter.num_seeks
 
 first_frame = counter_decoder[0]
 
-print("Decoding the first frame required "
-      f"{file_op_counter.num_reads - init_reads} additional reads and "
-      f"{file_op_counter.num_seeks - init_seeks} additional seeks.")
+print(
+    "Decoding the first frame required "
+    f"{file_op_counter.num_reads - init_reads} additional reads and "
+    f"{file_op_counter.num_seeks - init_seeks} additional seeks."
+)
 
 # %%
 # While we defined a simple class primarily for demonstration, it's actually
@@ -299,5 +301,6 @@ def decode_from_existing_open_file_object():
 # %%
 # Finally, let's clean up the local resources we created.
 import shutil
+
 shutil.rmtree(temp_dir)
 # %%

examples/decoding/parallel_decoding.py

@@ -24,14 +24,15 @@
 pool manager.
 """
 
+from __future__ import annotations
+
 # %%
 # Let's first define some utility functions for benchmarking and data
 # processing.  We'll also download a video and create a longer version by
 # repeating it multiple times. This simulates working with long videos that
 # require efficient processing. You can ignore that part and jump right below to
 # :ref:`start_parallel_decoding`.
 
-from typing import List
 import torch
 import requests
 import tempfile
@@ -74,16 +75,16 @@ def report_stats(times, unit="s"):
     return med
 
 
-def split_indices(indices: List[int], num_chunks: int) -> List[List[int]]:
+def split_indices(indices: list[int], num_chunks: int) -> list[list[int]]:
     """Split a list of indices into approximately equal chunks."""
     chunk_size = len(indices) // num_chunks
     chunks = []
 
     for i in range(num_chunks - 1):
-        chunks.append(indices[i * chunk_size:(i + 1) * chunk_size])
+        chunks.append(indices[i * chunk_size : (i + 1) * chunk_size])
 
     # Last chunk may be slightly larger
-    chunks.append(indices[(num_chunks - 1) * chunk_size:])
+    chunks.append(indices[(num_chunks - 1) * chunk_size :])
     return chunks
 
 
@@ -96,18 +97,22 @@ def generate_long_video(temp_dir: str):
         raise RuntimeError(f"Failed to download video. {response.status_code = }.")
 
     short_video_path = Path(temp_dir) / "short_video.mp4"
-    with open(short_video_path, 'wb') as f:
+    with open(short_video_path, "wb") as f:
         for chunk in response.iter_content():
             f.write(chunk)
 
     # Create a longer video by repeating the short one 50 times
     long_video_path = Path(temp_dir) / "long_video.mp4"
     ffmpeg_command = [
-        "ffmpeg", "-y",
-        "-stream_loop", "49",  # repeat video 50 times
-        "-i", str(short_video_path),
-        "-c", "copy",
-        str(long_video_path)
+        "ffmpeg",
+        "-y",
+        "-stream_loop",
+        "49",  # repeat video 50 times
+        "-i",
+        str(short_video_path),
+        "-c",
+        "copy",
+        str(long_video_path),
     ]
     subprocess.run(ffmpeg_command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 
@@ -122,7 +127,9 @@ def generate_long_video(temp_dir: str):
 
 short_duration = timedelta(seconds=VideoDecoder(short_video_path).metadata.duration_seconds)
 long_duration = timedelta(seconds=metadata.duration_seconds)
-print(f"Original video duration: {int(short_duration.total_seconds() // 60)}m{int(short_duration.total_seconds() % 60):02d}s")
+print(
+    f"Original video duration: {int(short_duration.total_seconds() // 60)}m{int(short_duration.total_seconds() % 60):02d}s"
+)
 print(f"Long video duration: {int(long_duration.total_seconds() // 60)}m{int(long_duration.total_seconds() % 60):02d}s")
 print(f"Video resolution: {metadata.width}x{metadata.height}")
 print(f"Average FPS: {metadata.average_fps:.1f}")
@@ -155,7 +162,8 @@ def generate_long_video(temp_dir: str):
 # Let's start with a sequential approach as our baseline. This processes
 # frames one by one without any parallelization.
 
-def decode_sequentially(indices: List[int], video_path=long_video_path):
+
+def decode_sequentially(indices: list[int], video_path=long_video_path):
     """Decode frames sequentially using a single decoder instance."""
     decoder = VideoDecoder(video_path, seek_mode="approximate")
     return decoder.get_frames_at(indices)
@@ -173,11 +181,8 @@ def decode_sequentially(indices: List[int], video_path=long_video_path):
 # via the ``num_ffmpeg_threads`` parameter. This approach uses multiple
 # threads within FFmpeg itself to accelerate decoding operations.
 
-def decode_with_ffmpeg_parallelism(
-    indices: List[int],
-    num_threads: int,
-    video_path=long_video_path
-):
+
+def decode_with_ffmpeg_parallelism(indices: list[int], num_threads: int, video_path=long_video_path):
     """Decode frames using FFmpeg's internal threading."""
     decoder = VideoDecoder(video_path, num_ffmpeg_threads=num_threads, seek_mode="approximate")
     return decoder.get_frames_at(indices)
@@ -197,11 +202,8 @@ def decode_with_ffmpeg_parallelism(
 #
 # Process-based parallelism distributes work across multiple Python processes.
 
-def decode_with_multiprocessing(
-    indices: List[int],
-    num_processes: int,
-    video_path=long_video_path
-):
+
+def decode_with_multiprocessing(indices: list[int], num_processes: int, video_path=long_video_path):
     """Decode frames using multiple processes with joblib."""
     chunks = split_indices(indices, num_chunks=num_processes)
 
@@ -226,11 +228,8 @@ def decode_with_multiprocessing(
 # Thread-based parallelism uses multiple threads within a single process.
 # TorchCodec releases the GIL, so this can be very effective.
 
-def decode_with_multithreading(
-    indices: List[int],
-    num_threads: int,
-    video_path=long_video_path
-):
+
+def decode_with_multithreading(indices: list[int], num_threads: int, video_path=long_video_path):
     """Decode frames using multiple threads with joblib."""
     chunks = split_indices(indices, num_chunks=num_threads)
 
@@ -261,4 +260,5 @@ def decode_with_multithreading(
 
 # %%
 import shutil
+
 shutil.rmtree(temp_dir)