diff --git a/src/transformers/pipelines/object_detection.py b/src/transformers/pipelines/object_detection.py index 0a4fba996d7d..eaa2fdfa3a83 100644 --- a/src/transformers/pipelines/object_detection.py +++ b/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