added option to use the frames - python demo2.py image (#166)

demo2.py

@@ -1,6 +1,14 @@
+##############################################
+# Demo file
+# python demo2.py --- will use video
+# python demo2.py image --- will use images
+#############################################
+
 import numpy as np
 import cv2
 import pybgs as bgs
+import sys
+import glob
 
 print("OpenCV Version: {}".format(cv2.__version__))
 
@@ -65,52 +73,92 @@ def check_opencv_version(major):
 algorithms.append(bgs.ViBe())
 algorithms.append(bgs.CodeBook())
 
+
+# check if we want to use the images
+image = False
+if (len(sys.argv) == 2):
+    if(sys.argv[1] == "image"):
+        image = True
+        img_folder = "dataset/frames"
+        img_array = sorted(glob.iglob(img_folder + '/*.png'))
+
 video_file = "dataset/video.avi"
 
+
 for algorithm in algorithms:
   print("Running ", algorithm.__class__)
-
-  capture = cv2.VideoCapture(video_file)
-  while not capture.isOpened():
-    capture = cv2.VideoCapture(video_file)
-    cv2.waitKey(1000)
-    print("Wait for the header")
-
-  #pos_frame = capture.get(cv2.cv.CV_CAP_PROP_POS_FRAMES)
-  #pos_frame = capture.get(cv2.CV_CAP_PROP_POS_FRAMES)
-  pos_frame = capture.get(1)
-  while True:
-    flag, frame = capture.read()
-
-    if flag:
-      cv2.imshow('video', frame)
+
+  if(image):
+    # loop x times as files in our folder
+    for x in range(0, len(img_array)):
+
+        # we can loop now through our array of images
+        img_path = img_array[x]
+
+        # read file into open cv and apply to algorithm to generate background model
+        img = cv2.imread(img_path)
+        img_output = algorithm.apply(img)
+        img_bgmodel = algorithm.getBackgroundModel()
+
+        # show images in python imshow window
+        cv2.imshow('image', img)
+        cv2.imshow('img_output', img_output)
+        cv2.imshow('img_bgmodel', img_bgmodel)
+
+        # we need waitKey otherwise it wont display the image
+        if 0xFF & cv2.waitKey(10) == 27:
+          break
+
+        # Comment out to save images to bg and fg folder
+        #img_bg = img_path.replace(img_folder, "output/bg")
+        #img_fg = img_path.replace(img_folder, "output/fg")
+        #cv2.imwrite(img_bg, img_bgmodel)
+        #cv2.imwrite(img_fg, img_output)
+
+        print("Frames left: " + str(len(img_array)-x))
+
+  else:
+
+      capture = cv2.VideoCapture(video_file)
+      while not capture.isOpened():
+        capture = cv2.VideoCapture(video_file)
+        cv2.waitKey(1000)
+        print("Wait for the header")
+
       #pos_frame = capture.get(cv2.cv.CV_CAP_PROP_POS_FRAMES)
       #pos_frame = capture.get(cv2.CV_CAP_PROP_POS_FRAMES)
       pos_frame = capture.get(1)
-      #print str(pos_frame)+" frames"
-
-      img_output = algorithm.apply(frame)
-      img_bgmodel = algorithm.getBackgroundModel()
-
-      cv2.imshow('img_output', img_output)
-      cv2.imshow('img_bgmodel', img_bgmodel)
-
-    else:
-      #capture.set(cv2.cv.CV_CAP_PROP_POS_FRAMES, pos_frame-1)
-      #capture.set(cv2.CV_CAP_PROP_POS_FRAMES, pos_frame-1)
-      #capture.set(1, pos_frame-1)
-      #print "Frame is not ready"
-      cv2.waitKey(1000)
-      break
-
-    if 0xFF & cv2.waitKey(10) == 27:
-      break
-
-    #if capture.get(cv2.cv.CV_CAP_PROP_POS_FRAMES) == capture.get(cv2.cv.CV_CAP_PROP_FRAME_COUNT):
-    #if capture.get(cv2.CV_CAP_PROP_POS_FRAMES) == capture.get(cv2.CV_CAP_PROP_FRAME_COUNT):
-    #if capture.get(1) == capture.get(cv2.CV_CAP_PROP_FRAME_COUNT):
-      #break
-
-  cv2.destroyAllWindows()
+      while True:
+        flag, frame = capture.read()
+
+        if flag:
+          cv2.imshow('video', frame)
+          #pos_frame = capture.get(cv2.cv.CV_CAP_PROP_POS_FRAMES)
+          #pos_frame = capture.get(cv2.CV_CAP_PROP_POS_FRAMES)
+          pos_frame = capture.get(1)
+          #print str(pos_frame)+" frames"
+
+          img_output = algorithm.apply(frame)
+          img_bgmodel = algorithm.getBackgroundModel()
+
+          cv2.imshow('img_output', img_output)
+          cv2.imshow('img_bgmodel', img_bgmodel)
+
+        else:
+          #capture.set(cv2.cv.CV_CAP_PROP_POS_FRAMES, pos_frame-1)
+          #capture.set(cv2.CV_CAP_PROP_POS_FRAMES, pos_frame-1)
+          #capture.set(1, pos_frame-1)
+          #print "Frame is not ready"
+          cv2.waitKey(1000)
+          break
+
+        if 0xFF & cv2.waitKey(10) == 27:
+          break
+
+        #if capture.get(cv2.cv.CV_CAP_PROP_POS_FRAMES) == capture.get(cv2.cv.CV_CAP_PROP_FRAME_COUNT):
+        #if capture.get(cv2.CV_CAP_PROP_POS_FRAMES) == capture.get(cv2.CV_CAP_PROP_FRAME_COUNT):
+        #if capture.get(1) == capture.get(cv2.CV_CAP_PROP_FRAME_COUNT):
+          #break
 
 print("Finished")
+cv2.destroyAllWindows()