dev(narugo): optimize ui

felinewhisker/tasks/classification/ui.py

@@ -8,6 +8,8 @@
 ]
 _DEFAULT = object()
 
+_HOTKEY_EMOJIS = set((str(i) for i in range(1, 10)))
+
 
 def create_ui_for_classification(repo, block: gr.Blocks, gr_output_state: gr.State, hotkey_maps=_DEFAULT):
     from ...repository import DatasetRepository
@@ -30,10 +32,17 @@ def create_ui_for_classification(repo, block: gr.Blocks, gr_output_state: gr.Sta
 
                     btn_label_ids = [f'btn_label_{label}' for i, label in enumerate(labels)]
                     btns = [gr.Button(
-                        value=f'{label} ({hotkey_maps[i].upper()})',
+                        value=(
+                            f'({hotkey_maps[i].upper()}) {label}'
+                            if hotkey_maps[i] not in _HOTKEY_EMOJIS else label
+                        ),
                         elem_id=btn_id,
                         elem_classes='btn-label',
                         interactive=False,
+                        icon=(
+                            emoji_image_file(f':keycap_{hotkey_maps[i]}:')
+                            if hotkey_maps[i] in _HOTKEY_EMOJIS else None
+                        ),
                     ) for i, (label, btn_id) in enumerate(zip(labels, btn_label_ids))]
 
                     def _annotation_transition(annotation, triggered_state):
@@ -63,19 +72,12 @@ def _annotation_transition(annotation, triggered_state):
 
                 def _annotation_changed(current_position_id, state):
                     new_btns = [
-                        gr.Button(
-                            value=f'{label} ({hotkey_maps[i].upper()})',
-                            elem_id=btn_id,
+                        gr.update(
                             elem_classes='btn-label btn-selected' if state and label == state else 'btn-label',
                             interactive=True,
                         ) for i, (label, btn_id) in enumerate(zip(labels, btn_label_ids))
                     ]
-                    unannotate_button = gr.Button(
-                        value='Unannotate',
-                        elem_id='btn_unannoate_label',
-                        icon=emoji_image_file(':no_entry:'),
-                        interactive=True,
-                    )
+                    unannotate_button = gr.update(interactive=True if state else False)
                     if state:
                         state_html = f'<p>Current Sample: #{current_position_id}</p>' \
                                      f'<p>Annotated: <b>{state}</b></p>'

felinewhisker/ui/annotate.py

@@ -0,0 +1,128 @@
+import pathlib
+from typing import Iterable
+
+import gradio as gr
+from hbutils.string import plural_word
+
+from ..datasource import ImageItem
+from ..repository import DatasetRepository, WriterSession
+from ..tasks import create_ui_for_annotator
+from ..utils import emoji_image_file
+
+_HOTKEY_JS_CODE = (pathlib.Path(__file__).parent / 'hotkeys.js').read_text()
+
+
+def create_annotation_tab(repo: DatasetRepository, demo: gr.Blocks,
+                          datasource: Iterable[ImageItem], write_session: WriterSession, **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_id_list = gr.State(value=[])
+
+    with gr.Row():
+        gr_state_input = create_ui_for_annotator(
+            repo=repo,
+            block=demo,
+            gr_output_state=gr_state_output,
+            **kwargs
+        )
+
+    with gr.Row():
+        gr_prev = gr.Button(
+            value='Prev',
+            elem_id='left-button',
+            icon=emoji_image_file(':left_arrow:'),
+            interactive=False,
+        )
+        gr_next = gr.Button(
+            value='Next',
+            elem_id='right-button',
+            icon=emoji_image_file(':right_arrow:'),
+        )
+        gr_save = gr.Button(
+            value='Save (Ctrl+S)',
+            elem_id='save-button',
+            icon=emoji_image_file(':floppy_disk:'),
+            interactive=False,
+        )
+
+    def _fn_prev(idx, ids):
+        if idx <= 0:
+            raise gr.Error('This is the first image, no previous sample.')
+        idx -= 1
+        return idx, ids
+
+    gr_prev.click(
+        fn=_fn_prev,
+        inputs=[gr_position_id, gr_id_list],
+        outputs=[gr_position_id, gr_id_list],
+    )
+
+    def _fn_next(idx, ids, max_length):
+        origin_idx = idx
+        idx += 1
+        if idx >= len(ids):
+            try:
+                gr.Info(f'Loading image #{idx} ...')
+                item: ImageItem = next(data_iterator)
+            except StopIteration:
+                gr.Warning('No more images in the data source, '
+                           'you have met the end.')
+                idx = origin_idx
+                max_length = len(ids)
+            else:
+                sample_id = item.id
+                with item.make_file(force_reencode=True) as image_file:
+                    write_session.add(
+                        id_=item.id,
+                        image_file=image_file,
+                        annotation=item.annotation,
+                    )
+                ids.append(sample_id)
+
+        return 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],
+    )
+
+    def _ch_change(state):
+        id_, annotation = state
+        write_session[id_] = annotation
+        print(f'{id_} --> {state!r}')
+
+    gr_state_output.change(
+        fn=_ch_change,
+        inputs=[gr_state_output],
+    )
+
+    def _fn_index_change(idx, ids, max_length):
+        sample_id = ids[idx]
+        annotation = write_session[sample_id]
+        image_file = write_session.get_image_path(sample_id)
+        return (idx, sample_id, image_file, annotation), \
+            gr.update(interactive=idx > 0), \
+            gr.update(interactive=max_length is None or idx < max_length - 1), \
+            gr.update(interactive=True)
+
+    gr_position_id.change(
+        fn=_fn_index_change,
+        inputs=[gr_position_id, gr_id_list, gr_datasource_length],
+        outputs=[gr_state_input, gr_prev, gr_next, gr_save],
+    )
+
+    def _fn_save():
+        save_count = write_session.get_annotated_count()
+        gr.Info(f'Saving {plural_word(save_count, "annotated sample")} ...')
+        write_session.save()
+        gr.Info(f'{plural_word(save_count, "sample")} saved!')
+
+    gr_save.click(
+        fn=_fn_save,
+    )
+
+    demo.load(None, js=_HOTKEY_JS_CODE)

felinewhisker/ui/dispatch.py

@@ -3,24 +3,22 @@
 from typing import Optional, ContextManager
 
 import gradio as gr
-from hbutils.string import titleize, plural_word
+from hbutils.string import titleize
 
-from ..datasource import BaseDataSource, ImageItem
+from .annotate import create_annotation_tab
+from ..datasource import BaseDataSource
 from ..repository import DatasetRepository
-from ..tasks import create_ui_for_annotator
-from ..utils import emoji_image_file
 
 _GLOBAL_CSS_CODE = (pathlib.Path(__file__).parent / 'global.css').read_text()
-_HOTKEY_JS_CODE = (pathlib.Path(__file__).parent / 'hotkeys.js').read_text()
 
 
 @contextmanager
-def create_annotator_app(repo: DatasetRepository, datasource: BaseDataSource, author: Optional[str] = None) \
+def create_annotator_app(repo: DatasetRepository, datasource: BaseDataSource, author: Optional[str] = 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)
-            data_iterator = iter(source)
 
             with gr.Blocks(css=_GLOBAL_CSS_CODE) as demo:
                 with gr.Row(elem_id='annotation_title'):
@@ -34,114 +32,12 @@ def create_annotator_app(repo: DatasetRepository, datasource: BaseDataSource, au
                 with gr.Row():
                     with gr.Tabs():
                         with gr.Tab('Annotation'):
-                            gr_state_output = gr.State(value=None)
-                            gr_position_id = gr.State(value=-1)
-                            gr_datasource_length = gr.State(value=None)
-                            gr_id_list = gr.State(value=[])
-
-                            with gr.Row():
-                                gr_state_input = create_ui_for_annotator(
-                                    repo=repo,
-                                    block=demo,
-                                    gr_output_state=gr_state_output,
-                                )
-
-                            with gr.Row():
-                                gr_prev = gr.Button(
-                                    value='Prev',
-                                    elem_id='left-button',
-                                    icon=emoji_image_file(':left_arrow:'),
-                                    interactive=False,
-                                )
-                                gr_next = gr.Button(
-                                    value='Next',
-                                    elem_id='right-button',
-                                    icon=emoji_image_file(':right_arrow:'),
-                                )
-                                gr_save = gr.Button(
-                                    value='Save (Ctrl+S)',
-                                    elem_id='save-button',
-                                    icon=emoji_image_file(':floppy_disk:'),
-                                    interactive=False,
-                                )
-
-                            def _fn_prev(idx, ids):
-                                if idx <= 0:
-                                    raise gr.Error('This is the first image, no previous sample.')
-                                idx -= 1
-                                return idx, ids
-
-                            gr_prev.click(
-                                fn=_fn_prev,
-                                inputs=[gr_position_id, gr_id_list],
-                                outputs=[gr_position_id, gr_id_list],
-                            )
-
-                            def _fn_next(idx, ids, max_length):
-                                origin_idx = idx
-                                idx += 1
-                                if idx >= len(ids):
-                                    try:
-                                        gr.Info(f'Loading image #{idx} ...')
-                                        item: ImageItem = next(data_iterator)
-                                    except StopIteration:
-                                        gr.Warning('No more images in the data source, '
-                                                   'you have met the end.')
-                                        idx = origin_idx
-                                        max_length = len(ids)
-                                    else:
-                                        sample_id = item.id
-                                        with item.make_file(force_reencode=True) as image_file:
-                                            write_session.add(
-                                                id_=item.id,
-                                                image_file=image_file,
-                                                annotation=item.annotation,
-                                            )
-                                        ids.append(sample_id)
-
-                                return 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],
-                            )
-
-                            def _ch_change(state):
-                                id_, annotation = state
-                                write_session[id_] = annotation
-                                print(f'{id_} --> {state!r}')
-
-                            gr_state_output.change(
-                                fn=_ch_change,
-                                inputs=[gr_state_output],
-                            )
-
-                            def _fn_index_change(idx, ids, max_length):
-                                sample_id = ids[idx]
-                                annotation = write_session[sample_id]
-                                image_file = write_session.get_image_path(sample_id)
-                                return (idx, sample_id, image_file, annotation), \
-                                    gr.update(interactive=idx > 0), \
-                                    gr.update(interactive=max_length is None or idx < max_length - 1), \
-                                    gr.update(interactive=True)
-
-                            gr_position_id.change(
-                                fn=_fn_index_change,
-                                inputs=[gr_position_id, gr_id_list, gr_datasource_length],
-                                outputs=[gr_state_input, gr_prev, gr_next, gr_save],
+                            create_annotation_tab(
+                                repo=repo,
+                                demo=demo,
+                                datasource=source,
+                                write_session=write_session,
+                                **(annotation_options or {}),
                             )
 
-                            def _fn_save():
-                                save_count = write_session.get_annotated_count()
-                                gr.Info(f'Saving {plural_word(save_count, "annotated sample")} ...')
-                                write_session.save()
-                                gr.Info(f'{plural_word(save_count, "sample")} saved!')
-
-                            gr_save.click(
-                                fn=_fn_save,
-                            )
-
-                            demo.load(None, js=_HOTKEY_JS_CODE)
-
             yield demo