epicure.trackmate_export

  1from datetime import datetime
  2from pathlib import Path
  3from typing import Dict, List
  4from xml.dom import minidom
  5import xml.etree.ElementTree as ET
  6#import cv2
  7#import matplotlib.pyplot as plt
  8import numpy as np
  9import pandas as pd
 10from scipy.cluster.hierarchy import DisjointSet
 11from skimage.measure import find_contours
 12
 13import epicure.Utils as ut
 14
 15
 16# TODO: deal with groups (store in MANUAL_SPOT_COLOR)
 17
 18SPOT_FEATS = [
 19    {"feature": "POSITION_X", "name": "X", "shortname": "X", "dimension": "POSITION", "isint": "false"},
 20    {"feature": "POSITION_Y", "name": "Y", "shortname": "Y", "dimension": "POSITION", "isint": "false"},
 21    {"feature": "POSITION_Z", "name": "Z", "shortname": "Z", "dimension": "POSITION", "isint": "false"},
 22    {"feature": "POSITION_T", "name": "T", "shortname": "T", "dimension": "TIME", "isint": "false"},
 23    {"feature": "FRAME", "name": "Frame", "shortname": "Frame", "dimension": "NONE", "isint": "true"},
 24    {"feature": "VISIBILITY", "name": "Visibility", "shortname": "Visibility", "dimension": "NONE", "isint": "true"},
 25    {"feature": "RADIUS", "name": "Radius", "shortname": "R", "dimension": "LENGTH", "isint": "false"},
 26    {"feature": "MANUAL_SPOT_COLOR", "name": "Manual spot color", "shortname": "Spot color", "dimension": "NONE", "isint": "true"},
 27]
 28
 29EDGE_FEATS = [
 30    {"feature": "SPOT_SOURCE_ID", "name": "Source spot ID", "shortname": "Source ID", "dimension": "NONE", "isint": "true"},
 31    {"feature": "SPOT_TARGET_ID", "name": "Target spot ID", "shortname": "Target ID", "dimension": "NONE", "isint": "true"},
 32]
 33
 34TRACK_FEATS = [
 35    {"feature": "TRACK_INDEX", "name": "Track index", "shortname": "Index", "dimension": "NONE", "isint": "true"},
 36    {"feature": "TRACK_ID", "name": "Track ID", "shortname": "ID", "dimension": "NONE", "isint": "true"},
 37]
 38
 39
 40def build_feat_declaration_tag():
 41    """Build the FeatureDeclarations tag for TrackMate XML."""
 42    feat_declarations = ET.Element("FeatureDeclarations")
 43    spot_feats = ET.SubElement(feat_declarations, "SpotFeatures")
 44    for feat in SPOT_FEATS:
 45        ET.SubElement(spot_feats, "Feature", feat)
 46    edge_feats = ET.SubElement(feat_declarations, "EdgeFeatures")
 47    for feat in EDGE_FEATS:
 48        ET.SubElement(edge_feats, "Feature", feat)
 49    track_feats = ET.SubElement(feat_declarations, "TrackFeatures")
 50    for feat in TRACK_FEATS:
 51        ET.SubElement(track_feats, "Feature", feat)
 52    return feat_declarations
 53
 54
 55def build_spots_df(epic):
 56    """Build a DataFrame representing the spots table for TrackMate XML."""
 57    df_spots = pd.DataFrame(epic.tracking.track_data, columns=["label", "FRAME", "pos_x", "pos_y"])
 58    df_spots["ID"] = df_spots.index
 59    df_spots["name"] = df_spots.apply(lambda row: f"LABEL{row['label']}_FRAME{row['FRAME']}", axis=1)
 60    # Invert X and Y to match TrackMate convention.
 61    df_spots["POSITION_X"] = df_spots["pos_y"] * epic.epi_metadata.get("ScaleXY", 1)
 62    df_spots["POSITION_Y"] = df_spots["pos_x"] * epic.epi_metadata.get("ScaleXY", 1)
 63    df_spots["POSITION_Z"] = 0.0  # 2D data. 
 64    df_spots["POSITION_T"] = df_spots["FRAME"] * epic.epi_metadata.get("ScaleT", 1)
 65    df_spots["VISIBILITY"] = 1
 66    df_spots.drop(columns=["pos_x", "pos_y"], inplace=True)
 67
 68    return df_spots
 69
 70
 71def get_cell_contour(seg_array, label, frame):
 72    """Get the contours of the cell with the given label in the given frame."""
 73    spot_seg = seg_array[frame, :, :] == label
 74    # print(spot_seg.dtype)
 75    # Pad to avoid contours on the border (skimage.measure.find_contours)
 76    spot_seg = np.pad(spot_seg, pad_width=1, mode='constant', constant_values=0)
 77    #spot_seg_uint8 = spot_seg.astype(np.uint8)
 78    # print(spot_seg_uint8.dtype)
 79
 80    if np.sum(spot_seg) > 0:
 81        contours = find_contours(spot_seg, level=0.5)
 82        #contours, _ = cv2.findContours(spot_seg_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
 83        # print(type(contours))
 84        # if label == 125 and frame == 0:
 85        #     contours = contours[:-1]  # hack to remove spurious contour
 86
 87        # plt.show()
 88        # print(contours)
 89        assert len(contours) > 0, f"No contour found for label {label} in frame {frame}"
 90        # assert len(contours) < 2, f"More than one contour found for label {label} in frame {frame}"
 91        if len(contours) > 1:
 92             print(f"Warning, found two cells for label {label} at frame {frame}")
 93             print(f"Keep only the first one in the export")
 94        #if label==6:
 95        #     fig, ax = plt.subplots(1, 1)
 96        #     ax.imshow(seg_array[frame,:,:]>0)
 97             #ax.imshow(spot_seg, cmap="grey")
 98             # Plot each contour
 99        #     for contour in contours:
