updating model options

Author: JensBlack

DeepLabStream.py

@@ -19,7 +19,8 @@
 
 from utils.configloader import RESOLUTION, FRAMERATE, OUT_DIR, MODEL,  MULTI_CAM, STACK_FRAMES, \
     ANIMALS_NUMBER, STREAMS, VIDEO, IPWEBCAM
-from utils.poser import load_deeplabcut, get_pose, find_local_peaks_new, calculate_skeletons
+from utils.poser import load_deeplabcut, get_pose, find_local_peaks_new, calculate_skeletons,\
+    get_ma_pose, calculate_ma_skeletons, calculate_skeletons_dlc_live
 from utils.plotter import plot_bodyparts, plot_metadata_frame
 
 
@@ -264,34 +265,47 @@ def write_video(self, frames: dict, index: int):
     ######################
     @staticmethod
     def get_pose_mp(input_q, output_q):
-        from dlclive import DLCLive
+        from utils.configloader import MODEL_ORIGIN
+        from utils.poser import get_ma_pose
 
         """
         Process to be used for each camera/DLC stream of analysis
         Designed to be run in an infinite loop
         :param input_q: index and corresponding frame
         :param output_q: index and corresponding analysis
         """
-        config, sess, inputs, outputs = load_deeplabcut()
-        while True:
-            if input_q.full():
-                index, frame = input_q.get()
-                scmap, locref, pose = get_pose(frame, config, sess, inputs, outputs)
-                peaks = find_local_peaks_new(scmap, locref, ANIMALS_NUMBER, config)
-                output_q.put((index, peaks))
-
-
-        dlc_live = DLCLive(DLC_LIVE)
-        while True:
-            if input_q.full():
-                index, frame = input_q.get()
-                if not dlc_live.is_initialized:
-                    peaks = dlc_live.init_inference(frame)
-                else:
-                    peaks = dlc_live.get_pose(frame)
 
-                output_q.put((index, peaks))
+        if MODEL_ORIGIN == 'DLC' or  MODEL_ORIGIN == 'MADLC':
+            config, sess, inputs, outputs = load_deeplabcut()
+            while True:
+                if input_q.full():
+                    index, frame = input_q.get()
+                    if MODEL_ORIGIN == 'DLC':
+                        scmap, locref, pose = get_pose(frame, config, sess, inputs, outputs)
+                        peaks = find_local_peaks_new(scmap, locref, ANIMALS_NUMBER, config)
+                    if MODEL_ORIGIN == 'MADLC':
+                        peaks = get_ma_pose(frame, config, sess, inputs, outputs)
+
+                    output_q.put((index, peaks))
+
+        elif MODEL_ORIGIN == 'DLC-LIVE':
+            from dlclive import DLCLive
+            from utils.configloader import MODEL_PATH
+            dlc_live = DLCLive(MODEL_PATH)
+            while True:
+                if input_q.full():
+                    index, frame = input_q.get()
+                    if not dlc_live.is_initialized:
+                        peaks = dlc_live.init_inference(frame)
+                    else:
+                        peaks = dlc_live.get_pose(frame)
+
+                    output_q.put((index, peaks))
+        elif MODEL_ORIGIN == 'DEEPPOSEKIT':
+            print('Not here yet...')
 
+        else:
+            raise ValueError(f'Model origin {MODEL_ORIGIN} not available.')
 
     @staticmethod
     def create_mp_tools(devices):

settings.ini

@@ -4,21 +4,20 @@ FRAMERATE = 30
 STREAMS = color, depth, infrared
 OUTPUT_DIRECTORY = /Output
 MULTIPLE_DEVICES = False
-CAMERA_SOURCE = 0
+CAMERA_SOURCE = 2
 
-[DeepLabCut]
-DLC_PATH = DLC_PATH
-MODEL = MODEL_NAME
+[Pose Estimation]
+MODEL_ORIGIN = DLC
+MODEL_PATH = MODEL_PATH
+MODEL_NAME = MODEL
+; only used in DLC-LIVE for now
+ALL_BODYPARTS = bp1, bp2, bp3, bp4
 
 [Experiment]
-EXP_ORIGIN = CUSTOM/BASE
+EXP_ORIGIN = BASE/CUSTOM
 EXP_NAME = CONFIG_NAME
 RECORD_EXP = True
 
 [Video]
 VIDEO_SOURCE = PATH_TO_PRERECORDED_VIDEO
 VIDEO = False
-
-[IPWEBCAM]
-PORT = 5555
-IPWEBCAM = True

utils/configloader.py

@@ -26,14 +26,21 @@ def get_script_path():
 # DeepLabCut
 deeplabcut_config = dict(dsc_config.items('DeepLabCut'))
 
+#poseestimation
+MODEL_ORIGIN = dsc_config['Pose Estimation'].get('MODEL_ORIGIN')
+MODEL_PATH = dsc_config['Pose Estimation'].get('MODEL_PATH')
+MODEL_NAME = dsc_config['Pose Estimation'].get('MODEL_NAME')
+ALL_BODYPARTS = tuple(part for part in dsc_config['Streaming'].get('ALL_BODYPARTS').split(','))
+
+
+
 # Streaming items
 try:
     RESOLUTION = tuple(int(part) for part in dsc_config['Streaming'].get('RESOLUTION').split(','))
 except ValueError:
     print('Incorrect resolution in config!\n'
           'Using default value "RESOLUTION = 848, 480"')
     RESOLUTION = (848, 480)
