Skip to content

Commit 52c61c5

Browse files
committed
fix(predictor): use image orieantion in batch grouping
Closes Kitware#92
1 parent 82dbb35 commit 52c61c5

File tree

1 file changed

+11
-10
lines changed

1 file changed

+11
-10
lines changed

src/nrtk_explorer/library/predictor.py

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -74,23 +74,24 @@ def eval(
7474
batches: dict = {}
7575
for image in images_with_ids:
7676
size = image.image.size
77+
exif_data = image.image.getexif()
78+
orientation = exif_data.get(274, None)
79+
# Swap dimensions if the orientation implies a 90° rotation
80+
if orientation in [5, 6, 7, 8]:
81+
size = (size[1], size[0])
7782
batches.setdefault(size, [])
7883
batches[size].append(image)
7984

8085
if batch_size != 0:
8186
self.batch_size = self.batch_size
8287
while self.batch_size > 0:
8388
try:
84-
predictions_in_baches = [
85-
zip(
86-
[image.id for image in imagesInBatch],
87-
self.pipeline(
88-
[image.image for image in imagesInBatch],
89-
batch_size=self.batch_size,
90-
),
91-
)
92-
for imagesInBatch in batches.values()
93-
]
89+
predictions_in_baches = []
90+
for imagesInBatch in batches.values():
91+
image_ids = [image.id for image in imagesInBatch]
92+
image_data = [image.image for image in imagesInBatch]
93+
pipeline_output = self.pipeline(image_data, batch_size=self.batch_size)
94+
predictions_in_baches.append(zip(image_ids, pipeline_output))
9495

9596
predictions_by_image_id = {
9697
image_id: predictions

0 commit comments

Comments
 (0)