From d024e7b1568668f5a8e5eaf6e1c30a49446ea27c Mon Sep 17 00:00:00 2001 From: Kaiyue Yang Date: Tue, 25 Jun 2024 11:55:53 -0700 Subject: [PATCH] [torchcodec]Remove usage of cv2 in smoke_test Summary: The smoke test relies on cv2. But in conda its not pre-installed. Use torchvision instead. Reviewed By: NicolasHug Differential Revision: D58980922 fbshipit-source-id: 7a71d4b555c6f6330af4ec3add207871a7a3bc01 --- test/decoders/manual_smoke_test.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test/decoders/manual_smoke_test.py b/test/decoders/manual_smoke_test.py index 7fe5878b..4a1b7b17 100644 --- a/test/decoders/manual_smoke_test.py +++ b/test/decoders/manual_smoke_test.py @@ -1,7 +1,7 @@ import os -import cv2 import torchcodec +from torchvision.io.image import write_png decoder = torchcodec.decoders._core.create_from_file( os.path.dirname(__file__) + "/../resources/nasa_13013.mp4" @@ -10,5 +10,5 @@ frame = torchcodec.decoders._core.get_frame_at_index( decoder, stream_index=3, frame_index=180 ) - -cv2.imwrite("frame180.png", frame.numpy()) +frame = frame.permute(2, 0, 1) +write_png(frame, "frame180.png")