Skip to content

Commit e295953

Browse files
committed
env. variable for colab, color flip openCV fixed for video analysis, etc.
1 parent 72ccb2a commit e295953

14 files changed

Lines changed: 96 additions & 59 deletions

File tree

README.md

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
1-
21
[![License: LGPL v3](https://img.shields.io/badge/License-LGPL%20v3-blue.svg)](https://www.gnu.org/licenses/lgpl-3.0)
3-
[![GitHub stars](https://img.shields.io/github/stars/AlexEMG/DeepLabCut.svg?style=social&label=Star)](https://github.com/AlexEMG/DeepLabCut)
42
[![Krihelimeter](http://krihelinator.xyz/badge/AlexEMG/DeepLabCut)](http://krihelinator.xyz/repositories/AlexEMG/DeepLabCut)
3+
[![GitHub forks](https://img.shields.io/github/forks/AlexEMG/DeepLabCut.svg?style=social&label=Fork)](https://github.com/AlexEMG/DeepLabCut)
4+
[![GitHub stars](https://img.shields.io/github/stars/AlexEMG/DeepLabCut.svg?style=social&label=Star)](https://github.com/AlexEMG/DeepLabCut)
5+
56

67
## DeepLabCut
78

@@ -31,6 +32,29 @@ Please check out [www.mousemotorlab.org/deeplabcut](https://www.mousemotorlab.or
3132
# [DEMO the code](/examples)
3233
We provide several Jupyter Notebooks: one that walks you through a demo dataset to test your installation, and another Notebook to run DeepLabCut from the begining on your own data. We also show you how to use the code in Docker, and on Google Colab.
3334

35+
# Why using DeepLabCut?:
36+
37+
- Due to transfer learning it requires **little training data** for multiple, challenging behaviors (see [Mathis et al.](https://www.nature.com/articles/s41593-018-0209-y) for details).
38+
<p align="center">
39+
<img src="docs/images/ErrorvsTrainingsetSize.png" width="60%">
40+
</p>
41+
42+
- Video anlysis is fast (see [Mathis/Warren](https://www.biorxiv.org/content/early/2018/10/30/457242) for details)
43+
<p align="center">
44+
<img src="docs/images/inferencespeed.png" width="60%">
45+
</p>
46+
47+
- The feature detectors are robust to video compression (see [Mathis/Warren](https://www.biorxiv.org/content/early/2018/10/30/457242) for details)
48+
<p align="center">
49+
<img src="docs/images/compressionrobustness.png" width="60%">
50+
</p>
51+
52+
- It allows 3D pose estimation with a single network and camera (see [Mathis/Warren](https://www.biorxiv.org/content/early/2018/10/30/457242) for details)
53+
<p align="center">
54+
<img src="docs/images/MouseLocomotion_warren.gif" width="25%">
55+
</p>
56+
57+
3458
# News (and in the news):
3559

3660
- Nov 2018: Various (post-hoc) analysis scripts contributed by users (and us) will be gathered at [DLCutils](https://github.com/AlexEMG/DLCutils). Feel free to contribute! In particular, there is a script guiding you through

deeplabcut/__init__.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
A Mathis, alexander.mathis@bethgelab.org
55
T Nath, nath@rowland.harvard.edu
66
M Mathis, mackenzie@post.harvard.edu
7+
8+
79
"""
810

911
import os
@@ -31,11 +33,14 @@
3133
from deeplabcut import pose_estimation_tensorflow
3234
from deeplabcut import utils
3335
from deeplabcut.create_project import create_new_project, add_new_videos, load_demo_data
34-
3536
from deeplabcut.generate_training_dataset import extract_frames
36-
from deeplabcut.refine_training_dataset import extract_outlier_frames,merge_datasets,filterpredictions
3737
from deeplabcut.generate_training_dataset import check_labels,create_training_dataset
3838

39+
if os.environ.get('Colab', default=False) == 'True':
40+
print("Project loaded in colab-mode. Apparently Colab has trouble loading statsmodels, so the smooting & outlier frame extraction is disabled. Sorry!")
41+
else:
42+
from deeplabcut.refine_training_dataset import extract_outlier_frames,merge_datasets,filterpredictions
43+
3944
#Direct import for convenience
4045
from deeplabcut.pose_estimation_tensorflow import train_network
4146
from deeplabcut.pose_estimation_tensorflow import evaluate_network

deeplabcut/generate_training_dataset/frame_extraction.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,7 @@ def extract_frames(config,mode,algo='uniform',crop=False,checkcropping=False):
9191

9292
#update to openCV
9393
clip = VideoFileClip(video)
94+
9495
indexlength = int(np.ceil(np.log10(clip.duration * clip.fps)))
9596

9697
if crop==True:

deeplabcut/pose_estimation_tensorflow/predict_videos.py

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -12,14 +12,11 @@
1212
####################################################
1313

1414
import os.path
15-
import sys
16-
1715
from deeplabcut.pose_estimation_tensorflow.nnet import predict
1816
from deeplabcut.pose_estimation_tensorflow.config import load_config
1917
from deeplabcut.pose_estimation_tensorflow.dataset.pose_dataset import data_to_input
2018

2119
from random import sample
22-
import pickle
2320
import time
2421
import pandas as pd
2522
import numpy as np
@@ -29,6 +26,8 @@
2926
from tqdm import tqdm
3027
import tensorflow as tf
3128
from deeplabcut.utils import auxiliaryfunctions
29+
import cv2
30+
from skimage.util import img_as_ubyte
3231

3332
####################################################
3433
# Loading data, and defining model folder
@@ -160,8 +159,8 @@ def analyze_videos(config,videos,shuffle=1,trainingsetindex=0,videotype='avi',gp
160159

161160

162161
def GetPoseF(cfg,dlc_cfg, sess, inputs, outputs,cap,nframes,batchsize,frame_buffer):
163-
''' note cfg here is for pose-tensorflow.'''
164-
from skimage.util import img_as_ubyte
162+
''' Batchwise prediction of pose '''
163+
165164
PredicteData = np.zeros((nframes, 3 * len(dlc_cfg['all_joints_names'])))
166165
batch_ind = 0 # keeps track of which image within a batch should be written to
167166
batch_num = 0 # keeps track of which batch you are at
@@ -188,6 +187,7 @@ def GetPoseF(cfg,dlc_cfg, sess, inputs, outputs,cap,nframes,batchsize,frame_buff
188187
pbar.update(step)
189188
ret, frame = cap.read()
190189
if ret:
190+
frame=cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
191191
if cfg['cropping']:
192192
frames[batch_ind] = img_as_ubyte(frame[cfg['y1']:cfg['y2'],cfg['x1']:cfg['x2']])
193193
else:
@@ -212,8 +212,7 @@ def GetPoseF(cfg,dlc_cfg, sess, inputs, outputs,cap,nframes,batchsize,frame_buff
212212
return PredicteData,nframes
213213

214214
def GetPoseS(cfg,dlc_cfg, sess, inputs, outputs,cap,nframes,frame_buffer):
215-
''' note cfg here is for pose-tensorflow.'''
216-
from skimage.util import img_as_ubyte
215+
''' Non batch wise pose estimation for video cap.'''
217216
if cfg['cropping']:
218217
print("Cropping based on the x1 = %s x2 = %s y1 = %s y2 = %s. You can adjust the cropping coordinates in the config.yaml file." %(cfg['x1'], cfg['x2'],cfg['y1'], cfg['y2']))
219218
nx=cfg['x2']-cfg['x1']
@@ -237,6 +236,7 @@ def GetPoseS(cfg,dlc_cfg, sess, inputs, outputs,cap,nframes,frame_buffer):
237236

238237
ret, frame = cap.read()
239238
if ret:
239+
frame=cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
240240
if cfg['cropping']:
241241
frame= img_as_ubyte(frame[cfg['y1']:cfg['y2'],cfg['x1']:cfg['x2']])
242242
else:
@@ -254,7 +254,7 @@ def GetPoseS(cfg,dlc_cfg, sess, inputs, outputs,cap,nframes,frame_buffer):
254254

255255
def AnalzyeVideo(video,DLCscorer,cfg,dlc_cfg,sess,inputs, outputs,pdindex,frame_buffer=10):
256256
#from moviepy.editor import VideoFileClip
257-
import cv2
257+
258258
print(video)
259259
#videotype = Path(video).suffix
260260
print("Starting % ", video)

deeplabcut/utils/make_labeled_video.py

Lines changed: 28 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,6 @@
3939
from deeplabcut.utils.video_processor import VideoProcessorCV as vp # used to CreateVideo
4040

4141

42-
4342
def get_cmap(n, name='hsv'):
4443
'''Returns a function that maps each index in 0, 1, ..., n-1 to a distinct
4544
RGB color; the keyword argument name must be a standard mpl colormap name.'''
@@ -148,7 +147,7 @@ def CreateVideoSlow(clip,Dataframe,tmpfolder,dotsize,colormap,alphavalue,pcutoff
148147
os.remove(file_name)
149148
os.chdir(start)
150149

151-
def create_labeled_video(config,video,shuffle=1,trainingsetindex=0,videotype='avi',save_frames=False,delete=False,displayedbodyparts='all'):
150+
def create_labeled_video(config,videos,shuffle=1,trainingsetindex=0,videotype='avi',save_frames=False,delete=False,displayedbodyparts='all',codec='X264'):
152151
"""
153152
Labels the bodyparts in a video. Make sure the video is already analyzed by the function 'analyze_video'
154153
@@ -157,7 +156,7 @@ def create_labeled_video(config,video,shuffle=1,trainingsetindex=0,videotype='av
157156
config : string
158157
Full path of the config.yaml file as a string.
159158
160-
video : list
159+
videos : list
161160
A list of string containing the full paths of the videos to analyze.
162161
163162
shuffle : int, optional
@@ -181,6 +180,8 @@ def create_labeled_video(config,video,shuffle=1,trainingsetindex=0,videotype='av
181180
from config.yaml are used orr a list of strings that are a subset of the full list.
182181
E.g. ['hand','Joystick'] for the demo Reaching-Mackenzie-2018-08-30/config.yaml to select only these two body parts.
183182
183+
codec: codec for labeled video. Options see http://www.fourcc.org/codecs.php [depends on your ffmpeg installation.]
184+
184185
Examples
185186
--------
186187
If you want to create the labeled video for only 1 video
@@ -211,16 +212,16 @@ def create_labeled_video(config,video,shuffle=1,trainingsetindex=0,videotype='av
211212

212213
bodyparts=auxiliaryfunctions.IntersectionofBodyPartsandOnesGivenbyUser(cfg,displayedbodyparts)
213214

214-
if [os.path.isdir(i) for i in video] == [True]:
215+
if [os.path.isdir(i) for i in videos] == [True]:
215216
print("Analyzing all the videos in the directory")
216-
videofolder= video[0]
217+
videofolder= videos[0]
217218
os.chdir(videofolder)
218-
videos = np.sort([fn for fn in os.listdir(os.curdir) if (videotype in fn)])
219-
print("Starting ", videofolder, videos)
219+
Videos = np.sort([fn for fn in os.listdir(os.curdir) if (videotype in fn)])
220+
print("Starting ", videofolder, Videos)
220221
else:
221-
videos = video
222+
Videos = videos
222223

223-
for video in videos:
224+
for video in Videos:
224225
videofolder= Path(video).parents[0] #where your folder with videos is.
225226
os.chdir(str(videofolder))
226227
videotype = Path(video).suffix
@@ -235,6 +236,7 @@ def create_labeled_video(config,video,shuffle=1,trainingsetindex=0,videotype='av
235236
Dataframe = pd.read_hdf(dataname)
236237
metadata=auxiliaryfunctions.LoadVideoMetadata(dataname)
237238
#print(metadata)
239+
datanames=[dataname]
238240
except FileNotFoundError:
239241
datanames=[fn for fn in os.listdir(os.curdir) if (vname in fn) and (".h5" in fn) and "resnet" in fn]
240242
if len(datanames)==0:
@@ -248,23 +250,24 @@ def create_labeled_video(config,video,shuffle=1,trainingsetindex=0,videotype='av
248250
Dataframe = pd.read_hdf(datanames[0])
249251
metadata=auxiliaryfunctions.LoadVideoMetadata(datanames[0])
250252

251-
#Loading cropping data used during analysis
252-
cropping=metadata['data']["cropping"]
253-
[x1,x2,y1,y2]=metadata['data']["cropping_parameters"]
254-
print(cropping,x1,x2,y1,y2)
255-
256-
if save_frames==True:
257-
tmpfolder = os.path.join(str(videofolder),'temp-' + vname)
258-
auxiliaryfunctions.attempttomakefolder(tmpfolder)
259-
clip = vp(video)
260-
#CreateVideoSlow(clip,Dataframe,tmpfolder,cfg["dotsize"],cfg["colormap"],cfg["alphavalue"],cfg["pcutoff"],cfg["cropping"],cfg["x1"],cfg["x2"],cfg["y1"],cfg["y2"],delete,DLCscorer,bodyparts)
261-
CreateVideoSlow(clip,Dataframe,tmpfolder,cfg["dotsize"],cfg["colormap"],cfg["alphavalue"],cfg["pcutoff"],cropping,x1,x2,y1,y2,delete,DLCscorer,bodyparts)
262-
else:
263-
clip = vp(fname = video,sname = os.path.join(vname + DLCscorer+'_labeled.mp4'))
264-
if cropping:
265-
print("Fast video creation has currently not been implemented for cropped videos. Please use 'save_frames=True' to get the video.")
253+
if len(datanames)>0:
254+
#Loading cropping data used during analysis
255+
cropping=metadata['data']["cropping"]
256+
[x1,x2,y1,y2]=metadata['data']["cropping_parameters"]
257+
print(cropping,x1,x2,y1,y2)
258+
259+
if save_frames==True:
260+
tmpfolder = os.path.join(str(videofolder),'temp-' + vname)
261+
auxiliaryfunctions.attempttomakefolder(tmpfolder)
262+
clip = vp(video)
263+
#CreateVideoSlow(clip,Dataframe,tmpfolder,cfg["dotsize"],cfg["colormap"],cfg["alphavalue"],cfg["pcutoff"],cfg["cropping"],cfg["x1"],cfg["x2"],cfg["y1"],cfg["y2"],delete,DLCscorer,bodyparts)
264+
CreateVideoSlow(clip,Dataframe,tmpfolder,cfg["dotsize"],cfg["colormap"],cfg["alphavalue"],cfg["pcutoff"],cropping,x1,x2,y1,y2,delete,DLCscorer,bodyparts)
266265
else:
267-
CreateVideo(clip,Dataframe,cfg["pcutoff"],cfg["dotsize"],cfg["colormap"],DLCscorer,bodyparts,cropping,x1,x2,y1,y2) #NEED TO ADD CROPPING!
266+
clip = vp(fname = video,sname = os.path.join(vname + DLCscorer+'_labeled.mp4'),codec=codec)
267+
if cropping:
268+
print("Fast video creation has currently not been implemented for cropped videos. Please use 'save_frames=True' to get the video.")
269+
else:
270+
CreateVideo(clip,Dataframe,cfg["pcutoff"],cfg["dotsize"],cfg["colormap"],DLCscorer,bodyparts,cropping,x1,x2,y1,y2) #NEED TO ADD CROPPING!
268271

269272
if __name__ == '__main__':
270273
parser = argparse.ArgumentParser()

deeplabcut/utils/plotting.py

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -136,15 +136,16 @@ def RunTrajectoryAnalysis(video,basefolder,scorer,videofolder,cfg,showfigures):
136136
# Looping analysis over video
137137
##################################################
138138

139-
def plot_trajectories(config,video,shuffle=1,trainingsetindex=0,videotype='.avi',showfigures=False):
139+
def plot_trajectories(config,videos,shuffle=1,trainingsetindex=0,videotype='.avi',showfigures=False):
140140
"""
141141
Plots the trajectories of various bodyparts across the video.
142142
143143
Parameters
144144
----------
145145
config : string
146146
Full path of the config.yaml file as a string.
147-
video : list
147+
148+
videos : list
148149
A list of strings containing the full paths of the videos to analyze.
149150
150151
shuffle: list, optional
@@ -163,7 +164,7 @@ def plot_trajectories(config,video,shuffle=1,trainingsetindex=0,videotype='.avi'
163164
Example
164165
--------
165166
for labeling the frames
166-
>>> deeplabcut.plot_trajectories('/analysis/project/reaching-task/config.yaml',['/analysis/project/videos/reachingvideo1.avi'])
167+
>>> deeplabcut.plot_trajectories('home/alex/analysis/project/reaching-task/config.yaml',['/home/alex/analysis/project/videos/reachingvideo1.avi'])
167168
--------
168169
169170
"""
@@ -173,18 +174,18 @@ def plot_trajectories(config,video,shuffle=1,trainingsetindex=0,videotype='.avi'
173174
DLCscorer = auxiliaryfunctions.GetScorerName(cfg,shuffle,trainFraction) #automatically loads corresponding model (even training iteration based on snapshot index)
174175

175176
#checks if input is a directory
176-
if [os.path.isdir(i) for i in video] == [True]:#os.path.isdir(video)==True:
177+
if [os.path.isdir(i) for i in videos] == [True]:#os.path.isdir(video)==True:
177178
"""
178179
Analyze all the videos in the directory
179180
"""
180181
print("Analyzing all the videos in the directory")
181-
videofolder= video[0]
182+
videofolder= videos[0]
182183
os.chdir(videofolder)
183-
videos = np.sort([fn for fn in os.listdir(os.curdir) if (videotype in fn) and ("labeled" not in fn)])
184+
Videos = np.sort([fn for fn in os.listdir(os.curdir) if (videotype in fn) and ("labeled" not in fn)])
184185
else:
185-
videos = video
186+
Videos = videos
186187

187-
for video in videos:
188+
for video in Videos:
188189
print(video)
189190
videofolder= str(Path(video).parents[0]) #where your folder with videos is.
190191
videotype = str(Path(video).suffix)

deeplabcut/utils/video_processor.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,11 @@ class VideoProcessor(object):
1818
Base class for a video processing unit,
1919
implementation is required for video loading and saving
2020
'''
21-
def __init__(self,fname='',sname='', nframes = -1, fps = 30):
21+
def __init__(self,fname='',sname='', nframes = -1, fps = 30,codec='X264'):
2222
self.fname = fname
2323
self.sname = sname
2424
self.nframes = nframes
25+
self.codec=codec
2526

2627
self.h = 0
2728
self.w = 0
@@ -128,10 +129,11 @@ def get_info(self):
128129
print(self.nframes)
129130

130131
def create_video(self):
131-
fourcc = cv2.VideoWriter_fourcc(*'XVID')
132+
fourcc = cv2.VideoWriter_fourcc(*self.codec)
132133
return cv2.VideoWriter(self.sname,fourcc, self.FPS, (self.w,self.h),True)
133134

134-
def _read_frame(self):
135+
def _read_frame(self): #return RGB (rather than BGR)!
136+
#return cv2.cvtColor(np.flip(self.vid.read()[1],2), cv2.COLOR_BGR2RGB)
135137
return np.flip(self.vid.read()[1],2)
136138

137139
def save_frame(self,frame):

deeplabcut/version.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,5 +7,5 @@
77
88
"""
99

10-
__version__ = '2.0.beta'
10+
__version__ = '2.0.0.beta'
1111
VERSION = __version__
24.5 KB
Loading
5.87 MB
Loading

0 commit comments

Comments
 (0)