@@ -213,7 +213,7 @@ def __init__(self, path, img_size=640, stride=32, auto=True, transforms=None):
213
213
self .auto = auto
214
214
self .transforms = transforms # optional
215
215
if any (videos ):
216
- self .new_video (videos [0 ]) # new video
216
+ self ._new_video (videos [0 ]) # new video
217
217
else :
218
218
self .cap = None
219
219
assert self .nf > 0 , f'No images or videos found in { p } . ' \
@@ -238,10 +238,11 @@ def __next__(self):
238
238
if self .count == self .nf : # last video
239
239
raise StopIteration
240
240
path = self .files [self .count ]
241
- self .new_video (path )
241
+ self ._new_video (path )
242
242
ret_val , im0 = self .cap .read ()
243
243
244
244
self .frame += 1
245
+ # im0 = self._cv2_rotate(im0) # for use if cv2 auto rotation is False
245
246
s = f'video { self .count + 1 } /{ self .nf } ({ self .frame } /{ self .frames } ) { path } : '
246
247
247
248
else :
@@ -260,10 +261,23 @@ def __next__(self):
260
261
261
262
return path , im , im0 , self .cap , s
262
263
263
- def new_video (self , path ):
264
+ def _new_video (self , path ):
265
+ # Create a new video capture object
264
266
self .frame = 0
265
267
self .cap = cv2 .VideoCapture (path )
266
268
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
267
281
268
282
def __len__ (self ):
269
283
return self .nf # number of files
0 commit comments