Merge branch 'dev_models'

Author: JensBlack

DeepLabStream.py

@@ -17,9 +17,10 @@
 import pandas as pd
 import click
 
-from utils.configloader import RESOLUTION, FRAMERATE, OUT_DIR, MODEL,  MULTI_CAM, STACK_FRAMES, \
+from utils.configloader import RESOLUTION, FRAMERATE, OUT_DIR, MODEL_NAME,  MULTI_CAM, STACK_FRAMES, \
     ANIMALS_NUMBER, STREAMS, STREAMING_SOURCE
-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, transform_2skeleton
 from utils.plotter import plot_bodyparts, plot_metadata_frame
 
 
@@ -266,19 +267,61 @@ def write_video(self, frames: dict, index: int):
     ######################
     @staticmethod
     def get_pose_mp(input_q, output_q):
+        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))
+
+        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)
+                        # TODO: REmove alterations to original
+                        #peaks = find_local_peaks_new(scmap, locref, ANIMALS_NUMBER, config)
+                        peaks = pose
+                    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':
+            from deepposekit.models import load_model
+            from utils.configloader import MODEL_PATH
+            model = load_model(MODEL_PATH)
+            predict_model = model.predict_model
+            while True:
+                if input_q.full():
+                    index, frame = input_q.get()
+                    frame = frame[..., 1][..., None]
+                    st_frame = np.stack([frame])
+                    prediction = predict_model.predict(st_frame, batch_size=1, verbose=True)
+                    peaks= prediction[0,:,:2]
+                    output_q.put((index, peaks))
+
+
+        else:
+            raise ValueError(f'Model origin {MODEL_ORIGIN} not available.')
 
     @staticmethod
     def create_mp_tools(devices):
@@ -366,7 +409,9 @@ def get_analysed_frames(self) -> tuple:
 
                     # Getting the analysed data
                     analysed_index, peaks = self._multiprocessing[camera]['output'].get()
-                    skeletons = calculate_skeletons(peaks, ANIMALS_NUMBER)
+                    #TODO: REMOVE IF USELESS
+                    skeletons = [transform_2skeleton(peaks)]
+                    #skeletons = calculate_skeletons(peaks, ANIMALS_NUMBER)
                     print('', end='\r', flush=True)  # this is the line you should not remove
                     analysed_frame, depth_map, input_time = self.get_stored_frames(camera)
                     analysis_time = time.time() - input_time
@@ -722,7 +767,7 @@ def show_benchmark_statistics():
 
     if benchmark_enabled:
         import re
-        short_model = re.split('[-_]', MODEL)
+        short_model = re.split('[-_]', MODEL_NAME)
         short_model = short_model[0] + '_' + short_model[2]
         np.savetxt(f'{OUT_DIR}/{short_model}_framerate_{FRAMERATE}_resolution_{RESOLUTION[0]}_{RESOLUTION[1]}.txt', np.transpose([fps_data, whole_loop_time_data]))
 

experiments/custom/triggers.py

@@ -11,6 +11,9 @@
 from utils.configloader import RESOLUTION
 from collections import deque
 import numpy as np
+
+import time
+
 """Single posture triggers"""
 
 class HeaddirectionROITrigger:
@@ -440,4 +443,4 @@ def check_skeleton(self, skeleton: dict):
                                                color=color)}}
         response = (result, response_body)
 
-        return response
+        return response

requirements.txt

@@ -1,3 +1,4 @@
+numba
 gpiozero
 pigpio
 pyserial

settings.ini

@@ -1,28 +1,32 @@
 [Streaming]
-RESOLUTION = 848, 480
+RESOLUTION = 1920, 1080
 FRAMERATE = 30
 OUTPUT_DIRECTORY = /Output
 #if you have connected multiple cameras (USB), you will need to select the number OpenCV has given them.
 #Default is "0", which takes the first available camera.
 CAMERA_SOURCE = 0
-#you can use camera, ipwebcam or video to select your input source
-STREAMING_SOURCE = camera
+#you can use "camera", "ipwebcam" or "video" to select your input source
+STREAMING_SOURCE = video
 
