Unverified 提交 9da6d0f9 authored 作者: Glenn Jocher's avatar Glenn Jocher 提交者: GitHub

Add LoadImages._cv2_rotate() (#9249)

Optional manual rotation code per iPhone rotation issue in https://github.com/ultralytics/yolov5/issues/8493Signed-off-by: 's avatarGlenn Jocher <glenn.jocher@ultralytics.com> Signed-off-by: 's avatarGlenn Jocher <glenn.jocher@ultralytics.com>
上级 ea981990
...@@ -213,7 +213,7 @@ class LoadImages: ...@@ -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: ...@@ -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: ...@@ -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
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论