Delete convert_image_to_tensor script

Author: Daniel Flores

test/convert_image_to_tensor.py

@@ -1,7 +1,11 @@
 #!/usr/bin/env python3
 import os
 import subprocess
-import sys
+
+import numpy as np
+
+import torch
+from PIL import Image
 
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
@@ -13,9 +17,23 @@
 # from source media already checked into the repo.
 
 
+def convert_image_to_tensor(image_path):
+    if not os.path.exists(image_path):
+        return
+    # Get base filename without extension
+    base_filename = os.path.splitext(image_path)[0]
+    pil_image = Image.open(image_path)
+    img_tensor = torch.from_numpy(np.asarray(pil_image))
+    print(img_tensor.shape)
+    print(img_tensor.dtype)
+    # Save tensor to disk
+    torch.save(img_tensor, base_filename + ".pt", _use_new_zipfile_serialization=True)
+    os.remove(image_path)
+
+
 def main():
     SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
-    TORCHCODEC_PATH = os.path.join(SCRIPT_DIR, "..")
+    TORCHCODEC_PATH = os.path.dirname(SCRIPT_DIR)
     RESOURCES_DIR = os.path.join(TORCHCODEC_PATH, "test", "resources")
     VIDEO_PATH = os.path.join(RESOURCES_DIR, "nasa_13013.mp4")
 
@@ -31,6 +49,7 @@ def main():
             # frames would result in 1-based index naming. We enforce 0-based index naming
             # so that the name of reference frames matches the index when accessing that
             # frame in the Python decoder.
+            output_bmp = f"{VIDEO_PATH}.stream{stream}.frame{frame:06d}.bmp"
             frame_name = f"{frame:06d}"
             cmd = [
                 "ffmpeg",
@@ -45,14 +64,16 @@ def main():
                 "vfr",
                 "-q:v",
                 "2",
-                f"{VIDEO_PATH}.stream{stream}.frame{frame_name}.bmp",
+                output_bmp,
             ]
             subprocess.run(cmd, check=True)
+            convert_image_to_tensor(output_bmp)
 
     # Extract individual frames at specific timestamps, including the last frame of the video.
     seek_timestamp = [6.0, 6.1, 10.0, 12.979633]
     timestamp_name = [f"{seek_timestamp:06f}" for seek_timestamp in seek_timestamp]
     for timestamp, name in zip(seek_timestamp, timestamp_name):
+        output_bmp = f"{VIDEO_PATH}.time{name}.bmp"
         cmd = [
             "ffmpeg",
             "-y",
@@ -65,6 +86,7 @@ def main():
             f"{VIDEO_PATH}.time{name}.bmp",
         ]
         subprocess.run(cmd, check=True)
+        convert_image_to_tensor(output_bmp)
 
     # This video was generated by running the following:
     # conda install -c conda-forge x265
@@ -75,6 +97,7 @@ def main():
     FRAMES = [5]
     for frame in FRAMES:
         frame_name = f"{frame:06d}"
+        output_bmp = f"{VIDEO_PATH}.stream0.frame{frame_name}.bmp"
         cmd = [
             "ffmpeg",
             "-y",
@@ -86,9 +109,10 @@ def main():
             "vfr",
             "-q:v",
             "2",
-            f"{VIDEO_PATH}.stream0.frame{frame_name}.bmp",
+            output_bmp,
         ]
         subprocess.run(cmd, check=True)
+        convert_image_to_tensor(output_bmp)
 
     # This video was generated by running the following:
     # ffmpeg -f lavfi -i testsrc=duration=5:size=640x360:rate=25,format=yuv420p -c:v libaom-av1 -crf 30 -colorspace bt709 -color_primaries bt709 -color_trc bt709 av1_video.mkv
@@ -98,6 +122,7 @@ def main():
 
     for frame in FRAMES:
         frame_name = f"{frame:06d}"
+        output_bmp = f"{VIDEO_PATH}.stream0.frame{frame_name}.bmp"
         cmd = [
             "ffmpeg",
             "-y",
@@ -109,20 +134,10 @@ def main():
             "vfr",
             "-q:v",
             "2",
-            f"{VIDEO_PATH}.stream0.frame{frame_name}.bmp",
+            output_bmp,
         ]
         subprocess.run(cmd, check=True)
-
-    for bmp in [f for f in os.listdir(RESOURCES_DIR) if f.endswith(".bmp")]:
-        bmp_path = os.path.join(RESOURCES_DIR, bmp)
-        subprocess.run(
-            [
-                sys.executable,
-                os.path.join(TORCHCODEC_PATH, "test", "convert_image_to_tensor.py"),
-                bmp_path,
-            ]
-        )
-        os.remove(bmp_path)
+        convert_image_to_tensor(output_bmp)
 
 
 if __name__ == "__main__":