Skip to content

Enhance testing and improve example handling #3

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
pgleeson merged 4 commits into from
Jun 5, 2026
Merged

Enhance testing and improve example handling #3

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3e53235
Much better testing and loads wcon.zip examples
pgleeson Jun 4, 2026
084e4c2
Handling of ox/oy & better handling when data is array
pgleeson Jun 4, 2026
7dd05de
Improved app.py
pgleeson Jun 4, 2026
84eb512
Refactor SimpleWCON to own file
pgleeson Jun 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
150 changes: 150 additions & 0 deletions SimpleWCON.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,150 @@
import json

import numpy as np


class SimpleWCON:
def __init__(self, wcon_file):

# Check whether the filename ends with wcon.zip, and if so, unzip it, read the wcon file it contains and load the contents of that using json
if wcon_file.endswith("wcon.zip"):
import zipfile

with zipfile.ZipFile(wcon_file, "r") as zip_ref:
wcon_files = [f for f in zip_ref.namelist() if f.endswith(".wcon")]
if len(wcon_files) == 0:
raise ValueError("No .wcon file found in the zip archive.")
elif len(wcon_files) > 1:
raise ValueError(
"Multiple .wcon files found in the zip archive. Not yet supported..."
)
wcon_file_in_zip = wcon_files[0]
print(f"Extracting {wcon_file_in_zip} from {wcon_file}...")
with zip_ref.open(wcon_file_in_zip) as f:
wcon = json.load(f)
else:
with open(wcon_file, "r") as f:
print(f" === Loading WCON from file: {wcon_file}...")
wcon = json.load(f)

print(
" - WCON file loaded. Keys found: "
+ ", ".join(wcon.keys())
+ ". Processing data..."
)

self.extras = {}
for key in wcon:
if key.startswith("@"):
self.extras[key] = wcon[key]

self.t_units = "Unknown units"
self.x_units = "Unknown units"
self.y_units = "Unknown units"

if "units" in wcon:
self.t_units = wcon["units"].get("t")

self.x_units = wcon["units"].get("x")
self.y_units = wcon["units"].get("y")
print(
f" Time units: {self.t_units}, x units: {self.x_units}, y units: {self.y_units}"
)

# "data" is arrayable: it may be a single record object or a list of
# records (multiple worms / timepoint chunks). Normalise to a list.
data = wcon["data"]
if isinstance(data, dict):
data = [data]

record = data[0]
if len(data) > 1:
print(
" - Note: %d data records found; this minimal viewer only displays the first (id=%s)."
% (len(data), record.get("id"))
)

print(" - Data records: %d" % len(data))
print(" - Data keys: %s" % list(record.keys()))
print(" - Data time: %s" % len(record["t"]))
print(" - Data x: %s" % len(record["x"]))
print(" - Data y: %s" % len(record["y"]))

self.times = np.array(record["t"])

factor = 1

if self.x_units == "millimeters" or self.x_units == "mm":
factor = 1
self.x_units_used = self.x_units
self.y_units_used = self.y_units
elif (
self.x_units == "micrometers"
or self.x_units == "um"
or self.x_units == "µm"
):
factor = 1e-3
self.x_units_used = "mm"
self.y_units_used = "mm"
else:
self.x_units_used = self.x_units
self.y_units_used = self.y_units

# Cast to float so that any None values in experimental wcon data
# become nan (np.array([1.0, None], dtype=float) -> [1., nan]).
self.x = np.array(record["x"], dtype=float).T * factor
self.y = np.array(record["y"], dtype=float).T * factor

# Low-resolution trackers may give a single xy point per timepoint, so
# x/y are 1-D (n_timepoints,). Reshape to (1, n_timepoints) so that the
# [:, ti] indexing used downstream works uniformly with the spine case.
if self.x.ndim == 1:
self.x = self.x.reshape(1, -1)
self.y = self.y.reshape(1, -1)