100        #         contour = contour.squeeze()
101        #         ax.plot(contour[:, 1], contour[:, 0], linewidth=1)
102
103        #     plt.show()
104
105        return np.flip(contours[0].squeeze(), axis=1)
106    return None
107
108
109def build_roi_contours(epic, df_spots):
110    """Build a mapping from spot ID to its ROI contour."""
111    roi_contours = {}
112    seg_layer = epic.seg
113    for _, spot in df_spots.iterrows():
114        contour = get_cell_contour(seg_layer, spot["label"], spot["FRAME"])
115        if contour is not None:
116            # Convert from pixels to space units.
117            contour = contour.astype(np.float64)  # in case the scale is a float
118            # Remove the padding
119            contour[:, 0] -= 1 
120            contour[:, 1] -= 1 
121            contour[:, 0] *= epic.epi_metadata.get("ScaleXY", 1)
122            contour[:, 1] *= epic.epi_metadata.get("ScaleXY", 1)
123            # Convert contour from absolute to relative coordinates (to the cell XY position).
124            contour[:, 0] -= spot["POSITION_X"]
125            contour[:, 1] -= spot["POSITION_Y"]
126            # Flatten contour to a 1D array as expected by TrackMate.
127            contour = contour.flatten()
128            roi_contours[spot["ID"]] = contour
129            # print(spot["ID"], spot["label"], spot["FRAME"], contour.shape)
130    return roi_contours
131
132
133def build_all_spots_tag(df_spots, roi_n_points):
134    """Build the AllSpots tag for TrackMate XML."""
135    all_spots = ET.Element("AllSpots", {"nspots": str(len(df_spots))})
136    frames = df_spots["FRAME"].unique()
137    # print(roi_n_points)
138    for frame in frames:
139        spots_in_frame = df_spots[df_spots["FRAME"] == frame]
140        frame_tag = ET.SubElement(all_spots, "SpotsInFrame", {"frame": str(frame)})
141        for _, spot in spots_in_frame.iterrows():
142            spot_attrib = {
143                "ID": str(spot["ID"]),
144                "name": spot["name"],
145                "EpiCure_label": str(spot["label"]),
146                "POSITION_X": str(spot["POSITION_X"]),
147                "POSITION_Y": str(spot["POSITION_Y"]),
148                "POSITION_Z": str(spot["POSITION_Z"]),
149                "POSITION_T": str(spot["POSITION_T"]),
150                "FRAME": str(spot["FRAME"]),
151                "VISIBILITY": str(spot["VISIBILITY"]),
152                "RADIUS": "-1",
153                "MANUAL_SPOT_COLOR": "0",  # TODO: use it to store group
154            }
155
156            # Add ROI contour if available.
157            if spot["ID"] in roi_n_points:
158                spot_attrib["ROI_N_POINTS"] = str(roi_n_points[spot["ID"]].shape[0] // 2)
159
160            spot_tag = ET.SubElement(frame_tag, "Spot", spot_attrib)
161            if spot["ID"] in roi_n_points:
162                spot_tag.text = " ".join(map(str, roi_n_points[spot["ID"]]))
163
164    return all_spots
165
166
167def create_label_to_track_mapping(divisions: Dict[int, List[int]], unique_labels: List[int]) -> Dict[int, int]:
168    """
169    Create a mapping from labels to track IDs using scipy's DisjointSet for efficient track grouping.
170
171    Args:
172        divisions: dict of {daughter_label: [mother_labels]} from epic.tracking.graph
173        unique_labels: list of unique labels present in the tracking data
174
175    Returns:
176        dict: {label: track_id} - mapping from each label to its track ID
177    """
178    if not divisions:
179        # No divisions - each unique label is its own track.
180        return {label: label for label in unique_labels}
181
182    ds = DisjointSet(unique_labels)
183
184    # Union connected labels based on mother-daughter relationships.
185    for daughter, mothers in divisions.items():
186        if daughter not in unique_labels:  # weirdly, this can happen
187            continue
188        for mother in mothers:
189            if mother in unique_labels:
190                ds.merge(daughter, mother)
191
192    # A connected component is a TrackMate track. We use the root as track ID.
193    # Create a mapping from label to track_id (root).
194    label_to_track_id = {}
195    for label in unique_labels:
196        root = ds[label]
197        label_to_track_id[label] = root
198
199    return label_to_track_id
200
201
202def build_all_tracks_data(divisions, df_spots):
203    """"""
204    print(f"Divisions: {divisions}")
205    for mothers in divisions.values():
206        if len(mothers) > 1:
207            print("FUSION")
208
209    # Generate and assign TRACK_IDs.
210    labels = df_spots["label"].unique()
211    label_to_track_id = create_label_to_track_mapping(divisions, labels)
212    df_spots["TRACK_ID"] = df_spots["label"].map(label_to_track_id)
213
214    # Division edges: for each daughter-mother pair, create an edge.
215    edges_data = [{"daughter": daughter, "mother": mother} for daughter, mothers in divisions.items() for mother in mothers]
216    df_edges = pd.DataFrame(edges_data)
217    # Labels stay the same until there is a division. But spots ID are unique.
218    # It means that in df_spots, labels appears multiple times. Because of this
219    # we cannot easily map between df_spots and df_edges. So we create intermediary
220    # columns to ease the mapping.
221    df_spots["first_frame"] = df_spots.groupby("label")["FRAME"].transform("min")
222    df_spots["last_frame"] = df_spots.groupby("label")["FRAME"].transform("max")
223    # A daughter is at the first frame of its label, a mother at the last frame of its label.
224    df_spots["daughter"] = df_spots["first_frame"] == df_spots["FRAME"]
225    df_spots["mother"] = df_spots["last_frame"] == df_spots["FRAME"]
226    df_spots.drop(columns=["first_frame", "last_frame"], inplace=True)
227    # Now we can map between df_spots and df_edges.
228    # The SPOT_SOURCE_ID is the spot ID of the matching label that is a mother,
229    # and the SPOT_TARGET_ID is the spot ID of the matching label that is a daughter.
230    df_edges["SPOT_SOURCE_ID"] = df_edges["mother"].map(df_spots[df_spots["mother"]].set_index("label")["ID"])
231    df_edges["SPOT_TARGET_ID"] = df_edges["daughter"].map(df_spots[df_spots["daughter"]].set_index("label")["ID"])
232    df_spots.drop(columns=["daughter", "mother"], inplace=True)
233
234    # Add TRACK_ID to division edges (needed for XML).
235    if not df_edges.empty:
236        df_edges["TRACK_ID"] = df_edges["mother"].map(label_to_track_id)
237        df_edges.drop(columns=["daughter", "mother"], inplace=True)
238
239    # Non-division edges: for each label, connect consecutive spots within that label.
240    non_division_edges = []
241    for label in df_spots["label"].unique():
242        label_spots = df_spots[df_spots["label"] == label].sort_values("FRAME")
243        if len(label_spots) > 1:
244            track_id = label_spots.iloc[0]["TRACK_ID"]
245            for i in range(len(label_spots) - 1):
246                current_spot = label_spots.iloc[i]
247                next_spot = label_spots.iloc[i + 1]
248                non_division_edges.append({"SPOT_SOURCE_ID": current_spot["ID"], "SPOT_TARGET_ID": next_spot["ID"], "TRACK_ID": track_id})
249
250    # Combine division and non-division edges.
251    df_non_division_edges = pd.DataFrame(non_division_edges)
252    if not df_edges.empty and not df_non_division_edges.empty:
253        # Make sure both dataframes have the same columns.
254        df_edges = df_edges[["SPOT_SOURCE_ID", "SPOT_TARGET_ID", "TRACK_ID"]]
255        df_edges = pd.concat([df_edges, df_non_division_edges], ignore_index=True)
256    elif not df_non_division_edges.empty:
257        df_edges = df_non_division_edges
258
259    # Final cleanup and type conversion.
260    if not df_edges.empty:
261        # We can have NaN if a label has no mother (appears at first frame)
262        # or no daughter (disappears at last frame).
263        df_edges.dropna(inplace=True)
264        # Convert to int in case of NaN.
265        df_edges["SPOT_SOURCE_ID"] = df_edges["SPOT_SOURCE_ID"].astype(int)
266        df_edges["SPOT_TARGET_ID"] = df_edges["SPOT_TARGET_ID"].astype(int)
267        df_edges["TRACK_ID"] = df_edges["TRACK_ID"].astype(int)
268
269    return df_edges
270
271
272def build_all_tracks_tag(df_edges):
273    """Build the AllTracks tag for TrackMate XML."""
274    all_tracks = ET.Element("AllTracks")
275    track_ids = sorted(df_edges["TRACK_ID"].unique())
276    track_id_to_index = {track_id: index for index, track_id in enumerate(track_ids)}
277
278    for track_id in track_ids:
279        track = ET.SubElement(all_tracks, "Track")
280        track.set("name", f"Track_{track_id}")
281        track.set("TRACK_ID", str(track_id))
282        track.set("TRACK_INDEX", str(track_id_to_index[track_id]))
283
284        track_edges = df_edges[df_edges["TRACK_ID"] == track_id]
285        for _, edge in track_edges.iterrows():
286            edge_attrib = {
287                "SPOT_SOURCE_ID": str(edge["SPOT_SOURCE_ID"]),
288                "SPOT_TARGET_ID": str(edge["SPOT_TARGET_ID"]),
289            }
290            ET.SubElement(track, "Edge", edge_attrib)
291
292    return all_tracks
293
294
295def build_filtered_tracks_tag(track_ids):
296    """Build the FilteredTracks tag for TrackMate XML."""
297    filtered_tracks = ET.Element("FilteredTracks")
298    for track_id in track_ids:
299        ET.SubElement(filtered_tracks, "TrackID", {"TRACK_ID": str(track_id)})
300    return filtered_tracks
301
302
303def build_model_tag(epic):
304    """Build the Model tag for TrackMate XML."""
305    model = ET.Element("Model")
306    model.set("spatialunits", epic.epi_metadata.get("UnitXY", "pixel"))
307    model.set("timeunits", epic.epi_metadata.get("UnitT", "frame"))
308    model.append(build_feat_declaration_tag())
309
310    print("Tracked?", epic.tracked)
311    df_spots = build_spots_df(epic)
312    cell_contours = build_roi_contours(epic, df_spots)
313    model.append(build_all_spots_tag(df_spots, cell_contours))
314    divisions = epic.tracking.graph  # dict of {daughter: mothers}
315    if divisions:
316        df_edges = build_all_tracks_data(divisions, df_spots)
317        model.append(build_all_tracks_tag(df_edges))
318        track_ids = sorted(df_edges["TRACK_ID"].unique())
319        model.append(build_filtered_tracks_tag(track_ids))
320    else:
321        model.append(ET.Element("AllTracks"))
322        model.append(ET.Element("FilteredTracks"))
323
324    return model
325
326
327def build_settings_tag(epic):
328    """Build the Settings tag for TrackMate XML."""
329    settings = ET.Element("Settings")
330    img_path = Path(epic.epi_metadata["MovieFile"])
331    ET.SubElement(
332        settings,
333        "ImageData",
334        {
335            "filename": img_path.name,
336            "folder": str(img_path.parent),  # TODO: missing / at the end compared to TM, is it an issue?
337            "width": str(epic.imgshape2D[1]),
338            "height": str(epic.imgshape2D[0]),
339            "nslices": "1",
340            "nframes": str(epic.nframes),
341            "pixelwidth": str(epic.epi_metadata.get("ScaleXY", 1)),
342            "pixelheight": str(epic.epi_metadata.get("ScaleXY", 1)),
343            "voxeldepth": "1",
344            "timeinterval": str(epic.epi_metadata.get("ScaleT", 1)),
345        },
346    )
347    ET.SubElement(
348        settings,
349        "BasicSettings",
350        {
351            "xstart": "0",
352            "xend": str(epic.imgshape2D[1] - 1),
353            "ystart": "0",
354            "yend": str(epic.imgshape2D[0] - 1),
355            "zstart": "0",
356            "zend": "0",
357            "tstart": "0",
358            "tend": str(epic.nframes - 1),
359        },
360    )
361    ET.SubElement(settings, "DetectorSettings", {"DETECTOR_NAME": "MANUAL_DETECTOR", "RADIUS": "5"})
362    ET.SubElement(settings, "InitialSpotFilter")
363    ET.SubElement(settings, "SpotFilterCollection")
364    ET.SubElement(settings, "TrackerSettings", {"TRACKER_NAME": "MANUAL_TRACKER"})
365    ET.SubElement(settings, "TrackFilterCollection")
366    return settings
367
368
369# def build_gui_state_tag():
370#     """Build the GUIState tag for TrackMate XML."""
371#     gui_state = ET.Element("GUIState")
372#     return gui_state
373
374
375def pretty_print_xml(element):
376    """Pretty print an XML element."""
377    rough_string = ET.tostring(element, encoding="utf-8")
378    parsed = minidom.parseString(rough_string)
379    return parsed.toprettyxml(indent="  ")
380
381
382def save_trackmate_xml(epic, outname):
383    """Save a TrackMate XML file."""
384    if epic.verbose == 3:
385        ut.show_info(f"ScaleXY: {epic.epi_metadata.get('ScaleXY')}")
386        ut.show_info(f"ScaleT: {epic.epi_metadata.get('ScaleT')}")
387        ut.show_info(f"imgshape2D: {epic.imgshape2D}")
388        ut.show_info(f"UnitXY: {epic.epi_metadata.get('UnitXY')}")
389        ut.show_info(f"UnitT: {epic.epi_metadata.get('UnitT')}")
390
391    root = ET.Element("TrackMate", {"version": "unknown"})
392    log = ET.SubElement(root, "Log")
393    now = datetime.now()
394    log.text = f"Created by EpiCure on {now.strftime('%Y-%m-%d %H:%M:%S')}"
395    model = build_model_tag(epic)
396    root.append(model)
397    settings = build_settings_tag(epic)
398    root.append(settings)
399    ET.SubElement(root, "GUIState", {"state": "ConfigureViews"})
400    ET.SubElement(root, "DisplaySettings")
401    # display_settings = ET.SubElement(root, "DisplaySettings")
402    # display_settings.text = "{'spotDisplayedAsRoi': true}"
403    # tree = ET.ElementTree(root)
404
405    pretty_xml = pretty_print_xml(root)
406    # TODO: check if ET.indent is better or more efficient than pretty_print_xml
407    # for Python 3.9+
408    # Pretty formatting - ET.indent is available from Python 3.9+
409    # try:
410    #     ET.indent(tree, space="  ", level=0)
411    # except AttributeError:
412    #     # Fallback for Python < 3.9
413    #     pass
414    with open(outname, "w", encoding="utf-8") as f:
415        f.write(pretty_xml)