# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
# default_exp capture_annotator
# -

# %load_ext autoreload
# %autoreload 2

from nbdev import *

# +
#exporti

from functools import partial
import math

from ipywidgets import (AppLayout, VBox, HBox, Button, GridBox, Layout, Checkbox, HTML, IntText, Valid, Output)
from ipyevents import Event

from pathlib import Path

from ipyannotator.navi_widget import Navi
from ipyannotator.storage import setup_project_paths, get_image_list_from_folder, AnnotationStorage

from traitlets import Dict, Int, HasTraits, observe, dlink, link, List, Unicode, Bool

from ipyannotator.image_button import ImageButton


# +
#export

class CaptureGrid(GridBox, HasTraits):
    """
    Represents grid of `ImageButtons` with state. 
    
    """
    debug_output = Output(layout={'border': '1px solid black'})
    current_state = Dict()
    
    def __init__(self, grid_item=ImageButton, image_width=150, image_height=150, 
                 n_rows=3, n_cols=3, display_label=False):
        
        
        self.image_width = image_width
        self.image_height = image_height
        self.n_rows = n_rows
        self.n_cols = n_cols
        self._screen_im_number = IntText(value=n_rows * n_cols, 
                                         description='screen_image_number', 
                                         disabled=False)
            
        self._labels = [grid_item(display_label=display_label,
                                  image_width='%dpx' % self.image_width,
                                  image_height='%dpx' % self.image_height) for _ in range(self._screen_im_number.value)]
        
        self.callback = None
        
        self.observe(self.on_state_change, 'current_state')
            
        
        gap = 40 if display_label else 15
        
        centered_settings = {
            'grid_template_columns': " ".join(["%dpx" % (self.image_width + gap) for i
                                               in range(self.n_cols)]),
            'grid_template_rows': " ".join(["%dpx" % (self.image_height + gap) for i
                                            in range(self.n_rows)]),
            'justify_content': 'center',
            'align_content': 'space-around'
        }
        
         
        super().__init__(children=self._labels, layout=Layout(**centered_settings))
        
        
    @debug_output.capture(clear_output=True)
    def on_state_change(self, change=None):
        print('on_state_change', change['new'])
        new_state = change['new']
        updated = 0
        iter_state = iter(new_state)
        
        for label in self._labels:
            p = next(iter_state, None)
            if p:
                label.image_path = str(p)
                label.label_value = Path(p).stem
                label.active = new_state[p].get('answer', False)
            else:
                label.clear()
            
        if self.callback:
            self.register_on_click()
            
            
    def on_click(self, cb):
        self.callback = cb
        self.register_on_click()

    @debug_output.capture(clear_output=True)
    def register_on_click(self):
        print('register_on_click')
        for l in self._labels:
            l.reset_callbacks()
            l.on_click(partial(self.callback, name=l.name))        
# -

ibg = CaptureGrid(grid_item=ImageButton, image_width=50, image_height=75)
ibg

# You should not see anythong at this step, until you set a correct visual state (see next step below)

ibg.debug_output

# Actually `CaptureGrid` does not have own `on_click` event listener, but grid elements itself should implement `on_click(ev)` and `reset_callbacks()` methods to register/reset onclick callback function respectively. Also grid element shoudl have a field `name` in order user can destinguish between grid children.
#
# In current implementation `ImageButton` is default grid element.
#
# While ipyevents implementation lacks `sender` or `source` in callback args, `functools.partial` used to back element `name` into return value. You can see example of on_click event handler `test_handler` below. 
# `name` of the button is printed out on click.

ibg.current_state = {'../data/projects/capture1/pics/pink25x25.png': {'answer': True},
                     '../data/mock/pics/test200x200.png': {'answer': True}, '': {'answer': False}}

# +
h = HTML('Event info')                         
display(h)

        
def test_handler(event, name=None):
    event.update({'label_name': name})
    lines = ['{}: {}'.format(k, v) for k, v in event.items()]
    content = '<br>'.join(lines)
    h.value = event['label_name']

