Add LoadImages._cv2_rotate() (#9249)

Optional manual rotation code per iPhone rotation issue in https://github.com/ultralytics/yolov5/issues/8493

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
pull/9258/head
Glenn Jocher 2022-09-01 22:30:26 +02:00 committed by GitHub
parent ea98199041
commit 9da6d0f9f5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 17 additions and 3 deletions

View File

@ -213,7 +213,7 @@ class LoadImages:
self.auto = auto
self.transforms = transforms # optional
if any(videos):
self.new_video(videos[0]) # new video
self._new_video(videos[0]) # new video
else:
self.cap = None
assert self.nf > 0, f'No images or videos found in {p}. ' \
@ -238,10 +238,11 @@ class LoadImages:
if self.count == self.nf: # last video
raise StopIteration
path = self.files[self.count]
self.new_video(path)
self._new_video(path)
ret_val, im0 = self.cap.read()
self.frame += 1
# im0 = self._cv2_rotate(im0) # for use if cv2 auto rotation is False
s = f'video {self.count + 1}/{self.nf} ({self.frame}/{self.frames}) {path}: '
else:
@ -260,10 +261,23 @@ class LoadImages:
return path, im, im0, self.cap, s
def new_video(self, path):
def _new_video(self, path):
# Create a new video capture object
self.frame = 0
self.cap = cv2.VideoCapture(path)
self.frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
self.orientation = int(self.cap.get(cv2.CAP_PROP_ORIENTATION_META)) # rotation degrees
# self.cap.set(cv2.CAP_PROP_ORIENTATION_AUTO, 0) # disable https://github.com/ultralytics/yolov5/issues/8493
def _cv2_rotate(self, im):
# Rotate a cv2 video manually
if self.orientation == 0:
return cv2.rotate(im, cv2.ROTATE_90_CLOCKWISE)
elif self.orientation == 180:
return cv2.rotate(im, cv2.ROTATE_90_COUNTERCLOCKWISE)
elif self.orientation == 90:
return cv2.rotate(im, cv2.ROTATE_180)
return im
def __len__(self):
return self.nf # number of files