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>
This commit is contained in:
1 parent
ea98199041
commit
9da6d0f9f5
1 file changed
+17
-3
+17
-3
@@ -213,7 +213,7 @@ class LoadImages:
|
|||||||
self.auto = auto
|
self.auto = auto
|
||||||
self.transforms = transforms # optional
|
self.transforms = transforms # optional
|
||||||
if any(videos):
|
if any(videos):
|
||||||
self.new_video(videos[0]) # new video
|
self._new_video(videos[0]) # new video
|
||||||
else:
|
else:
|
||||||
self.cap = None
|
self.cap = None
|
||||||
assert self.nf > 0, f'No images or videos found in {p}. ' \
|
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
|
if self.count == self.nf: # last video
|
||||||
raise StopIteration
|
raise StopIteration
|
||||||
path = self.files[self.count]
|
path = self.files[self.count]
|
||||||
self.new_video(path)
|
self._new_video(path)
|
||||||
ret_val, im0 = self.cap.read()
|
ret_val, im0 = self.cap.read()
|
||||||
|
|
||||||
self.frame += 1
|
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}: '
|
s = f'video {self.count + 1}/{self.nf} ({self.frame}/{self.frames}) {path}: '
|
||||||
|
|
||||||
else:
|
else:
|
||||||
@@ -260,10 +261,23 @@ class LoadImages:
|
|||||||
|
|
||||||
return path, im, im0, self.cap, s
|
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.frame = 0
|
||||||
self.cap = cv2.VideoCapture(path)
|
self.cap = cv2.VideoCapture(path)
|
||||||
self.frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
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):
|
def __len__(self):
|
||||||
return self.nf # number of files
|
return self.nf # number of files
|
||||||
|
|||||||
Reference in new issue
Block a user