'''
# TODO save with tracking data from all body parts for animation making
for ROI in arena.ROIs_dict.keys():
# create figure and draw ROI image
f = plt.figure(figsize=(20, 12))
f.suptitle(ROI)
f._save_name = f"roi_{ROI}"
axes = f.subplot_mosaic("""
AABC
AADE
FFGH
""")
for ax in "ABCDE":
draw.ROI(ROI, ax=axes[ax], set_ax=True if ax == "A" else False)
# fetch from database
crossings = pd.DataFrame((ROICrossing * ROICrossing.InitialCondition
& f'roi="{ROI}"'
& "mouse_exits=1").fetch())
logger.info(f"Loaded {len(crossings)} crossings")
color = arena.ROIs_dict[crossings.iloc[0].roi].color
# draw tracking
for i, cross in crossings.iterrows():
if i % 30 == 0:
axes["A"].scatter(
cross.x[0],
cross.y[0],
# switch and lock switch for 20 frames
lock = swc.switch_recognition(thresholded, fingers)
channels_index = swc.switch_channel(lock, channels_index,
len(all_channels))
frame_list = all_channels[channels_index]
# Draw contours around real hand in live stream and the number of fingers
visual_color = cv2.flip(visual_color, 1)
cv2.drawContours(visual_color,
[hand_segment + (roi_right, roi_top)], -1,
(255, 0, 0), 1)
draw.finger_number(frame_copy, fingers)
draw.finger_number(visual_color, fingers)
draw.ROI(visual_color, roi_left, roi_top, roi_right, roi_bottom)
cv2.imshow("Visual threshold", visual_threshold)
cv2.imshow("Visual color", visual_color)
frame_number_on_video = vm.control(fingers, frame_number_on_video)
draw.ROI(frame_copy, roi_left, roi_top, roi_right, roi_bottom)
number_of_frame += 1
cv2.imshow("Finger Count", frame_copy)
# Close windows with Esc
k = cv2.waitKey(1) & 0xFF
if k == 27:
break
print(f'Kept {len(_bouts)} bouts')
# create bouts after resampling in time
paths = time.time_rescale([Path(bout.x, bout.y).smooth() for i,bout in _bouts.iterrows()])
# %%
f = plt.figure(figsize=(16, 8))
axes = f.subplot_mosaic(
"""
AABBB
AACCC
"""
)
draw.ROI(ROI, ax=axes["A"])
SKIP = 0
N = 10 # len(_bouts)
S, T = np.full((80, N), np.nan), np.full((80, N), np.nan)
_S, _T = np.full((80, N), np.nan), np.full((80, N), np.nan)
simulated = []
for i in track(range(N)):
bout = paths[i]
trajectory, _time = MSD(
bout, skip=SKIP, start_frame=2, end_frame=-2
).simulate()
# plot
draw.Tracking(bout.x, bout.y, zorder=-1, ax=axes["A"], alpha=0.8)
def plot_roi_crossing(
crossing: pd.Series,
tracking: dict = None,
step: int = 5,
highlight: str = None,
arrow_scale=0.5,
) -> plt.Figure:
"""
Plots an entire ROI crossing, tracking, vectors and mouse
"""
if highlight == "velocity":
v_alpha, a_alpha = 1, 0.4
elif highlight == "acceleration":
v_alpha, a_alpha = 0.4, 1
else:
v_alpha, a_alpha = 1, 1
# path = Path(
# convolve_with_gaussian(crossing.x, kernel_width=11),
# convolve_with_gaussian(crossing.y, kernel_width=11))
path = Path(crossing.x, crossing.y)
f = plt.figure(figsize=(20, 10))
axes = f.subplot_mosaic(
"""
AACC
AADD
BBEE
"""
)
f._save_name = (
f"{crossing.roi}_{crossing.crossing_id}" + ""
if highlight is None
else highlight
)
# draw main crossing plot
draw.ROI(crossing.roi, ax=axes["A"], set_ax=True)
draw.Tracking(path.x, path.y, lw=0.5, color="k", ax=axes["A"])
# draw velocity and acceleartion vectors
draw.Arrows(
path.x,
path.y,
path.velocity.angle,
label="velocity",
L=arrow_scale * path.velocity.magnitude / 30,
step=step,
color=colors.velocity,
ax=axes["A"],
outline=True,
alpha=v_alpha,
)
draw.Arrows(
path.x,
path.y,
path.acceleration.angle,
label="acceleration",
L=arrow_scale * 6 * path.acceleration.magnitude / 30,
step=step,
color=colors.acceleration,
ax=axes["A"],
outline=True,
alpha=a_alpha,
)
draw.Tracking.scatter(
path.x[::step], path.y[::step], color="k", ax=axes["A"], zorder=200,
)
# draw speed traces
time = np.arange(len(path.speed)) / 60
axes["B"].fill_between(time, 0, path.speed, alpha=0.5, color=colors.speed)
axes["B"].plot(
time, path.speed, lw=4, color=colors.speed,
)
if tracking is not None:
for bp in PAWS:
axes["C"].plot(
tracking[bp]["bp_speed"], color=colors.bodyparts[bp], label=bp
)
axes["C"].plot(
tracking["body"]["bp_speed"],
color=colors.bodyparts["body"],
label="body",
)
axes["C"].plot(
np.mean(
np.vstack(
[tracking[bp]["bp_speed"] for bp in PAWS if "h" in bp]
),
0,
),
color="k",
label="mean",
)
# clean plots
axes["A"].legend()
axes["B"].set(xlabel="time (s)", ylabel="speed (cm/s)")
axes["C"].legend()
f.tight_layout()
return f
def animate_one(ROI: str, crossing_id: int, FPS: int):
scale = 1 / 30 # to scale velocity vectors
save_folder = (paths.analysis_folder / "behavior" /
"roi_crossings_animations")
# ----------------------------------- prep ----------------------------------- #
# load tracking
bouts = pd.read_hdf(paths.analysis_folder / "behavior" / "roi_crossings" /
f"{ROI}_crossings.h5")
bout = bouts.sort_values("duration").iloc[crossing_id]
logger.info(f"Animating bout {crossing_id} - {round(bout.duration, 2)}s")
# tracking: dict = ROICrossing.get_crossing_tracking(bout.crossing_id)
# generate a path from tracking
path = Path(bout.x, bout.y)
# create fig and start camera
if ROI == "T4":
f = plt.figure(figsize=(12, 12))
elif "S" in ROI:
f = plt.figure(figsize=(8, 14))
else:
f = plt.figure(figsize=(5, 14))
axes = f.subplot_mosaic("""
AA
AA
BB
""")
camera = Camera(f)
# ----------------------------------- plot ----------------------------------- #
n_interpolated_frames = int(60 / FPS)
n_frames = len(bout.x) - 1
trajectory = GrowingPath()
for frame in track(range(n_frames), total=n_frames):
# repeat each frame N times interpolating between frames
for interpol in np.linspace(0, 1, n_interpolated_frames):
# get interpolated quantities
x = interpolate_at_frame(path.x, frame, interpol)
y = interpolate_at_frame(path.y, frame, interpol)
speed = interpolate_at_frame(path.speed, frame, interpol)
trajectory.update(x, y, speed=speed)
# time elapsed
_time = (np.arange(len(trajectory.speed)) / n_interpolated_frames /
60)
# plot the arena
draw.ROI(ROI, set_ax=True, shade=False, ax=axes["A"])
# plot tracking so far
draw.Tracking(trajectory.x, trajectory.y, lw=3, ax=axes["A"])
draw.Dot(x, y, s=50, ax=axes["A"])
# plot speed and velocity vectors
draw.Arrow(
x,
y,
interpolate_at_frame(path.velocity.angle,
frame,
interpol,
method="step"),
L=scale *
interpolate_at_frame(path.velocity.magnitude, frame, interpol),
color=colors.velocity,
outline=True,
width=2,
ax=axes["A"],
)
draw.Arrow(
x,
y,
path.acceleration.angle[frame],
L=4 * scale * path.acceleration.magnitude[frame],
color=colors.acceleration,
outline=True,
width=2,
ax=axes["A"],
zorder=200,
)
# plot speed trace
axes["B"].fill_between(_time,
0,
trajectory.speed,
alpha=0.5,
color=colors.speed)
axes["B"].plot(_time, trajectory.speed, lw=4, color=colors.speed)
axes["A"].set(title=f"{ROI} - crossing: {crossing_id}")
axes["B"].set(xlabel="time (s)", ylabel="speed (cm/s)")
camera.snap()
# ------------------------------- video creation ------------------------------- #
save_path = save_folder / f"{ROI}_{bout.crossing_id}.mp4"
logger.info(f"Saving animation @: '{save_path}'")
animation = camera.animate(interval=1000 / FPS)
animation.save(save_path, fps=FPS)
logger.info("Done!")
plt.cla()
plt.close(f)
del camera
del animation
# ----------------------------------- plot ----------------------------------- #
f = plot_roi_crossing(bout, step=4)
recorder.add_figures(svg=False)
plt.close(f)
)
selected_bouts = bouts.sample(6).reset_index()
# draw stuff
f, axes = plt.subplots(figsize=(18, 12), ncols=3, nrows=2)
axes = axes.flatten()
for i, bout in selected_bouts.iterrows():
# generate a path from tracking
path = Path(bout.x, bout.y)
ax = axes[i]
# draw arena and tracking
draw.ROI(ROI, set_ax=True, shade=False, ax=ax)
# draw.Tracking(bouts.x, bouts.y, alpha=0.5)
# draw bout tracking and velocity/acceleration vectors
# draw.Tracking.scatter(path.x, path.y, c=path.tangent.angle, cmap='bwr', vmin=-180, vmax=180)
draw.Tracking(path.x, path.y, color=[0.6, 0.6, 0.6], lw=5, ax=ax)
draw.Arrows(
path.x,
path.y,
path.velocity.angle,
L=path.velocity.magnitude,
color=blue,
step=2,
outline=True,
width=2,
for i, bout in _bouts.iterrows():
crosses.append(LocomotionBout(bout))
# %%
# ---------------------------------------------------------------------------- #
# first plot - sanity checks #
# ---------------------------------------------------------------------------- #
"""
Plot tracking over threshold crossings for sanity checks
"""
colors = dict(velocity=blue_dark, acceleration=pink)
f, axes = plt.subplots(figsize=(16, 10), ncols=2)
for ax in axes:
draw.ROI(ROI, set_ax=True, ax=ax)
f._save_name = f"vector_field_{ROI}"
all_xy_vectors = {
"velocity": {},
"acceleration": {},
} # stores vectors at each XY position, for averaging
scale_factor = {"velocity": 20, "acceleration": 1}
for vn, (vec_name, xy_vectors) in enumerate(all_xy_vectors.items()):
ax = axes[vn]
cross_bins = []
for cross in crosses:
# round XY tracking to closest even number
X, Y = ((cross.x) / 2).round(0) * 2, ((cross.y) / 2).round(0) * 2
for t, (x, y) in enumerate(zip(X, Y)):
crosses.append(LocomotionBout(cross))
# %%
f = plt.figure(figsize=(20, 10))
axes = f.subplot_mosaic(
"""
AAABBDDFFF
AAACCEEFFF
"""
)
draw.ROI(
ROI,
ax=axes["A"],
set_ax=True,
# img_path=r'C:\Users\Federico\Documents\GitHub\pysical_locomotion\draw\hairpin.png'
img_path="/Users/federicoclaudi/Documents/Github/LocomotionControl/draw/hairpin.png",
)
draw.ROI(
ROI,
ax=axes["F"],
set_ax=True,
# img_path=r'C:\Users\Federico\Documents\GitHub\pysical_locomotion\draw\hairpin.png'
img_path="/Users/federicoclaudi/Documents/Github/LocomotionControl/draw/hairpin.png",
)
for cross in crosses:
time = cross.gcoord
axes["B"].plot(time, cross.speed, color="k", alpha=0.2)
# %%
# ---------------------------------------------------------------------------- #
# first plot - sanity checks #
# ---------------------------------------------------------------------------- #
"""
Plot tracking over threshold crossings for sanity checks
"""
f = plt.figure(figsize=(24, 10))
axes = f.subplot_mosaic("""
AAABBBBFFFG
AAACCCCFFFH
""")
draw.ROI(
ROI,
ax=axes["A"],
set_ax=True,
)
draw.ROI(
ROI,
ax=axes["F"],
set_ax=True,
)
for n, cross in enumerate(crosses):
time = np.linspace(0, cross.duration, len(cross.x))
# mark when the thresholds were exceeded
try:
acc_cross = th_crossess.loc[(th_crossess.variable == "acceleration")
& (th_crossess.n == n)].iloc[0]
aac_crooss = th_crossess.loc[
# %%
from data.data_utils import convolve_with_gaussian
import seaborn as sns
f = plt.figure(figsize=(20, 12), constrained_layout=True)
f.suptitle(f'{unit.brain_region} - {unit.unit_id}')
axes = f.subplot_mosaic('''
AABBCCFF
AADDEEGG
''')
unit = units.loc[units.unit_id == 1036].iloc[0]
frate = convolve_with_gaussian(unit.firing_rate)
draw.ROI(roi, ax=axes['A'])
X, Y = [], []
for bout in bouts:
draw.Tracking(bout.x, bout.y - 2, ax=axes['A'])
axes['B'].plot(bout.velocity.magnitude, color='k', lw=.5)
axes['C'].plot(bout.thetadot, color='k', lw=.5)
axes['D'].plot(frate[bout.start_frame:bout.end_frame])
X.extend(list(frate[bout.start_frame + 2:bout.end_frame - 2]))
Y.extend(list(bout.speed))
sns.regplot(X, Y, ax=axes['E'])
starts = [b.start_frame for b in bouts]
ends = [b.end_frame for b in bouts]
