Add check for rotation metadata

Workaround for OpenCV bug, where orientation of video is not detected properly.
opencv/opencv#15499

emotion_detection.py

@@ -2,6 +2,7 @@
 import os
 
 import cv2
+import ffmpeg
 import torch
 import torch.nn as nn
 from PIL import Image
@@ -66,6 +67,37 @@ def preprocess(image):
     image = image.unsqueeze(0)
     return image
 
+def check_rotation(video_file):
+    """Checks if given video file has rotation metadata
+
+    Used as workaround for OpenCV bug 
+    https://github.com/opencv/opencv/issues/15499, where the rotation metadata 
+    in a video is not used by cv2.VideoCapture. May have to be removed/tweaked 
+    when bug is fixed in a new opencv-python release.
+
+    Parameters
+    ----------
+    video_file : str
+        Path to video file to be checked
+
+    Returns
+    -------
+    rotate_code : enum or None
+        Flag fror cv2.rotate to decide how much to rotate the image. 
+        None if no rotation is required
+    """    
+    meta_dict = ffmpeg.probe(video_file)
+    rotation_angle = int(meta_dict["streams"][0]["tags"]["rotate"])
+
+    rotate_code = None
+    if rotation_angle == 90:
+        rotate_code = cv2.ROTATE_90_CLOCKWISE
+    elif rotation_angle == 180:
+        rotate_code = cv2.ROTATE_180
+    elif rotation_angle == 270:
+        rotate_code = cv2.ROTATE_90_COUNTERCLOCKWISE
+
+    return rotate_code
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = CNNModel().to(device)
@@ -81,11 +113,15 @@ def preprocess(image):
 
 if opt.file_type == "video":
     os.makedirs(opt.output_video_directory, exist_ok=True)
-    counter = 0
     cap = cv2.VideoCapture(opt.video_file)
+    rotate_code = check_rotation(opt.video_file)
+
+    counter = 0
     while cap.isOpened():
         ret, frame = cap.read()
         if ret:
+            if rotate_code is not None: # Rotate frame if necessary
+                frame = cv2.rotate(frame, rotate_code)
             gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
             faces = face_cascade.detectMultiScale(
                 gray, scaleFactor=1.5, minNeighbors=5)