ibg.on_click(test_handler)
# -

ibg.debug_output

ibg._screen_im_number.value


# +
#exporti

class CaptureAnnotatorGUI(AppLayout):
    def __init__(self, image_width=150, image_height=150, 
                 n_rows=3, n_cols=3):
        
        self._screen_im_number = IntText(value=n_rows * n_cols, 
                                    description='screen_image_number', 
                                    disabled=False)
        
        self.image_width = image_width
        self.image_height = image_height
        self.n_rows = n_rows
        self.n_cols = n_cols
        
        self._navi = Navi()
        
        self._save_btn = Button(description="Save",
                                layout=Layout(width='auto'))
        
        self._none_checkbox = Checkbox(description="Select none",
                                       indent=False,
                                       layout=Layout(width='100px'))

        self._controls_box = HBox([self._navi, self._save_btn, self._none_checkbox],
                                 layout=Layout(display='flex', justify_content='center', flex_flow='wrap', align_items='center'))
        
         
        self._grid_box = CaptureGrid(image_width=image_width, image_height=image_height,  n_rows=n_rows, n_cols=n_cols, display_label=False)


        self._grid_label = HTML()
        self._labels_box = VBox(children = [self._grid_label, self._grid_box],
                                layout=Layout(display='flex', justify_content='center', flex_wrap='wrap', align_items='center'))

        
        super().__init__(header=None,
                 left_sidebar=None,
                 center=self._labels_box,
                 right_sidebar=None,
                 footer=self._controls_box,
                 pane_widths=(2, 8, 0),
                 pane_heights=(1, 4, 1))
        
# -

c = CaptureAnnotatorGUI(image_width=50, image_height=50, n_rows=5, n_cols=5)
c._grid_label.value = 'Select smth'
c

c._grid_box.current_state = {'../data/projects/capture1/pics/pink25x25.png': {'answer': True}}


# +
#exporti

class CaptureAnnotatorLogic(HasTraits):
    debug_output = Output(layout={'border': '1px solid black'})
    index = Int(0) # state index
    disp_number = Int() # number of images on screen
    num_screens = Int() # number of screens
    current_state = Dict()
    question_value = Unicode()
    all_none = Bool()

    
    def __init__(self, project_path, question=None, image_dir='pics', filter_files=None, results_dir=None):
        self.project_path = Path(project_path)
        self.image_dir, self.annotation_file_path = setup_project_paths(self.project_path, image_dir=image_dir, results_dir=results_dir)
        
        self.image_paths = sorted(get_image_list_from_folder(self.image_dir)) #todo: use sorted in storage?

        if filter_files:
            self.image_paths = [p for p in self.image_paths if str(p) in filter_files]
        if not self.image_paths:
            raise UserWarning("No image files to display! Check image_dir of filter.")
        self.current_im_num = len(self.image_paths)
        self.annotations = AnnotationStorage(self.image_paths)
        if question:
            self.question_value = f'<center><p style="font-size:20px;">{question}</p></center>'
        self._update_state()
        
        
    @observe('disp_number')
    def _update_state(self, change=None): # from annotations
        state_images = self._get_state_names(self.index)
        current_state = {}
        for im_path in state_images:
            current_state[str(im_path)] = self.annotations.get(str(im_path)) or {}
        self.all_none = all(value == {'answer': False} for value in current_state.values())
        self.current_state = current_state

        
    def _update_annotations(self, index): # from screen  
        for p, anno in self.current_state.items():
            self.annotations[str(p)] = anno
        
    def _save_annotations(self, *args, **kwargs): # to disk
        index = kwargs.pop('old_index', self.index)
        self._update_annotations(index)
        self.annotations.save(self.annotation_file_path)
        
    def _get_state_names(self, index):
        start = index * self.disp_number
        end = start + self.disp_number
        im_names = self.image_paths[start:end]
        return im_names
        
    
    @observe('index')
    def _idx_changed(self, change):
        ''' On index change save old state 
            and update current state for visualisation
        '''
        self._save_annotations(old_index = change['old'])
        # update new screen
        self._update_state()
        
    @observe('disp_number')
    def _calc_screens_num(self, change=None):
        self.num_screens = math.ceil(self.current_im_num / self.disp_number)

        
    @debug_output.capture(clear_output=False)
    def _handle_grid_click(self, event, name=None):
        p = Path(self.image_dir, name)
        current_state = self.current_state.copy()
        if not p.is_dir():
            current_state[str(p)] = {'answer': not self.current_state[str(p)].get('answer', False)}
            if self.all_none:
                self.all_none = False
        else:
            return
        self.current_state = current_state
    
    def _select_none(self, change=None):
        self.current_state = {p: {'answer': False} for p in self.current_state}
