refactor websockets

cameras/image_converter.py

@@ -0,0 +1,11 @@
+import cv2
+from PIL import Image
+from io import BytesIO
+
+def convert_to_jpg(image):
+    """TBW."""
+    # assumes image is RGB provided
+    im = Image.fromarray(image)
+    mem_file = BytesIO()
+    im.save(mem_file, 'JPEG')
+    return mem_file.getvalue()

controller_listener.py

@@ -0,0 +1,71 @@
+import logging
+from controls import main_controller
+import websocket
+import json
+import _thread as thread
+
+logger = logging.getLogger(__name__)
+
+def start(websocket_url):
+
+    #websocket.enableTrace(True)
+    def update_controls(ws, message):
+        controls = json.loads(message)
+
+        logger.info(controls)
+
+        if 'request_type' in controls:
+
+            if controls['request_type'] == 'calibration':
+                main_controller.calibration = controls
+
+
+        if 'controls' in controls:
+
+            if 'camera_mode' in controls['controls']:
+                main_controller.camera_mode = controls['controls']['camera_mode']
+
+
+        if 'enable_calibration_feed' in controls:
+            main_controller.enable_calibration_feed = controls['enable_calibration_feed']
+
+        if 'enable_camera_feed' in controls:
+            main_controller.enable_camera_feed = controls['enable_camera_feed']
+
+        if 'enable_processing_feed' in controls:
+            main_controller.enable_processing_feed = controls['enable_processing_feed']
+
+        if 'color_profiles' in controls:
+            for (camera_mode, profile ) in controls['color_profiles'].items():
+                logger.info('updating %s ' % camera_mode)
+                current_profile = main_controller.color_profiles.get(camera_mode)
+                current_profile.update(profile)
+
+        if 'color_profile' in controls:
+            profile = controls['color_profile']
+            logger.info('updating %s ' % profile['camera_mode'])
+            current_profile = main_controller.color_profiles.get(profile['camera_mode'])
+            current_profile.update(profile)
+
+        #logger.info(main_controller.color_profiles)
+
+    def ws_closed(ws):
+        logger.info('closed socket')
+
+    def on_error(ws, error):
+        print(error)
+
+    def on_open(ws):
+        main_controller.enable_camera = True
+
+
+    def start_dashboard_socket(*args):
+
+        dashboard_ws = websocket.WebSocketApp(websocket_url,
+            on_message = update_controls,
+            on_close=ws_closed,
+            on_error = on_error)
+        dashboard_ws.on_open = on_open
+        dashboard_ws.run_forever()
+
+    thread.start_new_thread(start_dashboard_socket, ())

controls.py

@@ -10,14 +10,24 @@ class Controls():
 
     def __init__(self):
         self.enable_camera = True
-        self.enable_processing = False
-        self.enable_streaming = True
-        self.camera_mode = CAMERA_MODE_CALIBRATE
 
+        self.enable_camera_feed = False
+        self.enable_calibration_feed = False
+        self.enable_processing_feed = True
+
+
+        self.camera_mode = CAMERA_MODE_BALL
         self.enable_feed = True
-        self.turn_camera_off = False
+        self.color_profiles = {}
+
+
+        self.calibration = {}
+
 
     def connect(self):
         controller_listener.connect(self)
 
+    def update(message):
+        print(message)
+
 main_controller = Controls()

main.py

@@ -6,12 +6,24 @@
 
 from processing import colors
 import network as networktables
+
 from cameras import logitech_c270, generic
+from cameras import Camera
+from cameras import image_converter
+
 from profiles import color_profiles
 from processing import bay_tracker
 from processing import port_tracker
 from processing import ball_tracker
+from processing import color_calibrate
+
+
+import controls
 from controls import main_controller
+import controller_listener
+
+from profiles.color_profile import ColorProfile
+
 import _thread as thread
 import time
 
@@ -27,7 +39,6 @@
 from websocket import create_connection
 import ujson as json
 
