Skip to content
This repository has been archived by the owner on Jun 19, 2020. It is now read-only.

Commit

Permalink
merge
Browse files Browse the repository at this point in the history
  • Loading branch information
Your Name committed Jan 28, 2020
2 parents 3dc0fcb + 1ea35d5 commit ec54d2b
Show file tree
Hide file tree
Showing 24 changed files with 1,336 additions and 399 deletions.
11 changes: 11 additions & 0 deletions cameras/image_converter.py
Original file line number Diff line number Diff line change
@@ -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()
71 changes: 71 additions & 0 deletions controller_listener.py
Original file line number Diff line number Diff line change
@@ -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, ())
17 changes: 14 additions & 3 deletions controls.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,13 +10,24 @@ class Controls():

def __init__(self):
self.enable_camera = True
self.enable_processing = False
self.enable_streaming = True

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()
151 changes: 77 additions & 74 deletions main.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +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_tracker2
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

Expand All @@ -26,9 +39,6 @@
from websocket import create_connection
import ujson as json

from cameras import Camera
from web.handlers import main_controller


# initiate the top level logger
logging.basicConfig(
Expand All @@ -55,12 +65,15 @@ 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)

# Set camera properties
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
Expand All @@ -70,75 +83,51 @@ def main():
cap.set(cv2.CAP_PROP_EXPOSURE, 0.02)
cap.set(cv2.CAP_PROP_CONTRAST, 0.0)

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))

color_profile_map = {}
for profile in [controls.CAMERA_MODE_RAW,
controls.CAMERA_MODE_BALL,
controls.CAMERA_MODE_HEXAGON,
controls.CAMERA_MODE_LOADING_BAY]:

# 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(1)


#websocket.enableTrace(True)
color_profile_map[profile] = ColorProfile(profile)

def update_controls(ws, message):
logger.info(message)
main_controller.color_profiles = color_profile_map

def ws_closed(ws):
logger.info('closed socket')

def on_error(ws, error):
print(error)

# tracking_ws = create_connection("wss://localhost:8080/tracking/ws/")
#

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()
#frame = filters.resize(frame, camera.FRAME_WIDTH, camera.FRAME_HEIGHT)
_, raw_frame = cap.read()

rgb_frame = cv2.cvtColor(raw_frame, cv2.COLOR_BGR2RGB)

if main_controller.camera_mode == CAMERA_MODE_RAW:

frame = frame
Expand All @@ -151,19 +140,40 @@ def start_dashboard_socket(*args):

elif main_controller.camera_mode == CAMERA_MODE_BALL:

frame, tracking_data = ball_tracker2.process(frame, generic, frame_cnt)
dashboard.putStringArray(networktables.keys.vision_target_data, tracking_data)
color_profile=main_controller.color_profiles[CAMERA_MODE_BALL]

processed_frame, tracking_data = ball_tracker.process(rgb_frame,
camera,
frame_cnt,
color_profile)

tracking_ws.send(json.dumps(dict(targets=tracking_data)))

elif main_controller.camera_mode == CAMERA_MODE_HEXAGON:

frame, tracking_data = port_tracker.process(frame, generic)
dashboard.putStringArray(networktables.keys.vision_target_data, tracking_data)
tracking_ws.send(json.dumps(dict(targets=tracking_data)))
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:

# logger.info(main_controller.camera_mode)
if main_controller.enable_streaming:
cv2.putText(frame,
calibration_frame = raw_frame.copy()

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,
Expand All @@ -172,13 +182,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 )
#cv2.waitKey(1)

else:
logger.info('waiting for control socket')
# IDLE mode
if cap.isOpened():
print('closing camera')
Expand All @@ -190,17 +204,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()
Expand Down
Loading

0 comments on commit ec54d2b

Please sign in to comment.