updated autonaming and calculate_skeleton

Author: JensBlack

DeepLabStream.py

@@ -6,22 +6,21 @@
 https://github.com/SchwarzNeuroconLab/DeepLabStream
 Licensed under GNU General Public License v3.0
 """
-import time
+import multiprocessing as mp
 import os
 import sys
-import multiprocessing as mp
+import time
 from importlib.util import find_spec
 
+import click
 import cv2
 import numpy as np
 import pandas as pd
-import click
 
-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,\
-    get_ma_pose, calculate_ma_skeletons, calculate_skeletons_dlc_live, transform_2skeleton
-from utils.plotter import plot_bodyparts, plot_metadata_frame
+from utils.configloader import RESOLUTION,FRAMERATE,OUT_DIR,MODEL_NAME,MULTI_CAM,STACK_FRAMES, \
+    ANIMALS_NUMBER,STREAMS,STREAMING_SOURCE
+from utils.plotter import plot_bodyparts,plot_metadata_frame
+from utils.poser import load_deeplabcut,get_pose,calculate_skeletons
 
 
 def create_video_files(directory, devices, resolution, framerate, codec):
@@ -284,7 +283,7 @@ def get_pose_mp(input_q, output_q):
                     index, frame = input_q.get()
                     if MODEL_ORIGIN == 'DLC':
                         scmap, locref, pose = get_pose(frame, config, sess, inputs, outputs)
-                        # TODO: REmove alterations to original
+                        # TODO: Remove alterations to original
                         #peaks = find_local_peaks_new(scmap, locref, ANIMALS_NUMBER, config)
                         peaks = pose
                     if MODEL_ORIGIN == 'MADLC':
@@ -409,9 +408,7 @@ def get_analysed_frames(self) -> tuple:
 
                     # Getting the analysed data
                     analysed_index, peaks = self._multiprocessing[camera]['output'].get()
-                    #TODO: REMOVE IF USELESS
-                    skeletons = [transform_2skeleton(peaks)]
-                    #skeletons = calculate_skeletons(peaks, ANIMALS_NUMBER)
+                    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

utils/poser.py

@@ -85,7 +85,6 @@ def find_local_peaks_new(scoremap: np.ndarray, local_reference: np.ndarray, anim
     stride = config['stride']
     # filtering scoremap
     scoremap[scoremap < 0.1] = 0
-    plot_special_peak()
     for joint_num, joint in enumerate(all_joints_names):
         all_peaks[joint] = []
         # selecting the joint in scoremap and locref
@@ -243,23 +242,30 @@ def extract_to_animal_skeleton(coords):
 """DLC LIVE & DeepPoseKit"""
 
 
-
 def transform_2skeleton(pose):
+    """Transforms pose estimation into DLStream style "skeleton" posture. If ALL_BODYPARTS is not sufficient,
+     it will autoname the bodyparts in style bp1, bp2 ..."""
     from utils.configloader import ALL_BODYPARTS
-    skeleton = dict()
-    counter = 0
-    for bp in pose:
-        skeleton[ALL_BODYPARTS[counter]] = tuple(np.array(bp[0:2],dtype = int))
-        counter += 1
-    return skeleton
+    try:
+        skeleton = dict()
+        counter = 0
+        for bp in pose:
+            skeleton[ALL_BODYPARTS[counter]] = tuple(np.array(bp[0:2],dtype = int))
+            counter += 1
+    except KeyError:
+        skeleton = dict()
+        counter = 0
+        for bp in pose:
+            skeleton[f'bp{counter}'] = tuple(np.array(bp[0:2],dtype = int))
+            counter += 1
 
+    return skeleton
 
 def transform_2pose(skeleton):
     pose = np.array([*skeleton.values()])
     return pose
 
 
-
 def calculate_skeletons_dlc_live(pose ,animals_number: int = 1) -> list:
     """
     Creating skeletons from given pose