-MODEL = dsc_config['Streaming'].get('MODEL')
 FRAMERATE = dsc_config['Streaming'].getint('FRAMERATE')
 OUT_DIR = dsc_config['Streaming'].get('OUTPUT_DIRECTORY')
 STREAM = dsc_config['Streaming'].getboolean('STREAM')

utils/poser.py

@@ -16,7 +16,7 @@
 from scipy.ndimage.measurements import label, maximum_position
 from scipy.ndimage.morphology import generate_binary_structure, binary_erosion
 from scipy.ndimage.filters import maximum_filter
-from utils.configloader import deeplabcut_config
+from utils.configloader import deeplabcut_config, MODEL_ORIGIN
 
 MODEL = deeplabcut_config['model']
 DLC_PATH = deeplabcut_config['dlc_path']
@@ -25,6 +25,8 @@
 try:
     import deeplabcut.pose_estimation_tensorflow.nnet.predict as predict
     from deeplabcut.pose_estimation_tensorflow.config import load_config
+    from deeplabcut.pose_estimation_tensorflow.nnet import predict_multianimal
+
     models_folder = 'pose_estimation_tensorflow/models/'
 # if not DLC 2 is not installed, try import from DLC 1 the old way
 except ImportError:
@@ -118,7 +120,7 @@ def find_local_peaks_new(scoremap: np.ndarray, local_reference: np.ndarray, anim
     return all_peaks
 
 
-def calculate_skeletons(peaks: dict, animals_number: int) -> list:
+def calculate_dlstream_skeletons(peaks: dict, animals_number: int) -> list:
     """
     Creating skeletons from given peaks
     There could be no more skeletons than animals_number
@@ -188,6 +190,61 @@ def create_animal_skeleton(dots_cluster: tuple) -> dict:
 
     return animal_skeletons
 
+"""maDLC"""
+
+def get_ma_pose(image, config, session, inputs, outputs):
+    """
+    Gets scoremap, local reference and pose from DeepLabCut using given image
+    Pose is most probable points for each joint, and not really used later
+    Scoremap and local reference is essential to extract skeletons
+    :param image: frame which would be analyzed
+    :param config, session, inputs, outputs: DeepLabCut configuration and TensorFlow variables from load_deeplabcut()
+
+    :return: tuple of scoremap, local reference and pose
+    """
+    scmap, locref, paf,  pose = predict_multianimal.get_detectionswithcosts(image, config, session, inputs, outputs, outall=True,
+                nms_radius=5.0,
+                det_min_score=0.1,
+                c_engine=False)
+    return pose
+
+def calculate_ma_skeletons(pose: dict, animals_number: int) -> list:
+    """
+    Creating skeletons from given pose in maDLC
+    There could be no more skeletons than animals_number
+    Only unique skeletons output
+    """
+
+    def extract_to_animal_skeleton(coords):
+        """
+        Creating a easy to read skeleton from dots cluster
+        Format for each joint:
+        {'joint_name': (x,y)}
+        """
+        bodyparts = np.array(coords[0])
+        skeletons = {}
+        for bp in range(len(bodyparts)):
+            for animal_num in range(animals_number):
+                if 'Mouse'+str(animal_num+1) not in skeletons.keys():
+                    skeletons['Mouse' + str(animal_num + 1)] = {}
+                if len(bodyparts[bp]) >= animals_number:
+                    skeletons['Mouse'+str(animal_num+1)]['bp' + str(bp + 1)] = bodyparts[bp][animal_num].astype(int)
+                else:
+                    if animal_num < len(bodyparts[bp]):
+                        skeletons['Mouse'+str(animal_num+1)]['bp' + str(bp + 1)] = bodyparts[bp][animal_num].astype(int)
+                    else:
+                        skeletons['Mouse'+str(animal_num+1)]['bp' + str(bp + 1)] = np.array([0,0])
+
+        return skeletons
+    animal_skeletons = extract_to_animal_skeleton(pose['coordinates'])
+    # animal_skeletons = list(animal_skeletons.values())
+
+    return animal_skeletons
+
+
+"""DLC LIVE"""
+
+
 
 def transform_2skeleton(pose):
     from utils.configloader import ALL_BODYPARTS
@@ -215,3 +272,28 @@ def calculate_skeletons_dlc_live(pose ,animals_number: int = 1) -> list:
     skeletons = [transform_2skeleton(pose)]
 
     return skeletons
+
+
+def calculate_skeletons(peaks: dict, animals_number: int) -> list:
+    """
+    Creating skeletons from given peaks
+    There could be no more skeletons than animals_number
+    Only unique skeletons output
+    adaptive to chosen model origin
+    """
+
+    if MODEL_ORIGIN == 'DLC':
+        animal_skeletons = calculate_dlstream_skeletons(peaks, animals_number)
+
+    elif MODEL_ORIGIN == 'MADLC':
+        animal_skeletons = calculate_ma_skeletons(peaks, animals_number)
+
+    elif MODEL_ORIGIN == 'DLC-LIVE':
+        animal_skeletons = calculate_skeletons_dlc_live(peaks,  animals_number= 1)
+        if animals_number != 1:
+            raise ValueError('Multiple animals are currently not supported by DLC-LIVE.'
+                             ' If you are using differently colored animals, please refere to the bodyparts directly.')
+
+
+
+