Merge pull request #2 from SchwarzNeuroconLab/fix/various-clean-up-fixes

Cleaned up the imports

Author: alexfrostdesu

DeepLabStream.py

@@ -17,10 +17,12 @@
 import numpy as np
 import pandas as pd
 
-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
+from utils.generic import VideoManager, WebCamManager, GenericManager
+from utils.configloader import RESOLUTION, FRAMERATE, OUT_DIR, MODEL_NAME, MULTI_CAM, STACK_FRAMES, \
+    ANIMALS_NUMBER, STREAMS, STREAMING_SOURCE, MODEL_ORIGIN
+from utils.plotter import plot_bodyparts, plot_metadata_frame
+from utils.poser import load_deeplabcut, load_dpk, load_dlc_live, get_pose, calculate_skeletons,\
+    find_local_peaks_new, get_ma_pose
 
 
 def create_video_files(directory, devices, resolution, framerate, codec):
@@ -127,28 +129,19 @@ def set_camera_manager():
         :return: the chosen camera manager
         """
 
-        if STREAMING_SOURCE.lower() == 'video':
-            from utils.generic import VideoManager
-            manager = VideoManager()
-            return manager
-
-        elif STREAMING_SOURCE.lower() == 'ipwebcam':
-            from utils.generic import WebCamManager
-            manager = WebCamManager()
-            return manager
-
-        elif STREAMING_SOURCE.lower() == 'camera':
+        def select_camera_manager():
+            """
+            Function to select from all available camera managers
+            """
             manager_list = []
             # loading realsense manager, if installed
-            realsense = find_spec("pyrealsense2") is not None
-            if realsense:
+            if find_spec("pyrealsense2") is not None:
                 from utils.realsense import RealSenseManager
                 realsense_manager = RealSenseManager()
                 manager_list.append(realsense_manager)
 
             # loading basler manager, if installed
-            pylon = find_spec("pypylon") is not None
-            if pylon:
+            if find_spec("pypylon") is not None:
                 from utils.pylon import PylonManager
                 pylon_manager = PylonManager()
                 manager_list.append(pylon_manager)
@@ -170,9 +163,19 @@ def check_for_cameras(camera_manager):
                     return manager
             else:
                 # if no camera is found, try generic openCV manager
-                from utils.generic import GenericManager
                 generic_manager = GenericManager()
                 return generic_manager
+
+        MANAGER_SOURCE = {
+            'video': VideoManager,
+            'ipwebcam': WebCamManager,
+            'camera': select_camera_manager
+        }
+
+        # initialize selected manager
+        camera_manager = MANAGER_SOURCE.get(STREAMING_SOURCE)()
+        if camera_manager is not None:
+            return camera_manager
         else:
             raise ValueError(f'Streaming source {STREAMING_SOURCE} is not a valid option. \n'
                              f'Please choose from "video", "camera" or "ipwebcam".')
@@ -266,35 +269,30 @@ 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
         """
 
-        if MODEL_ORIGIN == 'DLC' or  MODEL_ORIGIN == 'MADLC':
-            from utils.poser import find_local_peaks_new
+        if MODEL_ORIGIN in ('DLC', '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
+                        # 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)
+            dlc_live = load_dlc_live()
             while True:
                 if input_q.full():
                     index, frame = input_q.get()
@@ -304,21 +302,17 @@ def get_pose_mp(input_q, output_q):
                         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
+            predict_model = load_dpk()
             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]
+                    peaks = prediction[0, :, :2]
                     output_q.put((index, peaks))
-
-
         else:
             raise ValueError(f'Model origin {MODEL_ORIGIN} not available.')
 
@@ -763,10 +757,14 @@ def show_benchmark_statistics():
                 print("[{0}/3000] Benchmarking in progress".format(len(analysis_time_data)))
 
     if benchmark_enabled:
-        import re
-        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]))
+        model_parts = MODEL_NAME.split('_')
+        if len(model_parts) == 3:
+            short_model = model_parts[0] + '_' + model_parts[2]
+        else:
+            short_model = MODEL_NAME
+        # the best way to save files
+        np.savetxt(f'{OUT_DIR}/{short_model}_framerate_{FRAMERATE}_resolution_{RESOLUTION[0]}_{RESOLUTION[1]}.txt',
+                   np.transpose([fps_data, whole_loop_time_data]))
 
 
 if __name__ == '__main__':