# -

cal = CaptureAnnotatorLogic(project_path='../data/projects/capture1')
cal.disp_number = 9 # should be synced from gui 
cal.image_dir, cal.annotation_file_path, cal.current_im_num

# We have 16 images in `capture1` project on disk, so first screen should load 9 images;
# 7 images (16-9) left for second screen.

assert len(cal._get_state_names(0)) == 9
assert len(cal._get_state_names(1)) == 7
assert len(cal._get_state_names(5)) == 0

# List of image names for the 1st screen:

cal.current_state

# Suppose state change from gui:

cal.current_state['../data/projects/capture1/pics/pink25x25.png'] = {'answer': False}

# ##### (Next-> button emulation) 
# Increment index to initiate annotation save and switch state for a new screen

cal.index=1
cal.current_state

# ##### (<-Prev button emulation) 
# Decrement index to initiate annotation save and switch state for previous screen, loading existing annotation

cal.index = 0
cal.current_state

cal._select_none()
cal.current_state

# +
#export


class CaptureAnnotator(CaptureAnnotatorGUI):
    """
    Represents capture annotator. 
    
    Gives an ability to itarate through image dataset, 
    select images of same class,
    export final annotations in json format
    
    """
    
    def __init__(self, project_path, image_dir='pics', image_width=150, image_height=150, 
                 n_rows=3, n_cols=3, question=None, filter_files=None, results_dir=None):
        
        super().__init__(image_width, image_height, n_rows, n_cols)
        
        self._model = CaptureAnnotatorLogic(project_path, question, image_dir,
                                           filter_files=filter_files, results_dir=results_dir)
        
        self._save_btn.on_click(self._model._save_annotations)        

        self._grid_box.on_click(self._model._handle_grid_click)
        
        link((self._model, 'all_none'), (self._none_checkbox, 'value'))
        
        self._none_checkbox.observe(self._model._select_none, 'value')
            
        # get correct screen image number from gui settings
        dlink((self._screen_im_number, 'value'), (self._model, 'disp_number'))
        
        # set correct number of screens slider value based on image number
        dlink((self._model, 'num_screens'), (self._navi.model, 'max_im_number'))
        
        # link current image index from controls to annotator model 
        link((self._navi.model, 'index'), (self._model, 'index'))
        
        # link annotation question 
        link((self._model, 'question_value'), (self._grid_label, 'value'))
        
        # link state of model and grid box visualizer
        link((self._model, 'current_state'), (self._grid_box, 'current_state'))
    
        
            
    def to_dict(self, only_annotated=True):
        return self._model.annotations.to_dict(only_annotated)
             

# +
ca = CaptureAnnotator('../data/projects/capture1', image_width=75, image_height=75,
                      n_rows=4, n_cols=3, question="Select pink squares")

ca
# -

ca._model.debug_output

assert ca._model.current_state == ca._grid_box.current_state

#hide
from nbdev.export import notebook2script
notebook2script()



StarcoderdataJupyterScriptsDedupFiltered