-[DeepLabCut]
-DLC_PATH = DLC_PATH
-MODEL = MODEL_NAME
+[Pose Estimation]
+#possible origins are: DLC, DLC-LIVE, MADLC, DEEPPOSEKIT
+MODEL_ORIGIN = ORIGIN
+MODEL_PATH = PATH_TO_MODEL
+MODEL_NAME = MODEL_NAME
+; only used in DLC-LIVE and DeepPoseKit for now; if left empty or to short,  auto-naming will be enabled in style bp1, bp2 ...
+ALL_BODYPARTS = bp1, bp2, bp3, bp4
 
 [Experiment]
 #Available parameters are "CUSTOM" and "BASE"
-EXP_ORIGIN = BASE
+EXP_ORIGIN = CUSTOM
 #Name of the experiment config in /experiments/configs or name of the custom experiment in /experiments/custom/experiments.py
-EXP_NAME = CONFIG_NAME
+EXP_NAME = ExampleExperiment
 #if you want the experiment to be recorded as a raw video set this to "True".
 RECORD_EXP = True
 
 [Video]
 #Full path to video that you want to use as input. Needs "STREAMING_SOURCE" set to "video"!
-VIDEO_SOURCE = PATH_TO_PRERECORDED_VIDEO
+VIDEO_SOURCE = PATH_TO_VIDEO
 
 [IPWEBCAM]
 #Standard Port is 5555 if you followed the SmoothStream setup

utils/VideoAnalyzer.py

@@ -18,6 +18,8 @@
 from utils.configloader import VIDEO_SOURCE, OUT_DIR, ANIMALS_NUMBER
 from experiments.custom.experiments import ExampleExperiment
 
+from deepposekit.models import load_model
+
 
 def create_dataframes(data_output):
     """
@@ -43,7 +45,7 @@ def create_dataframes(data_output):
 
 def start_videoanalyser():
     print("Starting DeepLabCut")
-    config, sess, inputs, outputs = load_deeplabcut()
+    model = load_model(r"D:\DeepPoseKit-Data-master\datasets\fly\best_model_densenet.h5")
 
     experiment_enabled = False
     video_output = True

utils/configloader.py

@@ -27,8 +27,12 @@ def get_script_path():
 adv_cfg_path = os.path.join(os.path.dirname(__file__), 'advanced_settings.ini')
 with open(adv_cfg_path) as adv_cfg_file:
     adv_dsc_config.read_file(adv_cfg_file)
-# 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['Pose Estimation'].get('ALL_BODYPARTS').split(','))
 
 # Streaming items
 try:
@@ -37,7 +41,7 @@ def get_script_path():
     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')
 CAMERA_SOURCE = dsc_config['Streaming'].get('CAMERA_SOURCE')
@@ -56,7 +60,6 @@ def get_script_path():
 
 START_TIME = time.time()
 
-
 """advanced settings"""
 STREAMS = [str(part).strip() for part in adv_dsc_config['Streaming'].get('STREAMS').split(',')]
 MULTI_CAM = adv_dsc_config['Streaming'].getboolean('MULTIPLE_DEVICES')

utils/plotter.py

@@ -7,7 +7,6 @@
 """
 
 import cv2
-import random
 
 
 def plot_dots(image, coordinates, color, cond=False):
@@ -55,6 +54,16 @@ def plot_metadata_frame(image, frame_width, frame_height, current_fps, current_e
                 (int(frame_width * 0.8), int(frame_height * 0.95)), font, 1, (255, 255, 0))
     return res_image
 
+def plot_dlc_bodyparts(image, bodyparts):
+    """
+    Plots dlc bodyparts on given image
+    adapted from plotter
+    """
+
+    for bp in bodyparts:
+        center = tuple(bp.astype(int))
+        cv2.circle(image, center=center, radius=3, color=(255, 0, 0), thickness=2)
+    return image
 
 def plot_triggers_response(image, response):
     """
@@ -74,6 +83,4 @@ def plot_triggers_response(image, response):
             cv2.rectangle(image, **plot['square'], thickness=2)
 
 
-def plot_special_peak():
-    if random.randint(0, 100000) == 1:
-        print(EGG)
+