diff --git a/utils/dataloaders.py b/utils/dataloaders.py index a4e6c0cfef18..5f86f83786db 100755 --- a/utils/dataloaders.py +++ b/utils/dataloaders.py @@ -213,7 +213,7 @@ def __init__(self, path, img_size=640, stride=32, auto=True, transforms=None): 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 @@ def __next__(self): 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 @@ def __next__(self): 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