# Variable origin: if ox/oy are present, all positional values at a
# timepoint are relative to that origin (spec: "Variable origin and
# centroid position"). A minimal reader must add it back to recover
# absolute coordinates. ox/oy are per-timepoint arrays (or a single
# local constant); x has shape (n_body, n_t) so (n_t,) broadcasts.
self.ox = record.get("ox")
self.oy = record.get("oy")
if self.ox is not None and self.oy is not None:
self.ox = np.atleast_1d(np.array(self.ox, dtype=float)) * factor
self.oy = np.atleast_1d(np.array(self.oy, dtype=float)) * factor
self.x = self.x + self.ox
self.y = self.y + self.oy

print(f"Times: {self.times}, shape: {self.times.shape}")
print(f"x: {self.x}, shape: {self.x.shape}")
print(f"y: {self.y}, shape: {self.y.shape}")

self.xmax = np.nanmax(self.x)
self.xmin = np.nanmin(self.x)
self.ymax = np.nanmax(self.y)
self.ymin = np.nanmin(self.y)
print(
f"Range of time: {self.times[0]}{self.t_units}->{self.times[-1]}{self.t_units}; x range: {self.xmin}{self.x_units}->{self.xmax}{self.x_units}; y range: {self.ymin}{self.y_units}->{self.ymax}{self.y_units}"
)

if "px" in record:
self.px = np.array(record["px"], dtype=float) * factor # (n_t, n_perim)
# px is arrayed per timepoint, so the origin (n_t,) is applied
# along axis 0 via a trailing axis.
if self.ox is not None and self.px.ndim == 2:
self.px = self.px + self.ox[:, None]
print(
f"px shape: {np.array(self.px).shape}, max: {np.nanmax(self.px)}, min: {np.nanmin(self.px)}"
)
else:
self.px = None

if "py" in record:
self.py = np.array(record["py"], dtype=float) * factor # (n_t, n_perim)
if self.oy is not None and self.py.ndim == 2:
self.py = self.py + self.oy[:, None]
print(
f"py shape: {np.array(self.py).shape}, max: {np.nanmax(self.py)}, min: {np.nanmin(self.py)}"
)
else:
self.py = None
107 changes: 35 additions & 72 deletions WormView.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,11 @@
from matplotlib import pyplot as plt
import numpy as np
import json
import math
import os
import argparse
import sys
from Player import Player
from SimpleWCON import SimpleWCON


def validate_file(file_path):
Expand All @@ -16,72 +16,16 @@ def validate_file(file_path):
return file_path


class SimpleWCON:
def __init__(self, wcon_file):

with open(wcon_file, "r") as f:
print(f" === Loading WCON from file: {wcon_file}...")
wcon = json.load(f)

print(
" - WCON file loaded. Keys found: "
+ ", ".join(wcon.keys())
+ ". Processing data..."
)

self.extras = {}
for key in wcon:
if key.startswith("@"):
self.extras[key] = wcon[key]

self.t_units = "??"
self.x_units = "??"
self.y_units = "??"

if "units" in wcon:
self.t_units = wcon["units"].get("t")
self.x_units = wcon["units"].get("x")
self.y_units = wcon["units"].get("y")
print(
f" Time units: {self.t_units}, x units: {self.x_units}, y units: {self.y_units}"
)

print(" - Data points: %d" % len(wcon["data"]))
print(" - Data keys: %s" % list(wcon["data"][0].keys()))
print(" - Data time: %s" % len(wcon["data"][0]["t"]))
print(" - Data x: %s" % len(wcon["data"][0]["x"]))
print(" - Data y: %s" % len(wcon["data"][0]["y"]))

self.times = np.array(wcon["data"][0]["t"])
self.x = np.array(wcon["data"][0]["x"]).T
self.y = np.array(wcon["data"][0]["y"]).T

# Required for expeerimental wcon data...
# replace any values in self.x and self.y which are None with nan
self.x = np.where(self.x == None, np.nan, self.x) # noqa: E711
self.y = np.where(self.y == None, np.nan, self.y) # noqa: E711

print(f"Times: {self.times}, shape: {self.times.shape}")
print(f"x: {self.x}, shape: {self.x.shape}")
print(f"y: {self.y}, shape: {self.y.shape}")

