Fix: ObjectDetectionPipeline batch inference only returns first image results

src/transformers/pipelines/object_detection.py

@@ -159,21 +159,23 @@ def unnormalize(bbox):
         else:
             # This is a regular ForObjectDetectionModel
             raw_annotations = self.image_processor.post_process_object_detection(model_outputs, threshold, target_size)
-            raw_annotation = raw_annotations[0]
-            scores = raw_annotation["scores"]
-            labels = raw_annotation["labels"]
-            boxes = raw_annotation["boxes"]
-
-            raw_annotation["scores"] = scores.tolist()
-            raw_annotation["labels"] = [self.model.config.id2label[label.item()] for label in labels]
-            raw_annotation["boxes"] = [self._get_bounding_box(box) for box in boxes]
-
-            # {"scores": [...], ...} --> [{"score":x, ...}, ...]
-            keys = ["score", "label", "box"]
-            annotation = [
-                dict(zip(keys, vals))
-                for vals in zip(raw_annotation["scores"], raw_annotation["labels"], raw_annotation["boxes"])
-            ]
+            annotation = []
+            for raw_annotation in raw_annotations:
+                scores = raw_annotation["scores"]
+                labels = raw_annotation["labels"]
+                boxes = raw_annotation["boxes"]
+
+                raw_annotation["scores"] = scores.tolist()
+                raw_annotation["labels"] = [self.model.config.id2label[label.item()] for label in labels]
+                raw_annotation["boxes"] = [self._get_bounding_box(box) for box in boxes]
+
+                keys = ["score", "label", "box"]
+                annotation.append(
+                    [
+                        dict(zip(keys, vals))
+                        for vals in zip(raw_annotation["scores"], raw_annotation["labels"], raw_annotation["boxes"])
+                    ]
+                )
 
         return annotation