Skip to content

Commit 406e336

Browse files
add opencv benchmark
1 parent 87b98e8 commit 406e336

File tree

2 files changed

+70
-0
lines changed

2 files changed

+70
-0
lines changed

benchmarks/decoders/benchmark_decoders.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
TorchCodecPublic,
2929
TorchCodecPublicNonBatch,
3030
TorchVision,
31+
OpenCVDecoder,
3132
)
3233

3334

@@ -61,6 +62,7 @@ class DecoderKind:
6162
{"backend": "video_reader"},
6263
),
6364
"torchaudio": DecoderKind("TorchAudio", TorchAudioDecoder),
65+
"opencv": DecoderKind("OpenCV", OpenCVDecoder),
6466
}
6567

6668

benchmarks/decoders/benchmark_decoders_library.py

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,74 @@ def decode_and_resize(self, video_file, pts_list, height, width, device):
145145
]
146146
return frames
147147

148+
class OpenCVDecoder(AbstractDecoder):
149+
def __init__(self):
150+
import cv2.videoio_registry as vr
151+
152+
self._print_each_iteration_time = False
153+
api_pref = None
154+
for backend in vr.getStreamBufferedBackends():
155+
if not vr.hasBackend(backend):
156+
continue
157+
if not vr.isBackendBuiltIn(backend):
158+
_, abi, api = vr.getStreamBufferedBackendPluginVersion(backend)
159+
if (abi < 1 or (abi == 1 and api < 2)):
160+
continue
161+
api_pref = backend
162+
break
163+
self._backend = api_pref
164+
165+
def decode_frames(self, video_file, pts_list):
166+
import cv2
167+
168+
cap = cv2.VideoCapture(video_file, self._backend, [])
169+
if not cap.isOpened():
170+
raise ValueError("Could not open video stream")
171+
172+
fps = cap.get(cv2.CAP_PROP_FPS)
173+
frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
174+
approx_frame_numbers = [int(pts * fps) for pts in pts_list]
175+
176+
current_frame = 0
177+
frames = []
178+
while True:
179+
ok = cap.grab()
180+
if not ok:
181+
break
182+
if current_frame in approx_frame_numbers: # only decompress needed
183+
ret, frame = cap.retrieve()
184+
if ret:
185+
frames.append(frame)
186+
187+
if len(frames) == len(approx_frame_numbers):
188+
break
189+
current_frame += 1
190+
cap.release()
191+
return frames
192+
193+
def decode_first_n_frames(self, video_file, n):
194+
import cv2
195+
196+
cap = cv2.VideoCapture(video_file, self._backend, [])
197+
if not cap.isOpened():
198+
raise ValueError("Could not open video stream")
199+
200+
frames = []
201+
for i in range(n):
202+
ok = cap.grab()
203+
if not ok:
204+
break
205+
ret, frame = cap.retrieve()
206+
if ret:
207+
frames.append(frame)
208+
cap.release()
209+
return frames
210+
211+
def decode_and_resize(self, video_file, pts_list, height, width, device):
212+
import cv2
213+
frames = [cv2.resize(frame, (width, height)) for frame in self.decode_frames(video_file, pts_list)]
214+
return frames
215+
148216

149217
class TorchCodecCore(AbstractDecoder):
150218
def __init__(self, num_threads=None, color_conversion_library=None, device="cpu"):

0 commit comments

Comments
 (0)