self.xmax = np.nanmax(self.x)
self.xmin = np.nanmin(self.x)
self.ymax = np.nanmax(self.y)
self.ymin = np.nanmin(self.y)
print(
f"Range of time: {self.times[0]}{self.t_units}->{self.times[-1]}{self.t_units}; x range: {self.xmax}{self.x_units}->{self.xmin}{self.x_units}; y range: {self.ymax}{self.y_units}->{self.ymin}{self.y_units}"
)

self.px = wcon["data"][0]["px"] if "px" in wcon["data"][0] else None
self.py = wcon["data"][0]["py"] if "py" in wcon["data"][0] else None


class WormView:
midline_plot = None
perimeter_plot = None
head_plot = None
times = None

def __init__(self, show_head=False):
self.show_head = show_head
print(" - Initializing WormView")

def get_perimeter(self, x, y, r):
n_bar = x.shape[0]
num_steps = x.shape[1]
Expand Down Expand Up @@ -131,6 +75,7 @@ def reset(self):
print(" - Resetting WormView")
self.midline_plot = None
self.perimeter_plot = None
self.head_plot = None
plt.close("all")

def get_plot(self, args):
Expand All @@ -143,8 +88,8 @@ def get_plot(self, args):

self.wcon = SimpleWCON(args.wcon_file)

self.ax.set_xlabel("x (%s)" % self.wcon.x_units)
self.ax.set_ylabel("y (%s)" % self.wcon.y_units)
self.ax.set_xlabel("x (%s)" % self.wcon.x_units_used)
self.ax.set_ylabel("y (%s)" % self.wcon.y_units_used)

factor = 0.05
if abs(self.wcon.xmax - self.wcon.xmin) > abs(self.wcon.ymax - self.wcon.ymin):
Expand Down Expand Up @@ -181,10 +126,6 @@ def get_plot(self, args):
else:
print("No objects found")

# Set the limits of the plot since we don't have any objects to help with autoscaling

self.ax.set_ylim([-1.5, 1.5])

if self.wcon.px is not None and self.wcon.py is not None:
if args.ignore_wcon_perimeter:
print(
Expand All @@ -196,7 +137,7 @@ def get_plot(self, args):
self.px = np.array(self.wcon.px).T
self.py = np.array(self.wcon.py).T
else:
if not args.suppress_automatic_generation:
if not args.suppress_automatic_generation and self.wcon.x.shape[0] >= 2:
print("Computing perimeter from midline")
self.px, self.py = self.get_perimeter(
self.wcon.x, self.wcon.y, args.minor_radius
Expand Down Expand Up @@ -241,7 +182,21 @@ def update(self, ti):
else:
self.perimeter_plot.set_data(self.px[:, ti], self.py[:, ti])

return self.midline_plot, self.perimeter_plot
if self.show_head:
if self.head_plot is None:
print("Adding head")
(self.head_plot,) = self.ax.plot(
[self.wcon.x[0, ti]],
[self.wcon.y[0, ti]],
color="r",
marker="o",
label="t=%sms" % self.wcon.times[ti],
linewidth=2,
)
else:
self.head_plot.set_data([self.wcon.x[0, ti]], [self.wcon.y[0, ti]])

return self.midline_plot, self.perimeter_plot, self.head_plot


def parse_args():
Expand Down Expand Up @@ -278,6 +233,12 @@ def parse_args():
help="Minor radius of the worm in millimeters (default: 40e-3)",
required=False,
)
parser.add_argument(
"-head",
"--show_head",
action="store_true",
help="Show the head of the worm.",
)

args = parser.parse_args()

Expand All @@ -289,7 +250,9 @@ def main():

args = parse_args()

wv = WormView()
print(" - Arguments parsed: %s" % args)

wv = WormView(show_head=args.show_head)

fig, ax = wv.get_plot(args)

Expand All @@ -313,7 +276,7 @@ def update(ti):
from matplotlib.animation import FFMpegWriter

FFwriter = FFMpegWriter(fps=10)
mp4_file = args.wcon_file.replace(".wcon", ".mp4")
mp4_file = args.wcon_file.replace(".wcon.zip", ".wcon").replace(".wcon", ".mp4")
print(f"Saving animation to: {mp4_file}")
anim.save(mp4_file, writer=FFwriter)

Expand Down
Loading