def __init__(self):
self.mp_array = SyncableMPArray((GRADIENT_HEIGHT, *MAP_DIMS[1:]))
# Constants
self.label_y = 15
self.label_xmin = (3, self.label_y)
self.label_xmax_shift = MAP_DIMS[1] - 10
# Main thread vars
self.last_min, self.last_max = -1, -1
def __init__(self):
self.mp_array = SyncableMPArray(MAP_DIMS)
# Constants
self.num_rows = 12
self.num_cols = 16
self.row_scale = int(MAP_DIMS[0] / self.num_rows)
self.col_scale = int(MAP_DIMS[1] / self.num_cols)
# Main thread vars
self.bins = np.zeros((self.num_rows, self.num_cols), dtype='uint32')
self.empty = np.zeros((self.num_rows, self.num_cols), dtype='uint8')
def __init__(self, initial_duration):
self.mp_array = SyncableMPArray((PROGBAR_HEIGHT, *VID_DIM_RGB[1:]))
# -- Constants -- #
# Total segments in progress bar (= horizontal length)
self.num_steps = VID_DIM_RGB[1]
# Text Constants
self.font = cv2.FONT_HERSHEY_SIMPLEX
self.text_vloc = 15
self.text_size = 0.45
self.text_thickness = 1
self.text_dims = cv2.getTextSize('00:00.000', self.font,
self.text_size,
self.text_thickness)[0]
spacing = self.text_dims[0] // 2
self.txt_left_lmt = spacing
self.txt_right_lmt = self.num_steps - spacing
# -- Main thread vars -- #
# Operation Params
self.curr_loc = -1
self.text_hloc = 0
self.start_time = None
self.output_array = None
self.image = None
self.displaying_error_image = False
self.targ_perim = CV2TargetAreaPerimeter()
# Progress bar segments for each element
self.pbar_slice = None
self.mouse_in_targ_slice = None
self.mouse_stim_slice = None
# Mouse Status
self.mouse_in_target = False # is mouse inside target region?
self.mouse_recv_stim = False # does mouse receive stimulation?
self.mouse_stim_timer = None # timer to make sure mouse receives STIM_ON secs stim, max every STIM_TOTAL secs
self.in_targ_stopwatch = StopWatch(
) # total time spent in target region
self.get_stim_stopwatch = StopWatch(
) # total time spent receiving stimulation
self.mouse_n_entries = 0 # num entries into target region
self.mouse_n_stims = 0 # num stimulations received
# -- Modifier vars (read-only for main thread) -- #
# Operation Params
self._running = False
self._duration = initial_duration
class Pathing(object):
"""Generates pathing map from coordinates"""
def __init__(self):
self.mp_array = SyncableMPArray(MAP_DIMS)
# Main thread vars
self.last_coord = None
# Initializing functions. Call once once new process starts
def init_unpickleable_objs(self):
"""These objects must be created in the process they will run in"""
self.output_array = self.mp_array.generate_np_array()
self.output_array.set_can_recv_img()
# Modifier Functions. Can call from other threads
# *** Non-underscored variables are READ ONLY
def reset(self):
self.last_coord = None
self.output_array.fill(0)
# Main Update Function. Run in Main Thread. Do NOT call from any other thread
# *** Underscored variables are READ ONLY
def update(self, coord):
"""Draw new pathing segment on pathing array"""
col, row = coord
if col is not None and row is not None:
# Scale coords
if MAP_DOWNSCALE > 1:
coord = round(col / MAP_DOWNSCALE), round(row / MAP_DOWNSCALE)
# Draw new path segment
if self.last_coord:
cv2.line(self.output_array, coord, self.last_coord,
(0, 255, 0), 1)
self.last_coord = coord
# Generate Output from List of Coords
@staticmethod
def get_pathmap(coord_list):
"""Provided a full list of coords, generate a full size map"""
last_path_coord = None
pathmap = np.zeros(VID_DIM_RGB, dtype='uint8')
for coord in coord_list:
if coord != (None, None):
if last_path_coord:
cv2.line(pathmap, coord, last_path_coord, (0, 255, 0), 1)
last_path_coord = coord
return pathmap
class Gradient(object):
"""Generates a gradient with variable labels from coordinates"""
def __init__(self):
self.mp_array = SyncableMPArray((GRADIENT_HEIGHT, *MAP_DIMS[1:]))
# Constants
self.label_y = 15
self.label_xmin = (3, self.label_y)
self.label_xmax_shift = MAP_DIMS[1] - 10
# Main thread vars
self.last_min, self.last_max = -1, -1
# Initializing functions. Call once once new process starts
def init_unpickleable_objs(self):
"""These objects must be created in the process they will run in"""
self.output_array = self.mp_array.generate_np_array()
self.init_gradient()
self.output_array.set_can_recv_img()
def init_gradient(self):
"""Create gradient indicator"""
num_grads = 32
raw = np.zeros((1, num_grads))
for col in range(num_grads):
raw[:, col] = col
empty = np.zeros((1, num_grads), dtype='uint8')
# Create red and green channels.
red = raw.copy()
green = raw.copy()
# Generate gradient. we use a black-yellow-red gradient.
red = (red / red.max()) * 2 * 255
red = np.clip(red, 0, 255)
green = (green / green.max()) * 2 * 255
green[green > 255] = 255 - (green[green > 255] - 255)
# Retype into 8 bits
red = red.astype('uint8')
green = green.astype('uint8')
# Create image
image = np.dstack((red, green, empty))
col_scale = int(MAP_DIMS[1] / num_grads)
self.gradient = np.kron(
image, np.ones((GRADIENT_HEIGHT, col_scale, 1), dtype='uint8'))
self.text_slice = self.gradient[(GRADIENT_HEIGHT // 2 -
10):(GRADIENT_HEIGHT // 2 + 10), :, :]
# Send gradient image
self.output_array.send_img(self.gradient)
# Modifier Functions. Can call from other threads
# *** Non-underscored variables are READ ONLY
def reset(self):
self.last_min, self.last_max = -1, -1
self.update(0, 0)
# Main Update Function. Run in Main Thread. Do NOT call from any other thread
# *** Underscored variables are READ ONLY
def update(self, minimum, maximum):
"""Update scale on gradient"""
if minimum == self.last_min and maximum == self.last_max:
return
# Reset text
self.text_slice[:] = self.gradient[:20, :, :]
# Find x location of maximum label
label_xmax = self.label_xmax_shift - (len(str(maximum)) * 5)
label_xmax = (label_xmax, self.label_y)
# Label gradient
cv2.putText(self.text_slice, str(minimum), self.label_xmin,
cv2.FONT_HERSHEY_SIMPLEX, 0.35, (255, 255, 255), 1)
cv2.putText(self.text_slice, str(maximum), label_xmax,
cv2.FONT_HERSHEY_SIMPLEX, 0.35, (0, 0, 0), 1)
# Send to shared array
self.output_array.send_img(self.gradient)
# Remember min and max
self.last_min, self.last_max = minimum, maximum
# Append a gradient to a heatmap, given min and max
@staticmethod
def append_gradient(minimum, maximum, heatmap):
"""Append gradient"""
num_grads = 64
raw = np.zeros((1, num_grads))
for col in range(num_grads):
raw[:, col] = col
empty = np.zeros((1, num_grads), dtype='uint8')
# Create red and green channels.
red = raw.copy()
green = raw.copy()
# Generate gradient. we use a black-yellow-red gradient.
red = (red / red.max()) * 2 * 255
red = np.clip(red, 0, 255)
green = (green / green.max()) * 2 * 255
green[green > 255] = 255 - (green[green > 255] - 255)
# Retype into 8 bits
red = red.astype('uint8')
green = green.astype('uint8')
# Create gradient (cv2 imwrite takes BGR)
image = np.dstack((empty, green, red))
col_scale = int(VID_DIM_RGB[1] / num_grads)
gradient = np.kron(
image, np.ones((GRADIENT_HEIGHT // 2, col_scale, 1),
dtype='uint8'))
# Add min max labels
y = 30
xmax = VID_DIM_RGB[1] - 8 - len(str(maximum) * 10), y
xmin = 3, y
# Label gradient
cv2.putText(gradient, str(minimum), xmin, cv2.FONT_HERSHEY_SIMPLEX,
0.5, (255, 255, 255), 1)
cv2.putText(gradient, str(maximum), xmax, cv2.FONT_HERSHEY_SIMPLEX,
0.5, (0, 0, 0), 1)
# Append to heatmap
shape = heatmap.shape[0] + GRADIENT_HEIGHT // 2, heatmap.shape[1], 3
output = np.zeros(shape, dtype='uint8')
output[:heatmap.shape[0], :, :] = heatmap
output[heatmap.shape[0]:, :, :] = gradient
return output
def __init__(self):
self.mp_array = SyncableMPArray(MAP_DIMS)
# Main thread vars
self.last_coord = None
class Heatmap(object):
"""Generates heatmap from coordinates"""
def __init__(self):
self.mp_array = SyncableMPArray(MAP_DIMS)
# Constants
self.num_rows = 12
self.num_cols = 16
self.row_scale = int(MAP_DIMS[0] / self.num_rows)
self.col_scale = int(MAP_DIMS[1] / self.num_cols)
# Main thread vars
self.bins = np.zeros((self.num_rows, self.num_cols), dtype='uint32')
self.empty = np.zeros((self.num_rows, self.num_cols), dtype='uint8')
# Initializing functions. Call once once new process starts
def init_unpickleable_objs(self):
"""These objects must be created in the process they will run in"""
self.output_array = self.mp_array.generate_np_array()
self.output_array.set_can_recv_img()
# Modifier Functions. Can call from other threads
# *** Non-underscored variables are READ ONLY
def reset(self):
self.bins.fill(0)
self.output_array.fill(0)
# Main Update Function. Run in Main Thread. Do NOT call from any other thread
# *** Underscored variables are READ ONLY
def update(self, coord):
"""Update heatmap with supplied coord"""
col, row = coord
# Find bins this coord belongs to, and add to bin
if row is not None and col is not None:
rowbin = int(row / (self.row_scale * MAP_DOWNSCALE))
colbin = int(col / (self.col_scale * MAP_DOWNSCALE))
self.bins[rowbin, colbin] += 1
# We use a black-yellow-red gradient.
red = self.bins.copy()
green = self.bins.copy()
# Create Gradient
red = (red / red.max()) * 2 * 255
red = np.clip(red, 0, 255)
green = (green / green.max()) * 2 * 255
green[green > 255] = 255 - (green[green > 255] - 255)
# Retype into 8 bits
red = red.astype('uint8')
green = green.astype('uint8')
# Create Image; blue is empty.
bins = np.dstack((red, green, self.empty))
heatmap = np.kron(
bins,
np.ones((self.row_scale, self.col_scale, 1), dtype='uint8'))
self.output_array.send_img(heatmap)
# Update Gradient
return self.bins.min(), self.bins.max()
# Generate Output from List of Coords
def get_heatmap(self, coord_list):
"""Provided a full list of coords, generate a full size map"""
bins = np.zeros((self.num_rows, self.num_cols), dtype='uint32')
empty = np.zeros((self.num_rows, self.num_cols), dtype='uint8')
row_scale = int(VID_DIM_RGB[0] / self.num_rows)
col_scale = int(VID_DIM_RGB[1] / self.num_cols)
for col, row in coord_list:
if row is not None and col is not None:
rowbin = int(row / row_scale)
colbin = int(col / col_scale)
bins[rowbin, colbin] += 1
red = bins.copy()
green = bins.copy()
# Create Gradient (black yellow red)
if bins.max() > 0:
red = (red / red.max()) * 2 * 255
red = np.clip(red, 0, 255)
green = (green / green.max()) * 2 * 255
green[green > 255] = 255 - (green[green > 255] - 255)
# Retype into 8 bits
red = red.astype('uint8')
green = green.astype('uint8')
# Create BGR Image (cv2 imwrite takes BGR)
stacked = np.dstack((empty, green, red))
heatmap = np.kron(stacked,
np.ones((row_scale, col_scale, 1), dtype='uint8'))
# Add bin text
for row in range(self.num_rows):
for col in range(self.num_cols):
num = int(bins[row, col])
if num < bins.max() / 3:
color = (255, 255, 255)
else:
color = (0, 0, 0)
cv2.putText(heatmap, str(num),
(col * col_scale + 5, row * row_scale + 25),
cv2.FONT_HERSHEY_SIMPLEX, 0.35, color, 1)
return bins.min(), bins.max(), heatmap
class ProgressBar(object):
"""Numpy Array based progress bar"""
def __init__(self, initial_duration):
self.mp_array = SyncableMPArray((PROGBAR_HEIGHT, *VID_DIM_RGB[1:]))
# -- Constants -- #
# Total segments in progress bar (= horizontal length)
self.num_steps = VID_DIM_RGB[1]
# Text Constants
self.font = cv2.FONT_HERSHEY_SIMPLEX
self.text_vloc = 15
self.text_size = 0.45
self.text_thickness = 1
self.text_dims = cv2.getTextSize('00:00.000', self.font,
self.text_size,
self.text_thickness)[0]
spacing = self.text_dims[0] // 2
self.txt_left_lmt = spacing
self.txt_right_lmt = self.num_steps - spacing
# -- Main thread vars -- #
# Operation Params
self.curr_loc = -1
self.text_hloc = 0
self.start_time = None
self.output_array = None
self.image = None
self.displaying_error_image = False
self.targ_perim = CV2TargetAreaPerimeter()
# Progress bar segments for each element
self.pbar_slice = None
self.mouse_in_targ_slice = None
self.mouse_stim_slice = None
# Mouse Status
self.mouse_in_target = False # is mouse inside target region?
self.mouse_recv_stim = False # does mouse receive stimulation?
self.mouse_stim_timer = None # timer to make sure mouse receives STIM_ON secs stim, max every STIM_TOTAL secs
self.in_targ_stopwatch = StopWatch(
) # total time spent in target region
self.get_stim_stopwatch = StopWatch(
) # total time spent receiving stimulation
self.mouse_n_entries = 0 # num entries into target region
self.mouse_n_stims = 0 # num stimulations received
# -- Modifier vars (read-only for main thread) -- #
# Operation Params
self._running = False
self._duration = initial_duration
# Initializing functions. Call once once new process starts
def init_unpickleable_objs(self):
"""These objects must be created in the process they will run in"""
self.output_array = self.mp_array.generate_np_array()
self.image = self.output_array.copy()
# Progress bar slices
self.pbar_slice = self.image[20:60, :, :1]
self.mouse_in_targ_slice = self.image[20:40, :, 1:2]
self.mouse_stim_slice = self.image[40:60, :, 2:3]
# Text slices
w, h = self.text_dims
self.main_timer_slice = self.image[:self.text_vloc + 5, :, :]
self.targ_timer_slice = self.image[92 - h:92 + 5, 95:95 + w, :]
self.stim_timer_slice = self.image[92 - h:92 + 5, 411:411 + w, :]
self.targ_count_slice = self.image[92 - h:92 + 5,
256:256 + int(w * 3 / 4), :]
self.stim_count_slice = self.image[92 - h:92 + 5,
563:563 + int(w * 3 / 4), :]
# Arduino object
cv2.putText(self.output_array, 'CONNECTING TO ARDUINO...', (30, 63),
cv2.FONT_HERSHEY_SIMPLEX, 1.3, (255, 255, 255), 1)
self.output_array.set_can_recv_img()
self.arduino = ArduinoDevice()
self.arduino.connect() # this step takes a few seconds
self.output_array.fill(0)
# Set Progress bar to initial conditions
self.reset_bar()
# Modifier Functions. Can call from other threads
# *** Non-underscored variables are READ ONLY
def set_duration(self, duration_in_secs):
self._duration = float(duration_in_secs)
def set_start(self):
self._running = True
self.reset_bar()
def set_stop(self):
self._running = False
# Main Update Function. Run in Main Thread. Do NOT call from any other thread
# *** Underscored variables are READ ONLY
def set_timer_text(self, reset):
"""Places cv2 text on output array"""
if reset:
main_timer = '00:00.000'
mouse_in_region_timer = '00:00.000'
mouse_recv_stim_timer = '00:00.000'
num_entries = '0)'
num_stims = '0)'
else:
main_timer = format_secs(time.perf_counter() - self.start_time,
'with_ms')
mouse_in_region_timer = format_secs(
self.in_targ_stopwatch.elapsed(), 'with_ms')
mouse_recv_stim_timer = format_secs(
self.get_stim_stopwatch.elapsed(), 'with_ms')
num_entries = '{})'.format(self.mouse_n_entries)
num_stims = '{})'.format(self.mouse_n_stims)
self.main_timer_slice.fill(0)
self.targ_timer_slice.fill(0)
self.stim_timer_slice.fill(0)
self.targ_count_slice.fill(0)
self.stim_count_slice.fill(0)
cv2.putText(self.main_timer_slice,
main_timer, (self.text_hloc, self.text_vloc),
fontFace=self.font,
fontScale=self.text_size,
color=(255, 255, 255))
cv2.putText(self.targ_timer_slice,
mouse_in_region_timer, (0, self.text_dims[1]),
fontFace=self.font,
fontScale=self.text_size,
color=(255, 255, 255))
cv2.putText(self.stim_timer_slice,
mouse_recv_stim_timer, (0, self.text_dims[1]),
fontFace=self.font,
fontScale=self.text_size,
color=(255, 255, 255))
cv2.putText(self.targ_count_slice,
num_entries, (0, self.text_dims[1]),
fontFace=self.font,
fontScale=self.text_size,
color=(255, 255, 255))
cv2.putText(self.stim_count_slice,
num_stims, (0, self.text_dims[1]),
fontFace=self.font,
fontScale=self.text_size,
color=(255, 255, 255))
# Image is now fully prepared, send
self.output_array.send_img(self.image)
def check_mouse_inside_target(self, coord):
"""Checks if mouse is inside target region"""
x, y = coord
self.mouse_in_target = False
if not self.targ_perim.draw or coord == (None, None):
return
x1, x2, y1, y2 = self.targ_perim.x1, self.targ_perim.x2, self.targ_perim.y1, self.targ_perim.y2
if x1 <= x <= x2 and y1 <= y <= y2:
self.mouse_in_target = True
def send_stim_to_mouse(self):
"""Stim mouse if inside region"""
if self.mouse_stim_timer:
elapsed = time.perf_counter() - self.mouse_stim_timer
if STIM_ON > elapsed:
return
elif STIM_ON <= elapsed < STIM_TOTAL:
self.mouse_recv_stim = False
return
elif STIM_TOTAL <= elapsed:
if self.mouse_in_target:
self.mouse_recv_stim = True
self.mouse_stim_timer = time.perf_counter()
return
elif not self.mouse_in_target:
self.mouse_recv_stim = False
self.mouse_stim_timer = None
return
else:
if self.mouse_in_target:
self.mouse_recv_stim = True
self.mouse_stim_timer = time.perf_counter()
def reset_bar(self):
"""Resets to initial conditions"""
# Reset Locations
self.curr_loc = -1
self.text_hloc = 0
# Reset Mouse Location timers and counters
self.in_targ_stopwatch.reset()
self.get_stim_stopwatch.reset()
self.mouse_n_entries = 0
self.mouse_n_stims = 0
self.mouse_stim_timer = None
# Reset Progress Bar Image
self.reset_progbar_img()
# Send Image
self.output_array.set_can_recv_img()
def reset_progbar_img(self):
"""Reset progressbar to initial image"""
self.image.fill(0)
# Add new progress bar at origin
self.pbar_slice[:, :1, :] = 255
# Add time indicators
self.image[61:62, :, :] = 255
num_chunks = int(self._duration / 30)
num_chunks = min([12, num_chunks])
num_chunks = max([2, num_chunks])
seg_size = int(self.image.shape[1] / num_chunks)
time_chunk = (self._duration / num_chunks)
for i in np.arange(1, num_chunks):
loc = seg_size * i
tloc = format_secs(time_chunk * i)
self.image[62:65, loc - 1:loc, :] = 255
cv2.putText(self.image, tloc, (loc - 15, 74),
cv2.FONT_HERSHEY_SIMPLEX, 0.3, (255, 255, 255))
# Add legends
self.image[79:96, 319:320, :] = 255
self.image[80:95, 3:18, 1] = 255
cv2.putText(self.image, 'In Region:', (19, 92),
cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 255))
cv2.putText(self.image, '(# Entries:', (175, 92),
cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 255))
self.image[80:95, 323:338, 2] = 255
cv2.putText(self.image, 'Get Stim:', (340, 92),
cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 255))
cv2.putText(self.image, '(# Stims:', (491, 92),
cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 255))
# add initializing timer text
self.set_timer_text(reset=True)
def can_update(self):
"""Checks we are allowed to proceed"""
finished = self.curr_loc >= self.num_steps
if finished or not self._running:
return False
return True
def update(self):
"""Draws next frame of progress bar"""
# Get elapsed time, get expected location, check mouse location/stim status
elapsed = time.perf_counter() - self.start_time
loc = int((elapsed / self._duration) * self.num_steps)
# Check if mouse in target region; also calculate total time inside
if self.mouse_in_target:
self.mouse_in_targ_slice[:, loc - 1:loc, :] = 255
if not self.in_targ_stopwatch.started:
self.mouse_n_entries += 1
self.in_targ_stopwatch.start()
else:
if self.in_targ_stopwatch.started:
self.in_targ_stopwatch.stop()
# Check if mouse receive stimulation; calculate total time receive
if self.mouse_recv_stim:
self.mouse_stim_slice[:, loc - 1:loc, :] = 255
if not self.get_stim_stopwatch.started:
self.mouse_n_stims += 1
self.get_stim_stopwatch.start()
if not self.arduino.manual_mode:
self.arduino.send_signal()
else:
if self.get_stim_stopwatch.started:
self.get_stim_stopwatch.stop()
# If enough time elapsed, update current location to expected location
if loc != self.curr_loc:
self.pbar_slice[:, self.curr_loc - 1:self.curr_loc, :] = 0
self.pbar_slice[:, loc - 1:loc + 1, :] = 255
if self.txt_left_lmt <= loc <= self.txt_right_lmt:
self.text_hloc = loc - self.txt_left_lmt
self.curr_loc = loc
# Test for arduino connection; send new progbar frame; attempt to reconnect lost devices
self.ping_arduino(updating=True)
def ping_arduino(self, updating):
"""Pings arduino, displays errors, attempt to reconnect"""
self.arduino.ping()
if self.output_array.can_send_img():
if not self.arduino.connected:
self.display_error_img()
elif updating:
self.set_timer_text(reset=False)
elif not updating and self.displaying_error_image:
self.displaying_error_image = False
self.output_array[:] = self.image
self.output_array.set_can_recv_img()
if not self.arduino.connected:
self.arduino.connect()
def display_error_img(self):
if not self.displaying_error_image:
self.displaying_error_image = True
self.output_array.fill(0)
cv2.putText(self.output_array, 'ARDUINO ERROR. RECONNECT DEVICE',
(10, 60), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255),
2)