diff --git a/.gitignore b/.gitignore index dbbd4a28bc..105bc5165b 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ +.idea/ + # Tensorflow checkpoints *.ckpt snapshot-* diff --git a/deeplabcut/create_project/new.py b/deeplabcut/create_project/new.py index 0c9f178b8d..65207c9df2 100644 --- a/deeplabcut/create_project/new.py +++ b/deeplabcut/create_project/new.py @@ -15,7 +15,7 @@ from deeplabcut import DEBUG import shutil -def create_new_project(project, experimenter, videos, working_directory=None, copy_videos=False,videotype='.avi'): +def create_new_project(project, experimenter, videos, working_directory=None, copy_videos=False,videotype='.avi', videoReader = None): """Creates a new project directory, sub-directories and a basic configuration file. The configuration file is loaded with the default values. Change its parameters to your projects need. Parameters @@ -134,7 +134,10 @@ def create_new_project(project, experimenter, videos, working_directory=None, co except: rel_video_path = os.readlink(str(video)) - vcap = cv2.VideoCapture(rel_video_path) + if videoReader is not None: + vcap = videoReader(rel_video_path) + else: + vcap = cv2.VideoCapture(rel_video_path) if vcap.isOpened(): width = int(vcap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(vcap.get(cv2.CAP_PROP_FRAME_HEIGHT)) diff --git a/deeplabcut/generate_training_dataset/frame_extraction.py b/deeplabcut/generate_training_dataset/frame_extraction.py index ea8aa298fc..3445707e9d 100644 --- a/deeplabcut/generate_training_dataset/frame_extraction.py +++ b/deeplabcut/generate_training_dataset/frame_extraction.py @@ -7,7 +7,7 @@ """ -def extract_frames(config,mode='automatic',algo='kmeans',crop=False,userfeedback=True,cluster_step=1,cluster_resizewidth=30,cluster_color=False,opencv=True): +def extract_frames(config,mode='automatic',algo='kmeans',crop=False,userfeedback=True,cluster_step=1,cluster_resizewidth=30,cluster_color=False,opencv=True, videoReader = None): """ Extracts frames from the videos in the config.yaml file. Only the videos in the config.yaml will be used to select the frames.\n Use the function ``add_new_video`` at any stage of the project to add new videos to the config file and extract their frames. @@ -88,6 +88,10 @@ def extract_frames(config,mode='automatic',algo='kmeans',crop=False,userfeedback from deeplabcut.utils import auxiliaryfunctions from matplotlib.widgets import RectangleSelector + # Use the OpenCV interface for the custom reader + if videoReader is not None: + opencv = True + if mode == "manual": wd = Path(config).resolve().parents[0] os.chdir(str(wd)) @@ -128,7 +132,12 @@ def extract_frames(config,mode='automatic',algo='kmeans',crop=False,userfeedback if askuser=='y' or askuser=='yes' or askuser=='Ja' or askuser=='ha': # multilanguage support :) #indexlength = int(np.ceil(np.log10(clip.duration * clip.fps))) - if opencv: + if videoReader is not None: + cap = videoReader(video) + fps = cap.get(5) + nframes = int(cap.get(7)) + duration = nframes * 1. / fps + elif opencv: cap=cv2.VideoCapture(video) fps = cap.get(5) #https://docs.opencv.org/2.4/modules/highgui/doc/reading_and_writing_images_and_video.html#videocapture-get nframes = int(cap.get(7)) diff --git a/deeplabcut/pose_estimation_tensorflow/predict_videos.py b/deeplabcut/pose_estimation_tensorflow/predict_videos.py index 73e3ffbd44..49f051d2d6 100755 --- a/deeplabcut/pose_estimation_tensorflow/predict_videos.py +++ b/deeplabcut/pose_estimation_tensorflow/predict_videos.py @@ -31,7 +31,7 @@ # Loading data, and defining model folder #################################################### -def analyze_videos(config,videos,videotype='avi',shuffle=1,trainingsetindex=0,gputouse=None,save_as_csv=False, destfolder=None): +def analyze_videos(config,videos,videotype='avi',shuffle=1,trainingsetindex=0,gputouse=None,save_as_csv=False, destfolder=None, videoReader=None, overwrite=False): """ Makes prediction based on a trained network. The index of the trained network is specified by parameters in the config file (in particular the variable 'snapshotindex') @@ -146,11 +146,12 @@ def analyze_videos(config,videos,videotype='avi',shuffle=1,trainingsetindex=0,gp ################################################## # Datafolder ################################################## + Videos=auxiliaryfunctions.Getlistofvideos(videos,videotype) if len(Videos)>0: #looping over videos for video in Videos: - AnalyzeVideo(video,DLCscorer,trainFraction,cfg,dlc_cfg,sess,inputs, outputs,pdindex,save_as_csv, destfolder) + AnalyzeVideo(video,DLCscorer,trainFraction,cfg,dlc_cfg,sess,inputs, outputs,pdindex,save_as_csv, destfolder, videoReader = videoReader, overwrite=overwrite) os.chdir(str(start_path)) print("The videos are analyzed. Now your research can truly start! \n You can create labeled videos with 'create_labeled_video'.") @@ -251,7 +252,7 @@ def GetPoseS(cfg,dlc_cfg, sess, inputs, outputs,cap,nframes): return PredicteData,nframes -def AnalyzeVideo(video,DLCscorer,trainFraction,cfg,dlc_cfg,sess,inputs, outputs,pdindex,save_as_csv, destfolder=None): +def AnalyzeVideo(video,DLCscorer,trainFraction,cfg,dlc_cfg,sess,inputs, outputs,pdindex,save_as_csv, destfolder=None, videoReader=None, overwrite=False): ''' Helper function for analyzing a video ''' print("Starting to analyze % ", video) vname = Path(video).stem @@ -260,11 +261,18 @@ def AnalyzeVideo(video,DLCscorer,trainFraction,cfg,dlc_cfg,sess,inputs, outputs, dataname = os.path.join(destfolder,vname + DLCscorer + '.h5') try: # Attempt to load data... - pd.read_hdf(dataname) - print("Video already analyzed!", dataname) + if overwrite: + raise FileNotFoundError + else: + pd.read_hdf(dataname) + print("Video already analyzed!", dataname) except FileNotFoundError: print("Loading ", video) - cap=cv2.VideoCapture(video) + # If a Video Reader was specified, use it instead of OpenCV + if videoReader is not None: + cap = videoReader(video) + else: + cap=cv2.VideoCapture(video) fps = cap.get(5) #https://docs.opencv.org/2.4/modules/highgui/doc/reading_and_writing_images_and_video.html#videocapture-get nframes = int(cap.get(7))