diff --git a/.gitignore b/.gitignore
index 1e76239..4851707 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,4 +2,5 @@
 __pycache__
 him_projects
 *.hef
-*.avi
\ No newline at end of file
+*.avi
+venv
\ No newline at end of file
diff --git a/heatmap_test.py b/heatmap_test.py
index 812a04e..2c38fdc 100644
--- a/heatmap_test.py
+++ b/heatmap_test.py
@@ -1,5 +1,5 @@
 import cv2
-from video_loader import VideoLoader
+from video_loader import VideoLoader, VideoType
 from background_heatmap import BackgroundHeatmap, BackgroundHeatmapFromGroup
 import numpy as np
 
@@ -30,7 +30,7 @@ if __name__ == '__main__':
         groupheatmap = BackgroundHeatmapFromGroup(group)
         for recording in group[1]:
             first = True
-            frames = VideoLoader.extract_frames(recording)
+            frames = VideoLoader.extract_frames(recording, video_type=VideoType.DEPTH)
             for frame in frames:
                 groupheatmap.update(frame)
                 cv2.imshow("IR", frame)
@@ -49,7 +49,7 @@ if __name__ == '__main__':
                 cv2.imshow("BGF diff", frame_with_keypoints)
                 cv2.moveWindow("Background filter", 600, 100)
                 cv2.moveWindow("BGF diff", 1200, 100)
-                if cv2.waitKey(1) & 0xFF == ord('q'):
+                if cv2.waitKey(1000) & 0xFF == ord('q'):
                     break
             heatmaps.append(groupheatmap.heatmap)
             # Reset current heatmap and sum
@@ -60,7 +60,7 @@ if __name__ == '__main__':
             imname = "Heatmap " + str(i)
             cv2.imshow(imname, bgh)
             cv2.moveWindow(imname, 500 * i, 0)
-        cv2.waitKey(2000)
+        cv2.waitKey()
         cv2.destroyAllWindows()
 
     # out.release()