diff --git a/sample_project/input_config.yaml b/sample_project/input_config.yaml index 095cdbb..2c8dcef 100644 --- a/sample_project/input_config.yaml +++ b/sample_project/input_config.yaml @@ -6,10 +6,10 @@ model_str: DLC_resnet50_jwasp_femaleandmaleSep12shuffle1_1000000 dashboard_export_data_path: ./sample_project/output ROI_tags: - enclosure - - left_box - - top_box - - right_box - - bottom_box + - nectar1_box + - control_box + - water_only_box + - nectar2_box - tube event_tags: - experiment_start diff --git a/sample_project/videos/jwaspE_nectar-open-close_control.metadata.yaml b/sample_project/videos/jwaspE_nectar-open-close_control.metadata.yaml index 9ff8683..c36ec93 100644 --- a/sample_project/videos/jwaspE_nectar-open-close_control.metadata.yaml +++ b/sample_project/videos/jwaspE_nectar-open-close_control.metadata.yaml @@ -23,19 +23,19 @@ ROIs: path: M348.0491395793499,93.62497131931165L356.4009560229445,232.82191204588906L361.9688336520076,480.5924665391969L361.9688336520076,625.3572848948373L361.9688336520076,850.8563288718926L373.1045889101338,959.4299426386231L378.67246653919693,1042.9481070745694L400.9439770554493,1073.5714340344166L690.4736137667304,1073.5714340344166L843.5902485659656,1048.5159847036325L1163.7432122370938,1037.3802294455063L1302.9401529636712,1012.3247801147224L1344.6992351816443,995.6211472275331L1366.9707456978967,917.6708604206498L1378.106500956023,811.881185468451L1389.2422562141492,736.7148374760992L1389.2422562141492,625.3572848948373L1350.2671128107074,491.72822179732304L1336.3474187380498,374.802791586998L1325.2116634799236,288.50068833652006L1322.427724665392,196.63070745697894L1289.0204588910133,65.78558317399617Z - drawn_on_frame: 38000 line_color: '#E15F99' - name: left_box + name: nectar1_box path: M484.23300516875133,565.133184774723L483.3063124386091,605.9076649009826L487.01308335917815,624.4415195038279L489.7931615496049,634.6351395353927L495.3533179304585,639.268603186104L532.4210271361491,639.268603186104L558.3684235801323,635.5618322655349L562.0751945007014,628.1482904243968L560.2218090404169,604.054279440698L557.4417308499901,570.6933411555766L557.4417308499901,563.2797993144385Z - drawn_on_frame: 38000 line_color: '#1CA71C' - name: top_box + name: control_box path: M756.6806678305767,278.34419881335197L761.314131481288,311.7051370984735L767.8009805922838,337.6525335424568L770.5810587827106,343.2126899233104L800.235226147263,337.6525335424568L832.6694717022423,336.7258408123146L834.5228571625267,322.8254498601806L826.1826225912464,286.68443338463237L825.2559298611042,267.22388605164485L826.1826225912464,275.5641206229252Z - drawn_on_frame: 38000 line_color: '#FB0D0D' - name: right_box + name: water-only_box path: M1061.562576047381,566.5456378875958L1071.756196078946,619.3671235057049L1078.2430451899418,633.2675144578387L1104.190441633925,630.487436267412L1143.1115362999,628.6340508071274L1145.891614490327,623.0738944262739L1143.1115362999,602.6866543631442L1136.6246871889043,573.0324869985917L1131.991223538193,554.4986323957464Z - drawn_on_frame: 38000 line_color: '#DA16FF' - name: bottom_box + name: nectar2_box path: M812.2822316391125,845.4801496604169L809.5021534486857,922.3956462622248L859.5435608763679,920.5422608019403L881.7841863997821,911.2753335005176L884.564264590209,894.5948643579569L885.4909573203512,861.2339260728354L882.7108791299245,835.2865296288521Z - drawn_on_frame: 38000 line_color: '#222A2A' diff --git a/wazp/callbacks/roi.py b/wazp/callbacks/roi.py index 690eb6e..d397bad 100644 --- a/wazp/callbacks/roi.py +++ b/wazp/callbacks/roi.py @@ -2,6 +2,7 @@ # import pdb import re +import time from typing import Optional import dash @@ -15,7 +16,6 @@ # TODO: other video extensions? have this in project config file instead? VIDEO_TYPES = [".avi", ".mp4"] -# TODO: make colomap this a project config parameter? ROI_CMAP = px.colors.qualitative.Dark24 @@ -140,6 +140,7 @@ def update_frame_slider( is selected. Read the parameters from storage if available, otherwise extract them from the video file (slower), and update the storage for future use. + Parameters ---------- video_path : str @@ -156,7 +157,7 @@ def update_frame_slider( Returns ------- int - Maximum frame input value. + Maximum frame index value. int Frame step size. int @@ -167,23 +168,36 @@ def update_frame_slider( video_name = pl.Path(video_path).name if video_name in frame_slider_storage.keys(): stored_video_params = frame_slider_storage[video_name] - num_frames = stored_video_params["max"] + max_frame_idx = stored_video_params["max"] frame_step = stored_video_params["step"] middle_frame = stored_video_params["value"] - return num_frames - 1, frame_step, middle_frame, dash.no_update + return max_frame_idx, frame_step, middle_frame, dash.no_update else: - num_frames = int(utils.get_num_frames(video_path)) - # Round the frame step to the nearest 1000 - frame_step = round(int(num_frames / 4), -3) + try: + num_frames = utils.get_num_frames(video_path) + except RuntimeError as e: + print(e) + # If the number of frames cannot be extracted, + # return a negative frame value. + # This will trigger an alert message in the app + return dash.no_update, dash.no_update, -1, dash.no_update + + # Divide the number of frames into 4 steps + frame_step = int(num_frames / 4) + # Round to the nearest 1000 if step is > 1000 + if frame_step > 1000: + frame_step = int(frame_step / 1000) * 1000 + # Default to the middle step middle_frame = frame_step * 2 + max_frame_idx = num_frames - 1 frame_slider_storage[video_name] = { - "max": num_frames - 1, + "max": max_frame_idx, "step": frame_step, "value": middle_frame, } return ( - num_frames - 1, + max_frame_idx, frame_step, middle_frame, frame_slider_storage, @@ -268,30 +282,42 @@ def set_roi_color_in_table( return cond_format @app.callback( - Output("roi-storage", "data"), + [ + Output("roi-storage", "data"), + Output("roi-table", "selected_rows"), + ], [ Input("frame-graph", "relayoutData"), Input("load-rois-button", "n_clicks"), + Input("delete-rois-button", "n_clicks"), ], [ State("video-select", "value"), State("frame-slider", "value"), State("roi-storage", "data"), State("roi-colors-storage", "data"), + State("roi-table", "data"), + State("roi-table", "selected_rows"), ], ) def update_roi_storage( graph_relayout: dict, load_clicks: int, + delete_clicks: int, video_path: str, frame_num: int, roi_storage: dict, roi_color_mapping: dict, - ) -> dict: + roi_table_rows: list, + roi_table_selected_rows: list, + ) -> tuple[dict, list]: """ - Update the ROI storage with the latest ROI shapes - drawn on the frame graph or with the ROI shapes - loaded from the video's .metadata.yaml file. + Update the ROI storage, when: + - Shapes are added/removed from the frame graph + - Shapes are edited on the frame graph + - Shapes are loaded from file + - Shapes are deleted from the ROI table + Parameters ---------- graph_relayout : dict @@ -299,6 +325,8 @@ def update_roi_storage( changes to the frame graph. load_clicks : int Number of times the load ROIs button has been clicked. + delete_clicks : int + Number of times the delete ROIs button has been clicked. video_path : str Path to the video file. frame_num : int @@ -309,10 +337,17 @@ def update_roi_storage( Dictionary with the following keys: - roi2color: dict mapping ROI names to colors - color2roi: dict mapping colors to ROI names + roi_table_rows : list + List of dictionaries with ROI table data. + roi_table_selected_rows : list + List of indices for the selected rows in the ROI table. + Returns ------- dict Updated dictionary storing ROI data for each video. + list + List of indices for the selected rows in the ROI table. """ trigger = dash.callback_context.triggered[0]["prop_id"] @@ -386,7 +421,24 @@ def update_roi_storage( yaml_path=metadata_path ) - return roi_storage + # If triggered by the delete ROIs button click + # Delete the selected ROIs from the roi-storage + elif trigger == "delete-rois-button.n_clicks": + if delete_clicks > 0 and roi_table_selected_rows: + deleted_roi_names = [ + roi_table_rows[idx]["name"] + for idx in roi_table_selected_rows + ] + stored_shapes = roi_storage[video_name]["shapes"] + roi_storage[video_name]["shapes"] = [ + sh + for sh in stored_shapes + if sh["roi_name"] not in deleted_roi_names + ] + # Clear the row selection + roi_table_selected_rows = [] + + return roi_storage, roi_table_selected_rows @app.callback( [ @@ -404,24 +456,27 @@ def update_roi_storage( [ State("frame-graph", "figure"), State("roi-colors-storage", "data"), + State("frame-slider", "max"), ], ) def update_frame_graph( video_path: str, - frame_num: int, + shown_frame_idx: int, roi_name: str, roi_storage: dict, current_fig: go.Figure, roi_color_mapping: dict, + max_frame_idx: int, ) -> tuple[go.Figure, str, str, bool]: """ Update the frame graph + Parameters ---------- video_path : str Path to the video file. - frame_num : int - Frame number (which video frame to display). + shown_frame_idx : int + Index of the frame to be shown. roi_name : str Name of the next ROI to be drawn. roi_storage : dict @@ -432,6 +487,9 @@ def update_frame_graph( Dictionary with the following keys: - roi2color: dict mapping ROI names to colors - color2roi: dict mapping colors to ROI names + max_frame_idx : int + Maximum frame index (num_frames - 1) + Returns ------- plotly.graph_objects.Figure.Figure @@ -442,8 +500,18 @@ def update_frame_graph( Color of the frame status alert. bool Whether to open the frame status alert. + int + Maximum frame index (num_frames - 1) """ + # If a negative frame index is passed, it means that the video + # could not be read correctly. So don't update the frame, + # but display an alert message + if shown_frame_idx < 0: + alert_msg = f"Could not read from '{video_path}'. " + alert_msg += "Is this a valid video file?" + return dash.no_update, alert_msg, "danger", True + # Get the video path and file name video_path_pl = pl.Path(video_path) video_name = video_path_pl.name @@ -478,36 +546,83 @@ def update_frame_graph( # Load the frame into a new figure else: try: - frame_filepath = utils.cache_frame(video_path_pl, frame_num) - new_frame = Image.open(frame_filepath) - # Put the frame in a figure - new_fig = px.imshow(new_frame) - # Add the stored shapes and set the next ROI color - new_fig.update_layout( - shapes=graph_shapes, - newshape_line_color=next_shape_color, - dragmode="drawclosedpath", - margin=dict(l=0, r=0, t=0, b=0), - yaxis={"visible": False, "showticklabels": False}, - xaxis={"visible": False, "showticklabels": False}, + frame_filepath = utils.cache_frame( + video_path_pl, shown_frame_idx ) - alert_msg = f"Showing frame {frame_num} from {video_name}." - alert_color = "success" - alert_open = True - return new_fig, alert_msg, alert_color, alert_open - except Exception as e: - print(e) - alert_msg = ( - f"Could not extract frames from {video_name}. " - f"Make sure that it is a valid video file." - ) - alert_color = "danger" - alert_open = True - return dash.no_update, alert_msg, alert_color, alert_open + except RuntimeError as e: + return dash.no_update, str(e), "danger", True + + new_frame = Image.open(frame_filepath) + new_fig = px.imshow(new_frame) + # Add the stored shapes and set the nextROI color + new_fig.update_layout( + shapes=graph_shapes, + newshape_line_color=next_shape_color, + dragmode="drawclosedpath", + margin=dict(l=0, r=0, t=0, b=0), + yaxis={"visible": False, "showticklabels": False}, + xaxis={"visible": False, "showticklabels": False}, + ) + alert_msg = f"Showing frame {shown_frame_idx}/{max_frame_idx}" + return new_fig, alert_msg, "light", True + + @app.callback( + Output("save-rois-button", "download"), + Input("save-rois-button", "n_clicks"), + [ + State("video-select", "value"), + State("roi-storage", "data"), + ], + ) + def save_rois_to_file( + save_clicks: int, + video_path: str, + roi_storage: dict, + ) -> str: + """ + Save the ROI shapes to a metadata YAML file. + + Parameters + ---------- + save_clicks : int + Number of times the save ROIs button has been clicked. + video_path : str + Path to the video file. + roi_storage : dict + Dictionary storing ROI shapes in the app. + + Returns + ------- + str + Download link for the metadata YAML file. + """ + if save_clicks > 0: + # Get the paths to the video and metadata files + video_path_pl = pl.Path(video_path) + video_name = video_path_pl.name + metadata_filepath = video_path_pl.with_suffix(".metadata.yaml") + + # Get the metadata from the YAML file + with open(metadata_filepath, "r") as yaml_file: + metadata = yaml.safe_load(yaml_file) + + # Get the video's ROI shapes in the app + rois_in_app = roi_storage[video_name]["shapes"] + # Add the ROI shapes to the metadata and save + with open(metadata_filepath, "w") as yaml_file: + metadata["ROIs"] = [ + utils.stored_shape_to_yaml_entry(shape) + for shape in rois_in_app + ] + yaml.safe_dump(metadata, yaml_file, sort_keys=False) + + # Return the download link + return metadata_filepath.as_posix() + else: + return dash.no_update @app.callback( [ - Output("save-rois-button", "download"), Output("rois-status-alert", "children"), Output("rois-status-alert", "color"), Output("rois-status-alert", "is_open"), @@ -518,26 +633,26 @@ def update_frame_graph( Input("video-select", "value"), ], ) - def save_rois_and_update_status_alert( + def update_roi_status_alert( save_clicks: int, roi_storage: dict, video_path: str, - ) -> tuple[str, str, str, bool]: + ) -> tuple[str, str, bool]: """ - Save the ROI shapes to a metadata YAML file - and update the ROI status alert accordingly. + Update the ROI status alert to inform the user about ROIs defined + in the app and their relation to those saved in the metadata file. + Parameters ---------- save_clicks : int Number of times the save ROIs button has been clicked. roi_storage : dict - Dictionary storing already drawn ROI shapes. + Dictionary storing ROI shapes in the app. video_path : str Path to the video file. + Returns ------- - str - Download path to the metadata YAML file. str Message to display in the ROI status alert. str @@ -550,62 +665,154 @@ def save_rois_and_update_status_alert( # Get the paths to the video and metadata files video_path_pl = pl.Path(video_path) video_name = video_path_pl.name - metadata_filepath = video_path_pl.with_suffix(".metadata.yaml") - - # Check if the metadata file exists - # and load any previously saved ROIs - if metadata_filepath.exists(): - with open(metadata_filepath, "r") as yaml_file: - metadata = yaml.safe_load(yaml_file) - if "ROIs" in metadata.keys(): - saved_rois = metadata["ROIs"] - else: - saved_rois = [] - else: - alert_msg = f"Could not find {metadata_filepath.name}" + metadata_path = video_path_pl.with_suffix(".metadata.yaml") + + # If triggered by a click on the save ROIs button + # return with a success message + if trigger == "save-rois-button.n_clicks": + alert_msg = f"Saved ROIs to '{metadata_path.name}'" + alert_color = "success" + return alert_msg, alert_color, True + + # Load saved ROIs from the metadata file, if any + try: + rois_in_file = utils.load_rois_from_yaml(metadata_path) + except FileNotFoundError: + alert_msg = f"Could not find {metadata_path.name}" alert_color = "danger" - return dash.no_update, alert_msg, alert_color, True + return alert_msg, alert_color, True + except KeyError: + rois_in_file = [] - # Get the stored ROI shapes for this video + # Get the app's ROI shapes for this video + rois_in_app = [] if video_name in roi_storage.keys(): - roi_shapes = roi_storage[video_name]["shapes"] - rois_to_save = [ - utils.stored_shape_to_yaml_entry(shape) for shape in roi_shapes - ] - else: - rois_to_save = [] + rois_in_app = roi_storage[video_name]["shapes"] - if rois_to_save: - if trigger == "roi-storage.data" and rois_to_save == saved_rois: - # This means that the ROIs have just been loaded - alert_msg = f"Loaded ROIs from {metadata_filepath.name}" - alert_color = "success" - return dash.no_update, alert_msg, alert_color, True - elif trigger == "roi-storage.data" and rois_to_save != saved_rois: - # This means that the ROIs have been modified - # in respect to the metadata file - alert_msg = "Detected unsaved changes to ROIs." - alert_color = "warning" - return dash.no_update, alert_msg, alert_color, True + if not rois_in_app: + alert_color = "light" + if rois_in_file: + alert_msg = ( + f"Found {len(rois_in_file)} ROIs in '{metadata_path.name}'" + ) else: - if save_clicks > 0: - # This means that the user wants to save the ROIs - # This overwrites any previously saved ROIs - with open(metadata_filepath, "w") as yaml_file: - metadata["ROIs"] = rois_to_save - yaml.safe_dump(metadata, yaml_file) - alert_msg = f"Saved ROIs to {metadata_filepath.name}" - alert_color = "success" - return ( - metadata_filepath.as_posix(), - alert_msg, - alert_color, - True, - ) - else: - return dash.no_update, dash.no_update, dash.no_update, True + alert_msg = "No ROIs defined in the metadata file." else: - alert_msg = "No ROIs to save." - alert_color = "light" - return dash.no_update, alert_msg, alert_color, True + # Some ROIs exist in the app + if rois_in_app == rois_in_file: + alert_color = "success" + if trigger == "roi-storage.data": + alert_msg = f"Loaded ROIs from '{metadata_path.name}'" + else: + alert_msg = ( + f"Shown ROIs match those in '{metadata_path.name}'" + ) + + else: + alert_color = "warning" + alert_msg = "Detected unsaved changes to ROIs." + + return alert_msg, alert_color, True + + @app.callback( + Output("save-rois-button", "disabled"), + Input("roi-storage", "data"), + Input("video-select", "value"), + ) + def disable_save_rois_button( + roi_storage: dict, + video_path: str, + ) -> bool: + """If there are no ROIs in the app or if the metadata file + does not exist, disable the save ROIs button. + + Parameters + ---------- + roi_storage : dict + Dictionary storing ROI shapes in the app. + video_path : str + Path to the selected video file. + + Returns + ------- + bool + Whether to enable the save ROIs button. + """ + + video_path_pl = pl.Path(video_path) + video_name = video_path_pl.name + metadata_path = video_path_pl.with_suffix(".metadata.yaml") + + rois_in_app = [] + if video_name in roi_storage.keys(): + rois_in_app = roi_storage[video_name]["shapes"] + + no_rois_to_save = len(rois_in_app) == 0 + metadata_file_does_not_exist = not metadata_path.is_file() + return no_rois_to_save or metadata_file_does_not_exist + + @app.callback( + Output("load-rois-button", "disabled"), + [ + Input("save-rois-button", "n_clicks"), + Input("video-select", "value"), + ], + ) + def disable_load_rois_button( + save_clicks: int, + video_path: str, + ) -> bool: + """If there are no ROIs saved in the metadata file, + disable the 'Load all from file' button. + + Parameters + ---------- + save_clicks : int + Number of times the save ROIs button has been clicked. + video_path : str + Path to the selected video file. + + Returns + ------- + bool + Whether to enable the load ROIs button. + """ + + video_path_pl = pl.Path(video_path) + metadata_path = video_path_pl.with_suffix(".metadata.yaml") + + # If triggered by a click on the save ROIs button, + # wait a few seconds before checking for ROIs in the file + trigger = dash.callback_context.triggered[0]["prop_id"] + if trigger == "save-rois-button.n_clicks" and save_clicks > 1: + time.sleep(2) + + try: + saved_shapes = utils.load_rois_from_yaml(yaml_path=metadata_path) + return len(saved_shapes) == 0 + except (FileNotFoundError, KeyError): + return True + + @app.callback( + Output("delete-rois-button", "disabled"), + Input("roi-table", "selected_rows"), + ) + def disable_delete_rois_button( + selected_rows: list, + ) -> bool: + """If there are no ROIs selected in the ROI table, + disable the 'Delete selected' button + + + Parameters + ---------- + selected_rows : list + List of selected rows. + + Returns + ------- + bool + Whether to enable the delete ROIs button. + """ + return len(selected_rows) == 0 diff --git a/wazp/pages/02_ROI.py b/wazp/pages/02_ROI.py index 13846a0..1458db2 100644 --- a/wazp/pages/02_ROI.py +++ b/wazp/pages/02_ROI.py @@ -31,18 +31,6 @@ # Initialize the ROI status alert init_roi_status: dict = {"message": "No ROIs to save.", "color": "light"} -# Instructions for the user -instructions = ( - "#### Instructions\n" - "1. Select the video from the top dropdown menu. \n" - "3. Select an ROI from the right dropdown menu. \n" - "4. Draw the ROI on the frame (adjust the shown frame if necessary). \n" - "5. You may also select an existing ROI " - "to edit it or delete it.\n" - "6. Repeat steps 3-5 for each ROI.\n" - "7. Save the ROIs and move on to the next video.\n" -) - ############################### # Graph showing a video frame # @@ -115,6 +103,7 @@ id="roi-table", columns=[dict(name=c, id=c) for c in init_roi_table_columns], data=[], + selected_rows=[], editable=False, style_data={"height": 40}, style_cell={ @@ -125,6 +114,7 @@ }, style_data_conditional=[], fill_width=True, + row_selectable="multi", ) # Dropdown for ROI selection @@ -137,44 +127,37 @@ ) # Buttons for saving/loading ROIs -roi_save_button = dbc.Button( - "Save ROIs", +disabled_button_style = { + "n_clicks": 0, + "outline": False, + "color": "dark", + "disabled": True, + "class_name": "w-100", +} + +save_rois_button = dbc.Button( + "Save all", id="save-rois-button", download="rois.yaml", - n_clicks=0, - outline=False, - color="dark", - active=True, + **disabled_button_style, ) -roi_load_button = dbc.Button( - "Load ROIs", - id="load-rois-button", - n_clicks=0, - outline=False, - color="dark", - active=True, +load_rois_button = dbc.Button( + "Load all", id="load-rois-button", **disabled_button_style ) -infer_rois_button = dbc.Button( - "Infer ROIs", - id="infer-rois-button", - n_clicks=0, - outline=True, - color="dark", - active=False, +delete_rois_button = dbc.Button( + "Delete selected", id="delete-rois-button", **disabled_button_style ) # Tooltips for ROI buttons save_rois_tooltip = dbc.Tooltip( - "Save the ROIs to the video's " ".metadata.yaml file", + "Save all ROIs to the video's .metadata.yaml file. " + "This will overwrite any existing ROIs in the file!", target="save-rois-button", ) load_rois_tooltip = dbc.Tooltip( - "Load ROIs from the video's metadata.yaml file", + "Load all ROIs from the video's metadata.yaml file. " + "This will overwrite any existing ROIs in the app!", target="load-rois-button", ) -infer_rois_tooltip = dbc.Tooltip( - "NOT IMPLEMENTED YET! " "Infer ROI positions based on defined ROIs", - target="infer-rois-button", -) # ROI status alert roi_status_alert = dbc.Alert( @@ -212,49 +195,26 @@ dbc.Col(dcc.Loading(frame_slider), width=9), ] ), + dbc.Row( + [ + dbc.Col(dcc.Markdown("Draw ROI for"), width=3), + dbc.Col(dcc.Loading(roi_dropdown), width=9), + ] + ), ] ), dbc.CardBody(frame_graph), - dbc.CardFooter( - dbc.Row( - [ - dcc.Loading(frame_status_alert), - dcc.Markdown(instructions), - ], - ) - ), + dbc.CardFooter(dcc.Loading(frame_status_alert)), ], ) # ROI table card table_card = dbc.Card( [ - dbc.CardHeader( - dbc.Row( - [ - dbc.Col(roi_save_button, width=4), - dbc.Col(roi_load_button, width=4), - dbc.Col(infer_rois_button, width=4), - save_rois_tooltip, - load_rois_tooltip, - infer_rois_tooltip, - ], - ), - ), + dbc.CardHeader(html.H3("Defined ROIs")), dbc.CardBody( [ - dbc.Row(dbc.Col(html.H4("Defined ROIs"))), dbc.Row(dbc.Col(roi_table)), - dbc.Row( - dbc.Col( - [ - html.Br(), - html.H4("Create new ROI for"), - roi_dropdown, - ], - align="center", - ) - ), dcc.Store(data={}, id="roi-colors-storage"), dcc.Store( id="roi-storage", @@ -263,7 +223,21 @@ ), ] ), - dbc.CardFooter(roi_status_alert), + dbc.CardFooter( + [ + dbc.Row( + [ + dbc.Col(delete_rois_button, width={"size": "auto"}), + dbc.Col(save_rois_button, width={"size": "auto"}), + dbc.Col(load_rois_button, width={"size": "auto"}), + save_rois_tooltip, + load_rois_tooltip, + ], + ), + html.Br(), + dcc.Loading(dbc.Row(roi_status_alert)), + ] + ), ] ) diff --git a/wazp/utils.py b/wazp/utils.py index 91f760a..4bc02eb 100644 --- a/wazp/utils.py +++ b/wazp/utils.py @@ -529,7 +529,7 @@ def assign_roi_colors( } -def get_num_frames(video_path): +def get_num_frames(video_path) -> int: """ Get the number of frames in a video. @@ -543,13 +543,17 @@ def get_num_frames(video_path): int Number of frames in the video """ - vidcap = cv2.VideoCapture(video_path) - return int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT)) + vidcap = cv2.VideoCapture(video_path, apiPreference=cv2.CAP_FFMPEG) + num_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT)) + if num_frames < 1: + raise RuntimeError( + f"Could not read from '{video_path}'. " + "Is this a valid video file?" + ) + return num_frames -def extract_frame( - video_path: str, frame_number: int, output_path: str -) -> None: +def extract_frame(video_path: str, frame_idx: int, output_path: str) -> None: """ Extract a single frame from a video and save it. @@ -557,25 +561,28 @@ def extract_frame( ---------- video_path : str Path to the video file - frame_number : int - Number of the frame to extract + frame_idx : int + Index of the frame to extract output_path : str Path to the output image file """ - print(f"Extracting frame {frame_number} from video {video_path}") - vidcap = cv2.VideoCapture(video_path) - vidcap.set(cv2.CAP_PROP_POS_FRAMES, frame_number) + print(f"Extracting frame {frame_idx} from video {video_path}") + + vidcap = cv2.VideoCapture(video_path, apiPreference=cv2.CAP_FFMPEG) + vidcap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx) success, image = vidcap.read() if success: cv2.imwrite(output_path, image) - print(f"Saved frame to {output_path}") + print(f"Saved frame {frame_idx} to {output_path}") else: - print("Error extracting frame from video") + raise RuntimeError( + f"Could not extract frame {frame_idx} from {video_path}." + ) def cache_frame( video_path: pl.Path, - frame_num: int, + frame_idx: int, cache_dir: pl.Path = pl.Path.home() / ".WAZP" / "roi_frames", frame_suffix: str = "png", ) -> pl.Path: @@ -586,8 +593,8 @@ def cache_frame( ---------- video_path : pl.Path Path to the video file - frame_num : int - Number of the frame to extract + frame_idx : int + Index of the frame to extract cache_dir : pl.Path Path to the cache directory frame_suffix : str, optional @@ -601,12 +608,12 @@ def cache_frame( cache_dir.mkdir(parents=True, exist_ok=True) frame_filepath = ( - cache_dir / f"{video_path.stem}_frame-{frame_num}.{frame_suffix}" + cache_dir / f"{video_path.stem}_frame-{frame_idx}.{frame_suffix}" ) # Extract frame if it is not already cached if not frame_filepath.exists(): extract_frame( - video_path.as_posix(), frame_num, frame_filepath.as_posix() + video_path.as_posix(), frame_idx, frame_filepath.as_posix() ) # Remove old frames from cache remove_old_frames_from_cache( @@ -754,12 +761,12 @@ def load_rois_from_yaml(yaml_path: pl.Path) -> list: if yaml_path.exists(): with open(yaml_path, "r") as yaml_file: metadata = yaml.safe_load(yaml_file) - if "ROIs" in metadata.keys(): + if "ROIs" in metadata: shapes_to_store = [ yaml_entry_to_stored_shape(roi) for roi in metadata["ROIs"] ] else: - raise ValueError(f"Could not find key 'ROIs' in {yaml_path}") + raise KeyError(f"Could not find key 'ROIs' in {yaml_path}") else: raise FileNotFoundError(f"Could not find {yaml_path}")