From 1d3c59c3511987a2d9fb9221678ac59379bf5e27 Mon Sep 17 00:00:00 2001
From: torzdf <36920800+torzdf@users.noreply.github.com>
Date: Wed, 20 Mar 2024 15:24:30 +0000
Subject: [PATCH] Bugfix: Extract error on rotate faces

---
 plugins/extract/detect/_base.py | 23 ++++++++++++++++++++---
 1 file changed, 20 insertions(+), 3 deletions(-)

diff --git a/plugins/extract/detect/_base.py b/plugins/extract/detect/_base.py
index c85f75776b..a256e67c45 100644
--- a/plugins/extract/detect/_base.py
+++ b/plugins/extract/detect/_base.py
@@ -283,8 +283,26 @@ def _predict(self, batch: BatchType) -> DetectorBatch:
                 if angle == 0:
                     batch.prediction = pred
                 else:
-                    batch.prediction = np.array([b if b.any() else p
-                                                 for b, p in zip(batch.prediction, pred)])
+                    try:
+                        batch.prediction = np.array([b if b.any() else p
+                                                    for b, p in zip(batch.prediction, pred)])
+                    except ValueError as err:
+                        # If batches are different sizes after rotation Numpy will error, so we
+                        # need to explicitly set the dtype to 'object' rather than let it infer
+                        # numpy error:
+                        # ValueError: setting an array element with a sequence. The requested array
+                        # has an inhomogeneous shape after 1 dimensions. The detected shape was
+                        # (8,) + inhomogeneous part
+                        if "inhomogeneous" in str(err):
+                            batch.prediction = np.array([b if b.any() else p
+                                                         for b, p in zip(batch.prediction, pred)],
+                                                        dtype="object")
+                            logger.trace(  # type:ignore[attr-defined]
+                                "Mismatched array sizes, setting dtype to object: %s",
+                                [p.shape for p in batch.prediction])
+                        else:
+                            raise
+
                 logger.trace("angle: %s, filenames: %s, "  # type:ignore[attr-defined]
                              "prediction: %s",
                              angle, batch.filename, pred)
@@ -307,7 +325,6 @@ def _predict(self, batch: BatchType) -> DetectorBatch:
             found_faces = T.cast(list[np.ndarray], ([face if not found.any() else found
                                                      for face, found in zip(batch.prediction,
                                                                             found_faces)]))
-
             if all(face.any() for face in found_faces):
                 logger.trace("Faces found for all images")  # type:ignore[attr-defined]
                 break