From 50d76bbd1161f3317db15bac232023efb043ffd4 Mon Sep 17 00:00:00 2001
From: Frankie Robertson <frankie@robertson.name>
Date: Mon, 22 Nov 2021 11:18:46 +0200
Subject: [PATCH] Fix up writeprotos to be robust to DECORDErrors

---
 skelshop/cmd/iden/writeprotos.py | 10 +++++++++-
 skelshop/utils/video.py          | 21 ++++++++++++++++-----
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/skelshop/cmd/iden/writeprotos.py b/skelshop/cmd/iden/writeprotos.py
index ccb9c7d..92dd5c1 100644
--- a/skelshop/cmd/iden/writeprotos.py
+++ b/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)
diff --git a/skelshop/utils/video.py b/skelshop/utils/video.py
index f1a9290..8486151 100644
--- a/skelshop/utils/video.py
+++ b/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