Added functionality to dataset generator; now works with multiple videos, and...

Added functionality to dataset generator; now works with multiple videos, and extracts a given percentage per video.

src/python/utilities/GenerateDataset.py

@@ -5,7 +5,10 @@ import numpy as np
 import sys
 from pathlib import Path
 
-setFPS = 2
+#   Skips every 1/n frames, according to this ratio. 
+#   Maybe do random frame n*frameCount instead?
+percentExtractedPerVideo = 0.3
+firstNvideos = 6
 resizeWidth = 480
 resizeHeight = 270
 
@@ -13,11 +16,14 @@ groundsPath = "src/python/groundTruth/"
 
 ###     Generates a dataset of labelled images from one video (soon a whole folder) by comparing to ground truth.
 ###         Takes about 5 minutes for a 30 minute input.
-def generateDatasetFromFrames(inputVideoPath, desiredFPS):
-    decoder = cv2.VideoCapture(inputVideoPath)
-    inputVideo = Path(inputVideoPath).name
+def generateDatasetFromFrames(inputVideosPath, extractionRatio):
+    for i in range(firstNvideos):
+        curVideo = os.listdir(inputVideosPath)[i]
+        decoder = cv2.VideoCapture(inputVideosPath + curVideo)
 
-    sourceNoExt = os.path.splitext(inputVideo)[0]
+        desiredFPS = int(cv2.CAP_PROP_FPS * extractionRatio)
+
+        sourceNoExt = os.path.splitext(curVideo)[0]
         groundTruthName = "groundTruth_" + sourceNoExt + ".csv"
    
         # Looks for ground truth csv.
@@ -29,8 +35,8 @@ def generateDatasetFromFrames(inputVideoPath, desiredFPS):
                 break
 
         if groundFound == False:
-        print("No ground truth file found to label from.")
-        return False
+            print("No ground truth file found to label from, for file {file}.\n".format(file = curVideo))
+            continue
 
         step = int(decoder.get(cv2.CAP_PROP_FPS) / desiredFPS)
         try:
@@ -60,15 +66,16 @@ def generateDatasetFromFrames(inputVideoPath, desiredFPS):
                     groundTruthCSV.seek(0)
                     if found == False:
                         cv2.imwrite("dataset/noFish/" + "frame{number}{source}.jpg".format(number=frameNo, source=sourceNoExt), frame)
-        print("Frames saved as labelled images in dataset folder.")
-        return True
+            print("Frames saved as labelled images in dataset folder.\n")
+        except:
+            print("Something went wrong in reading file {file}".format(file = curVideo))
 
         finally:
             decoder.release()
             groundTruthCSV.close()  
 
 def main(video,fps):
-    print(generateDatasetFromFrames(video,fps))
+    generateDatasetFromFrames(video,fps)
 
 if __name__ == "__main__":
-    main("sourceVideos/Myggbukta-[2021-05-21_10-47-06]-408.mp4", setFPS) 
+    main("sourceVideos/", percentExtractedPerVideo) 
\ No newline at end of file