dev(narugo): add state recover

felinewhisker/ui/annotate.py

@@ -1,3 +1,5 @@
+import json
+import os.path
 import pathlib
 from typing import Iterable, Callable, Optional, Any
 
@@ -13,15 +15,50 @@
 
 
 def create_annotation_tab(repo: DatasetRepository, demo: gr.Blocks,
-                          datasource: Iterable[ImageItem], write_session: WriterSession,
-                          fn_annotate_assist: Optional[Callable[[str], Any]] = None, **kwargs):
+                          datasource: Iterable[ImageItem], write_session: WriterSession, state_file: str,
+                          fn_annotate_assist: Optional[Callable[[str], Any]] = None,
+                          **kwargs):
     data_iterator = iter(datasource)
 
     gr_state_output = gr.State(value=None)
     gr_position_id = gr.State(value=-1)
-    gr_datasource_length = gr.State(value=None)
+    gr_max_length = gr.State(value=None)
     gr_id_list = gr.State(value=[])
 
+    def _fn_state_save(position_id, max_length, id_list):
+        with open(state_file, 'w') as f:
+            json.dump({
+                'position_id': position_id,
+                'max_length': max_length,
+                'id_list': id_list,
+            }, f, indent=4, sort_keys=True)
+
+    gr_position_id.change(
+        fn=_fn_state_save,
+        inputs=[gr_position_id, gr_max_length, gr_id_list],
+    )
+    gr_max_length.change(
+        fn=_fn_state_save,
+        inputs=[gr_position_id, gr_max_length, gr_id_list],
+    )
+    gr_id_list.change(
+        fn=_fn_state_save,
+        inputs=[gr_position_id, gr_max_length, gr_id_list],
+    )
+
+    def _fn_load_state():
+        if os.path.exists(state_file):
+            with open(state_file, 'r') as f:
+                state = json.load(f)
+                return state['position_id'], state['max_length'], state['id_list']
+        else:
+            return -1, None, []
+
+    demo.load(
+        fn=_fn_load_state,
+        outputs=[gr_position_id, gr_max_length, gr_id_list],
+    )
+
     with gr.Row():
         gr_state_input = create_ui_for_annotator(
             repo=repo,
@@ -87,8 +124,8 @@ def _fn_next(idx, ids, max_length):
 
     gr_next.click(
         fn=_fn_next,
-        inputs=[gr_position_id, gr_id_list, gr_datasource_length],
-        outputs=[gr_position_id, gr_id_list, gr_datasource_length],
+        inputs=[gr_position_id, gr_id_list, gr_max_length],
+        outputs=[gr_position_id, gr_id_list, gr_max_length],
     )
 
     def _ch_change(state):
@@ -121,7 +158,7 @@ def _fn_index_change(idx, ids, max_length):
 
     gr_position_id.change(
         fn=_fn_index_change,
-        inputs=[gr_position_id, gr_id_list, gr_datasource_length],
+        inputs=[gr_position_id, gr_id_list, gr_max_length],
         outputs=[gr_state_input, gr_prev, gr_next, gr_save],
     )
 

felinewhisker/ui/dispatch.py

@@ -1,9 +1,11 @@
+import os
 import pathlib
 from contextlib import contextmanager
 from typing import Optional, ContextManager, Callable, Any
 
 import gradio as gr
 from hbutils.string import titleize
+from hbutils.system import TemporaryDirectory
 
 from .annotate import create_annotation_tab
 from ..datasource import BaseDataSource
@@ -18,29 +20,31 @@ def create_annotator_app(
         fn_annotate_assist: Optional[Callable[[str], Any]] = None,
         annotation_options: Optional[dict] = None
 ) -> ContextManager[gr.Blocks]:
-    with repo.write(author=author) as write_session:
-        with datasource as source:
-            source.set_fn_contains_id(write_session.is_id_duplicated)
-
-            with gr.Blocks(css=_GLOBAL_CSS_CODE) as demo:
-                with gr.Row(elem_id='annotation_title'):
-                    gr_title = gr.HTML(
-                        f'<p class="title">'
-                        f'<u>{titleize(repo.meta_info["task"])}</u> - {repo.meta_info["name"]}'
-                        f'</p>'
-                    )
-                    _ = gr_title
-
-                with gr.Row():
-                    with gr.Tabs():
-                        with gr.Tab('Annotation'):
-                            create_annotation_tab(
-                                repo=repo,
-                                demo=demo,
-                                datasource=source,
-                                write_session=write_session,
-                                fn_annotate_assist=fn_annotate_assist,
-                                **(annotation_options or {}),
-                            )
-
-            yield demo
+    with TemporaryDirectory(prefix='felinewhisker_') as td_state, \
+            repo.write(author=author) as write_session, datasource as source:
+        state_file = os.path.join(td_state, 'state.json')
+        source.set_fn_contains_id(write_session.is_id_duplicated)
+
+        with gr.Blocks(css=_GLOBAL_CSS_CODE) as demo:
+            with gr.Row(elem_id='annotation_title'):
+                gr_title = gr.HTML(
+                    f'<p class="title">'
+                    f'<u>{titleize(repo.meta_info["task"])}</u> - {repo.meta_info["name"]}'
+                    f'</p>'
+                )
+                _ = gr_title
+
+            with gr.Row():
+                with gr.Tabs():
+                    with gr.Tab('Annotation'):
+                        create_annotation_tab(
+                            repo=repo,
+                            demo=demo,
+                            datasource=source,
+                            write_session=write_session,
+                            state_file=state_file,
+                            fn_annotate_assist=fn_annotate_assist,
+                            **(annotation_options or {}),
+                        )
+
+        yield demo