Fix up writeprotos to be robust to DECORDErrors

skelshop/cmd/iden/writeprotos.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import sys
 from csv import DictReader
 from functools import lru_cache
 from itertools import groupby
@@ -85,7 +86,14 @@ def bestcands(video_idx: int):
     for video_idx, frames in frame_grouped.items():
         video = corpus[video_idx]["video"]
         vid_read = decord_video_reader(str(video))
-        for frame_num, frame in read_numpy_chunks(vid_read, frames.keys()):
+        numpy_chunks = read_numpy_chunks(vid_read, frames.keys())
+        for frame_num, frame in numpy_chunks:
+            if frame is None:
+                print(
+                    f"Skipping {video}, frame {frame_num: 5d} due to error",
+                    file=sys.stderr,
+                )
+                continue
             for clus_dir, skel, extractor in frames[frame_num]:
                 PREFIX = "openpose-"
                 assert extractor.startswith(PREFIX)

skelshop/utils/video.py

@@ -76,9 +76,20 @@ def decord_video_reader(path):
 def read_numpy_chunks(
     video_reader, frame_idxs, batch_size=DEFAULT_FRAME_BATCH_SIZE, offset=0
 ):
+    from decord import DECORDError
+
     for frame_idx_batch in chunked(frame_idxs, batch_size):
-        batch_frames = video_reader.get_batch(
-            [frame_idx + offset for frame_idx in frame_idx_batch]
-        )
-        for frame_idx, frame in zip(frame_idx_batch, batch_frames.asnumpy()):
-            yield frame_idx, frame
+        frame_nums = [frame_idx + offset for frame_idx in frame_idx_batch]
+        try:
+            batch_frames = video_reader.get_batch(frame_nums)
+        except DECORDError:
+            video_reader.seek(0)
+            for frame_idx, frame_num in zip(frame_idx_batch, frame_nums):
+                try:
+                    yield frame_idx, video_reader[frame_num].asnumpy()
+                except DECORDError:
+                    yield frame_idx, None
+                    video_reader.seek(0)
+        else:
+            for frame_idx, frame in zip(frame_idx_batch, batch_frames.asnumpy()):
+                yield frame_idx, frame