-from cameras import Camera
 
 # initiate the top level logger
 logging.basicConfig(
@@ -50,84 +61,61 @@ def main():
 
     cap = cv2.VideoCapture(config.video_source_number)
 
-    # out_pipeline = gst_utils.get_udp_streamer_pipeline2(config.gstreamer_client_ip,
-    #                                          config.gstreamer_client_port,
-    #                                          config.gstreamer_bitrate)
+    enable_gstreamer_pipeline = False
+
+    out = None
+    if enable_gstreamer_pipeline:
+        out_pipeline = gst_utils.get_udp_sender(config.gstreamer_client_ip, config.gstreamer_client_port)
 
-    out_pipeline = gst_utils.get_udp_sender(config.gstreamer_client_ip,
-                                            config.gstreamer_client_port)
+        # out_pipeline = gst_utils.get_udp_streamer_pipeline2(config.gstreamer_client_ip,
+        #                                          config.gstreamer_client_port,
+        #                                          config.gstreamer_bitrate)
 
 
+        out = cv2.VideoWriter(out_pipeline, 0,
+                              camera.FPS,
+                              (camera.FRAME_WIDTH, camera.FRAME_HEIGHT),
+                              True)
+
     # Set camera properties
     camera = Camera(cap.get(cv2.CAP_PROP_FRAME_WIDTH),
                     cap.get(cv2.CAP_PROP_FRAME_HEIGHT),
                     cap.get(cv2.CAP_PROP_FPS))
 
-    # print([camera.FRAME_WIDTH])
-    # print([camera.FRAME_HEIGHT])
-    # print([camera.FPS])
-
-    out = cv2.VideoWriter(out_pipeline, 0,
-                          camera.FPS,
-                          (camera.FRAME_WIDTH, camera.FRAME_HEIGHT),
-                          True)
-
-    #TODO: if no camera, exit and msg no camera
-    time.sleep(5)
-
-
-    #websocket.enableTrace(True)
-
-    def update_controls(ws, message):
-        logger.info(message)
-
-    def ws_closed(ws):
-        logger.info('closed socket')
+    color_profile_map = {}
+    for profile in [controls.CAMERA_MODE_RAW,
+                    controls.CAMERA_MODE_BALL,
+                    controls.CAMERA_MODE_HEXAGON,
+                    controls.CAMERA_MODE_LOADING_BAY]:
 
-    def on_error(ws, error):
-        print(error)
+        color_profile_map[profile] = ColorProfile(profile)
 
-    # tracking_ws = create_connection("wss://localhost:8080/tracking/ws/")
-    #
+    main_controller.color_profiles = color_profile_map
 
-    def on_open(ws):
-        def run(*args):
-            for i in range(3):
-                time.sleep(1)
-                ws.send("Hello %d" % i)
-            time.sleep(1)
-            ws.close()
-            print("thread terminating...")
-        thread.start_new_thread(run, ())
+    time.sleep(5)
 
     tracking_ws = create_connection("ws://localhost:8080/tracking/ws")
+    camera_ws = create_connection("ws://localhost:8080/camera/ws")
+    processed_ws = create_connection("ws://localhost:8080/processed/ws")
+    calibration_ws = create_connection("ws://localhost:8080/calibration/ws")
 
-    def start_dashboard_socket(*args):
-        dashboard_ws = websocket.WebSocketApp("ws://localhost:8080/dashboard/ws",
-            on_message = update_controls,
-            on_close=ws_closed,
-            on_error = on_error)
-        dashboard.on_open = on_open
-        dashboard_ws.run_forever()
-
-    thread.start_new_thread(start_dashboard_socket, ())
+    controller_listener.start("ws://localhost:8080/dashboard/ws")
 
     logger.info('starting main loop ')
-
     frame_cnt = 0
     while(True):
 
         frame_cnt += 1
 
-        if True or main_controller.enable_camera:
+        if main_controller.enable_camera:
 
             if not cap.isOpened():
                 print('opening camera')
                 cap.open(config.video_source_number)
 
-            _, frame = cap.read()
+            _, raw_frame = cap.read()
 
-            #frame = filters.resize(frame, camera.FRAME_WIDTH, camera.FRAME_HEIGHT)
+            rgb_frame = cv2.cvtColor(raw_frame, cv2.COLOR_BGR2RGB)
 
             if main_controller.camera_mode == CAMERA_MODE_RAW:
 
@@ -141,19 +129,40 @@ def start_dashboard_socket(*args):
 
             elif main_controller.camera_mode == CAMERA_MODE_BALL:
 
-                frame, tracking_data = ball_tracker.process(frame,
+                color_profile=main_controller.color_profiles[CAMERA_MODE_BALL]
+
+                processed_frame, tracking_data = ball_tracker.process(rgb_frame,
                                                             camera,
-                                                            frame_cnt)
+                                                            frame_cnt,
+                                                            color_profile)
 
                 tracking_ws.send(json.dumps(dict(targets=tracking_data)))
 
             elif main_controller.camera_mode == CAMERA_MODE_HEXAGON:
 
-                frame = port_tracker.process(frame, generic, color_profiles.ReflectiveProfile())
+                processed_frame = port_tracker.process(frame, generic, color_profiles.ReflectiveProfile())
+
+
+            if main_controller.enable_camera_feed:
+
+                jpg=image_converter.convert_to_jpg(rgb_frame)
+                camera_ws.send_binary(jpg)
+
+            if main_controller.enable_calibration_feed:
 
+                calibration_frame = raw_frame.copy()
 
-            if main_controller.enable_streaming:
-                cv2.putText(frame,
+                calibration_frame = color_calibrate.process(calibration_frame,
+                                                            camera_mode = main_controller.calibration.get('camera_mode', 'RAW'),
+                                                            color_mode = main_controller.calibration.get('color_mode'),
+                                                            apply_mask = main_controller.calibration.get('apply_mask', False))
+
+                jpg=image_converter.convert_to_jpg(calibration_frame)
+                calibration_ws.send_binary(jpg)
+
+            if main_controller.enable_processing_feed:
+
+                cv2.putText(processed_frame,
                             'Tracking Mode %s' % main_controller.camera_mode,
                             (10,10),
                             cv2.FONT_HERSHEY_DUPLEX,
@@ -162,13 +171,17 @@ def start_dashboard_socket(*args):
                             1,
                             cv2.LINE_AA)
 
+                jpg=image_converter.convert_to_jpg(processed_frame)
+                processed_ws.send_binary(jpg)
 
-                out.write(frame)
+                # if out is not None:
+                #     out.write(frame)
 
             #cv2.imshow('frame', frame )
             #v2.waitKey(1)
 
         else:
+            logger.info('waiting for control socket')
             # IDLE mode
             #if cap.isOpened():
                 #print('closing camera')
@@ -180,17 +193,6 @@ def start_dashboard_socket(*args):
 
 
 
-
-
-def single_frame(debug=False):
-
-    img = cv2.imread("frc_cube.jpg")
-    img = cube_tracker.process(img,
-                               generic)
-
-    cv2.imshow('Objects Detected',img)
-    cv2.waitKey()
-
 if __name__ == '__main__':
     p = Process(target=start_web.main)
     p.start()