requirements.txt

@@ -1,11 +1,13 @@
-gpiozero
-pigpio
-pyserial
-nidaqmx>=0.5.7
-click>=7.0
-opencv-python>=3.4.5.20
+gpiozero==1.5.1
+pigpio==1.78
+pyserial==3.5
+nidaqmx==0.5.7
+click==7.1.2
+opencv-python==3.4.5.20
+opencv-contrib-python==4.4.0.46
 numpy>=1.14.5
-pandas>=0.21.0
-matplotlib>=3.0.3
-scikit-image>=0.14.2
-scipy>=1.1.0
+pandas==1.1.4
+matplotlib==3.0.3
+scikit-image==0.17.2
+scipy==1.5.4
+pyzmq==20.0.0

utils/generic.py

@@ -5,11 +5,15 @@
 https://github.com/SchwarzNeuroconLab/DeepLabStream
 Licensed under GNU General Public License v3.0
 """
+import time
+import base64
 
 import cv2
-from utils.configloader import CAMERA_SOURCE, VIDEO_SOURCE, RESOLUTION, FRAMERATE, PORT
-import time
 import numpy as np
+import zmq
+
+from utils.configloader import CAMERA_SOURCE, VIDEO_SOURCE, RESOLUTION, FRAMERATE, PORT
+
 
 class GenericManager:
     """
@@ -86,7 +90,6 @@ def get_name(self) -> str:
         return self._manager_name
 
 
-
 class VideoManager(GenericManager):
 
     """
@@ -97,13 +100,12 @@ def __init__(self):
         Generic video manager from video files
         Uses pure opencv
         """
-        self._manager_name = "generic"
+        super().__init__()
         self._camera = cv2.VideoCapture(VIDEO_SOURCE)
         self._camera_name = "Video"
         self.initial_wait = False
         self.last_frame_time = time.time()
 
-
     def get_frames(self) -> tuple:
         """
         Collect frames for camera and outputs it in 'color' dictionary
@@ -116,7 +118,6 @@ def get_frames(self) -> tuple:
         infra_frames = {}
         ret, image = self._camera.read()
         self.last_frame_time = time.time()
-        #print(ret)
         if ret:
             if not self.initial_wait:
                 cv2.waitKey(1000)
@@ -127,6 +128,10 @@ def get_frames(self) -> tuple:
             if running_time <= 1 / FRAMERATE:
                 sleepy_time = int(np.ceil(1000/FRAMERATE - running_time / 1000))
                 cv2.waitKey(sleepy_time)
+        else:
+            # cycle the video for testing purposes
+            self._camera.set(cv2.CAP_PROP_POS_FRAMES, 0)
+            return self.get_frames()
 
         return color_frames, depth_maps, infra_frames
 
@@ -138,24 +143,23 @@ def __init__(self):
         Binds the computer to a ip address and starts listening for incoming streams.
         Adapted from StreamViewer.py https://github.com/CT83/SmoothStream
         """
-        import zmq
+        super().__init__()
         self._context = zmq.Context()
         self._footage_socket = self._context.socket(zmq.SUB)
         self._footage_socket.bind('tcp://*:' + PORT)
         self._footage_socket.setsockopt_string(zmq.SUBSCRIBE, np.unicode(''))
 
-        self._manager_name = "generic"
         self._camera = None
         self._camera_name = "webcam"
         self.initial_wait = False
         self.last_frame_time = time.time()
 
-    def string_to_image(self, string):
-        """ Taken from https://github.com/CT83/SmoothStream"""
+    @ staticmethod
+    def string_to_image(string):
+        """
+        Taken from https://github.com/CT83/SmoothStream
+        """
 
-        import numpy as np
-        import cv2
-        import base64
         img = base64.b64decode(string)
         npimg = np.fromstring(img, dtype=np.uint8)
         return cv2.imdecode(npimg, 1)