Detecting and removing corrupt video files from the dataset

DPI/utils.py

@@ -216,3 +216,41 @@ def video_from_array(arr, high_noise, filename):
                 frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
                 out.write(frame)
             out.release()
+
+
+class CorruptVideos:
+    def __init__(self, dir_path):
+        self.dir_path = dir_path
+
+    def _is_video_corrupt(self,filepath):
+        """
+        Check if a video file is corrupt.
+
+        Args:
+            filepath (str): Path to the video file.
+
+        Returns:
+            bool: True if the video is corrupt, False otherwise.
+        """
+        # Open the video file
+        cap = cv2.VideoCapture(filepath)
+        if not cap.isOpened():
+            return True
+        ret, frame = cap.read()
+        if not ret:
+            return True
+        cap.release()
+        return False
+
+    def _delete_corrupt_video(self, filepath):
+        os.remove(filepath)
+
+    def is_video_corrupt(self, delete=False):
+        for filename in os.listdir(self.dir_path):
+            filepath = os.path.join(self.dir_path, filename)
+            if filepath.endswith(".mp4"):
+                if self._is_video_corrupt(filepath):
+                    print(f"{filepath} is corrupt.")
+                    if delete:
+                        self._delete_corrupt_video(filepath)
+                        print(f"Deleted {filepath}")