Skip to content

Commit 9da6d0f

Browse files
authored
Add LoadImages._cv2_rotate() (#9249)
Optional manual rotation code per iPhone rotation issue in #8493 Signed-off-by: Glenn Jocher <[email protected]> Signed-off-by: Glenn Jocher <[email protected]>
1 parent ea98199 commit 9da6d0f

File tree

1 file changed

+17
-3
lines changed

1 file changed

+17
-3
lines changed

utils/dataloaders.py

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -213,7 +213,7 @@ def __init__(self, path, img_size=640, stride=32, auto=True, transforms=None):
213213
self.auto = auto
214214
self.transforms = transforms # optional
215215
if any(videos):
216-
self.new_video(videos[0]) # new video
216+
self._new_video(videos[0]) # new video
217217
else:
218218
self.cap = None
219219
assert self.nf > 0, f'No images or videos found in {p}. ' \
@@ -238,10 +238,11 @@ def __next__(self):
238238
if self.count == self.nf: # last video
239239
raise StopIteration
240240
path = self.files[self.count]
241-
self.new_video(path)
241+
self._new_video(path)
242242
ret_val, im0 = self.cap.read()
243243

244244
self.frame += 1
245+
# im0 = self._cv2_rotate(im0) # for use if cv2 auto rotation is False
245246
s = f'video {self.count + 1}/{self.nf} ({self.frame}/{self.frames}) {path}: '
246247

247248
else:
@@ -260,10 +261,23 @@ def __next__(self):
260261

261262
return path, im, im0, self.cap, s
262263

263-
def new_video(self, path):
264+
def _new_video(self, path):
265+
# Create a new video capture object
264266
self.frame = 0
265267
self.cap = cv2.VideoCapture(path)
266268
self.frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
269+
self.orientation = int(self.cap.get(cv2.CAP_PROP_ORIENTATION_META)) # rotation degrees
270+
# self.cap.set(cv2.CAP_PROP_ORIENTATION_AUTO, 0) # disable https://github.com/ultralytics/yolov5/issues/8493
271+
272+
def _cv2_rotate(self, im):
273+
# Rotate a cv2 video manually
274+
if self.orientation == 0:
275+
return cv2.rotate(im, cv2.ROTATE_90_CLOCKWISE)
276+
elif self.orientation == 180:
277+
return cv2.rotate(im, cv2.ROTATE_90_COUNTERCLOCKWISE)
278+
elif self.orientation == 90:
279+
return cv2.rotate(im, cv2.ROTATE_180)
280+
return im
267281

268282
def __len__(self):
269283
return self.nf # number of files

0 commit comments

Comments
 (0)