|
| 1 | +import yaml |
| 2 | +import datajoint as dj |
| 3 | +import numpy as np |
| 4 | +import plotly.express as px |
| 5 | +import plotly.graph_objects as go |
| 6 | +from dash import no_update |
| 7 | +from dash_extensions.enrich import ( |
| 8 | + DashProxy, |
| 9 | + Input, |
| 10 | + Output, |
| 11 | + State, |
| 12 | + html, |
| 13 | + dcc, |
| 14 | + Serverside, |
| 15 | + ServersideOutputTransform, |
| 16 | +) |
| 17 | + |
| 18 | +from .utilities import * |
| 19 | + |
| 20 | + |
| 21 | +logger = dj.logger |
| 22 | + |
| 23 | + |
| 24 | +def draw_rois(db_prefix: str): |
| 25 | + scan = dj.create_virtual_module("scan", f"{db_prefix}scan") |
| 26 | + imaging = dj.create_virtual_module("imaging", f"{db_prefix}imaging") |
| 27 | + all_keys = (imaging.MotionCorrection).fetch("KEY") |
| 28 | + |
| 29 | + colors = {"background": "#111111", "text": "#00a0df"} |
| 30 | + |
| 31 | + app = DashProxy(transforms=[ServersideOutputTransform()]) |
| 32 | + app.layout = html.Div( |
| 33 | + [ |
| 34 | + html.H2("Draw ROIs", style={"color": colors["text"]}), |
| 35 | + html.Label( |
| 36 | + "Select data key from dropdown", style={"color": colors["text"]} |
| 37 | + ), |
| 38 | + dcc.Dropdown( |
| 39 | + id="toplevel-dropdown", options=[str(key) for key in all_keys] |
| 40 | + ), |
| 41 | + html.Br(), |
| 42 | + html.Div( |
| 43 | + [ |
| 44 | + html.Button( |
| 45 | + "Load Image", |
| 46 | + id="load-image-button", |
| 47 | + style={"margin-right": "20px"}, |
| 48 | + ), |
| 49 | + dcc.RadioItems( |
| 50 | + id="image-type-radio", |
| 51 | + options=[ |
| 52 | + {"label": "Average Image", "value": "average_image"}, |
| 53 | + { |
| 54 | + "label": "Max Projection Image", |
| 55 | + "value": "max_projection_image", |
| 56 | + }, |
| 57 | + ], |
| 58 | + value="average_image", |
| 59 | + labelStyle={"display": "inline-block", "margin-right": "10px"}, |
| 60 | + style={"display": "inline-block", "color": colors["text"]}, |
| 61 | + ), |
| 62 | + html.Div( |
| 63 | + [ |
| 64 | + html.Button("Submit Curated Masks", id="submit-button"), |
| 65 | + ], |
| 66 | + style={ |
| 67 | + "textAlign": "right", |
| 68 | + "flex": "1", |
| 69 | + "display": "inline-block", |
| 70 | + }, |
| 71 | + ), |
| 72 | + ], |
| 73 | + style={ |
| 74 | + "display": "flex", |
| 75 | + "justify-content": "flex-start", |
| 76 | + "align-items": "center", |
| 77 | + }, |
| 78 | + ), |
| 79 | + html.Br(), |
| 80 | + html.Br(), |
| 81 | + html.Div( |
| 82 | + [ |
| 83 | + dcc.Graph( |
| 84 | + id="avg-image", |
| 85 | + config={ |
| 86 | + "modeBarButtonsToAdd": [ |
| 87 | + "drawclosedpath", |
| 88 | + "drawrect", |
| 89 | + "drawcircle", |
| 90 | + "drawline", |
| 91 | + "eraseshape", |
| 92 | + ], |
| 93 | + }, |
| 94 | + style={"width": "100%", "height": "100%"}, |
| 95 | + ) |
| 96 | + ], |
| 97 | + style={ |
| 98 | + "display": "flex", |
| 99 | + "justify-content": "center", |
| 100 | + "align-items": "center", |
| 101 | + "padding": "0.0", |
| 102 | + "margin": "auto", |
| 103 | + }, |
| 104 | + ), |
| 105 | + html.Pre(id="annotations"), |
| 106 | + html.Div(id="button-output"), |
| 107 | + dcc.Store(id="store-key"), |
| 108 | + dcc.Store(id="store-mask"), |
| 109 | + dcc.Store(id="store-movie"), |
| 110 | + html.Div(id="submit-output"), |
| 111 | + ] |
| 112 | + ) |
| 113 | + |
| 114 | + @app.callback( |
| 115 | + Output("store-key", "value"), |
| 116 | + Input("toplevel-dropdown", "value"), |
| 117 | + ) |
| 118 | + def store_key(value): |
| 119 | + if value is not None: |
| 120 | + return Serverside(value) |
| 121 | + else: |
| 122 | + return no_update |
| 123 | + |
| 124 | + @app.callback( |
| 125 | + Output("avg-image", "figure"), |
| 126 | + Output("store-movie", "average_images"), |
| 127 | + State("store-key", "value"), |
| 128 | + Input("load-image-button", "n_clicks"), |
| 129 | + Input("image-type-radio", "value"), |
| 130 | + prevent_initial_call=True, |
| 131 | + ) |
| 132 | + def create_figure(value, render_n_clicks, image_type): |
| 133 | + if render_n_clicks is not None: |
| 134 | + if image_type == "average_image": |
| 135 | + summary_images = ( |
| 136 | + imaging.MotionCorrection.Summary & yaml.safe_load(value) |
| 137 | + ).fetch("average_image") |
| 138 | + else: |
| 139 | + summary_images = ( |
| 140 | + imaging.MotionCorrection.Summary & yaml.safe_load(value) |
| 141 | + ).fetch("max_proj_image") |
| 142 | + average_images = [image.astype("float") for image in summary_images] |
| 143 | + roi_contours = get_contours(yaml.safe_load(value), db_prefix) |
| 144 | + logger.info("Generating figure.") |
| 145 | + fig = px.imshow( |
| 146 | + np.asarray(average_images), |
| 147 | + animation_frame=0, |
| 148 | + binary_string=True, |
| 149 | + labels=dict(animation_frame="plane"), |
| 150 | + ) |
| 151 | + for contour in roi_contours: |
| 152 | + # Note: contour[:, 1] are x-coordinates, contour[:, 0] are y-coordinates |
| 153 | + fig.add_trace( |
| 154 | + go.Scatter( |
| 155 | + x=contour[:, 1], # Plotly uses x, y order for coordinates |
| 156 | + y=contour[:, 0], |
| 157 | + mode="lines", # Display as lines (not markers) |
| 158 | + line=dict(color="white", width=0.5), # Set line color and width |
| 159 | + showlegend=False, # Do not show legend for each contour |
| 160 | + ) |
| 161 | + ) |
| 162 | + fig.update_layout( |
| 163 | + dragmode="drawrect", |
| 164 | + autosize=True, |
| 165 | + height=550, |
| 166 | + newshape=dict(opacity=0.6, fillcolor="#00a0df"), |
| 167 | + plot_bgcolor=colors["background"], |
| 168 | + paper_bgcolor=colors["background"], |
| 169 | + font_color=colors["text"], |
| 170 | + ) |
| 171 | + fig.update_annotations(bgcolor="#00a0df") |
| 172 | + else: |
| 173 | + return no_update |
| 174 | + return fig, Serverside(average_images) |
| 175 | + |
| 176 | + @app.callback( |
| 177 | + Output("store-mask", "annotation_list"), |
| 178 | + Input("avg-image", "relayoutData"), |
| 179 | + prevent_initial_call=True, |
| 180 | + ) |
| 181 | + def on_relayout(relayout_data): |
| 182 | + if not relayout_data: |
| 183 | + return no_update |
| 184 | + else: |
| 185 | + if "shapes" in relayout_data: |
| 186 | + global shape_type |
| 187 | + try: |
| 188 | + shape_type = relayout_data["shapes"][-1]["type"] |
| 189 | + return Serverside(relayout_data) |
| 190 | + except IndexError: |
| 191 | + return no_update |
| 192 | + elif any(["shapes" in key for key in relayout_data]): |
| 193 | + return Serverside(relayout_data) |
| 194 | + |
| 195 | + @app.callback( |
| 196 | + Output("submit-output", "children"), |
| 197 | + Input("submit-button", "n_clicks"), |
| 198 | + State("store-mask", "annotation_list"), |
| 199 | + State("store-key", "value"), |
| 200 | + ) |
| 201 | + def submit_annotations(n_clicks, annotation_list, value): |
| 202 | + x_mask_li = [] |
| 203 | + y_mask_li = [] |
| 204 | + if n_clicks is not None: |
| 205 | + if annotation_list: |
| 206 | + if "shapes" in annotation_list: |
| 207 | + logger.info("Creating Masks.") |
| 208 | + shapes = [d["type"] for d in annotation_list["shapes"]] |
| 209 | + for shape, annotation in zip(shapes, annotation_list["shapes"]): |
| 210 | + mask = create_mask(annotation, shape) |
| 211 | + y_mask_li.append(mask[0]) |
| 212 | + x_mask_li.append(mask[1]) |
| 213 | + print("Masks created") |
| 214 | + insert_into_database( |
| 215 | + scan, imaging, yaml.safe_load(value), x_mask_li, y_mask_li |
| 216 | + ) |
| 217 | + else: |
| 218 | + logger.warn( |
| 219 | + "Incorrect annotation list format. This is a known bug. Please draw a line anywhere on the image and click `Submit Curated Masks`. It will be ignored in the final submission but will format the list correctly." |
| 220 | + ) |
| 221 | + return no_update |
| 222 | + else: |
| 223 | + logger.warn("No annotations to submit.") |
| 224 | + return no_update |
| 225 | + else: |
| 226 | + return no_update |
| 227 | + |
| 228 | + return app |
0 commit comments