Spaces:
Runtime error
Runtime error
import streamlit as st | |
from PIL import Image | |
import streamlit_nested_layout | |
from streamlit_sparrow_labeling import st_sparrow_labeling | |
from streamlit_sparrow_labeling import DataProcessor | |
import json | |
import math | |
import os | |
from natsort import natsorted | |
from tools import agstyler | |
from tools.agstyler import PINLEFT | |
import pandas as pd | |
from toolbar_main import component_toolbar_main | |
class DataAnnotation: | |
class Model: | |
pageTitle = "Data Annotation" | |
img_file = None | |
rects_file = None | |
labels_file = "docs/labels.json" | |
groups_file = "docs/groups.json" | |
assign_labels_text = "Assign Labels" | |
text_caption_1 = "Check 'Assign Labels' to enable editing of labels and values, move and resize the boxes to annotate the document." | |
text_caption_2 = "Add annotations by clicking and dragging on the document, when 'Assign Labels' is unchecked." | |
labels = ["", "invoice_no", "invoice_date", "seller", "client", "seller_tax_id", "client_tax_id", "iban", "item_desc", | |
"item_qty", "item_net_price", "item_net_worth", "item_vat", "item_gross_worth", "total_net_worth", "total_vat", | |
"total_gross_worth"] | |
groups = ["", "items_row1", "items_row2", "items_row3", "items_row4", "items_row5", "items_row6", "items_row7", | |
"items_row8", "items_row9", "items_row10", "summary"] | |
selected_field = "Selected Field: " | |
save_text = "Save" | |
saved_text = "Saved!" | |
subheader_1 = "Select" | |
subheader_2 = "Upload" | |
annotation_text = "Annotation" | |
no_annotation_file = "No annotation file selected" | |
no_annotation_mapping = "Please annotate the document. Uncheck 'Assign Labels' and draw new annotations" | |
download_text = "Download" | |
download_hint = "Download the annotated structure in JSON format" | |
annotation_selection_help = "Select an annotation file to load" | |
upload_help = "Upload a file to annotate" | |
upload_button_text = "Upload" | |
upload_button_text_desc = "Choose a file" | |
assign_labels_text = "Assign Labels" | |
assign_labels_help = "Check to enable editing of labels and values" | |
export_labels_text = "Export Labels" | |
export_labels_help = "Create key-value pairs for the labels in JSON format" | |
done_text = "Done" | |
grouping_id = "ID" | |
grouping_value = "Value" | |
completed_text = "Completed" | |
completed_help = "Check to mark the annotation as completed" | |
error_text = "Value is too long. Please shorten it." | |
selection_must_be_continuous = "Please select continuous rows" | |
def view(self, model, ui_width, device_type, device_width): | |
with open(model.labels_file, "r") as f: | |
labels_json = json.load(f) | |
labels_list = labels_json["labels"] | |
labels = [''] | |
for label in labels_list: | |
labels.append(label['name']) | |
model.labels = labels | |
with open(model.groups_file, "r") as f: | |
groups_json = json.load(f) | |
groups_list = groups_json["groups"] | |
groups = [''] | |
for group in groups_list: | |
groups.append(group['name']) | |
model.groups = groups | |
with st.sidebar: | |
st.markdown("---") | |
st.subheader(model.subheader_1) | |
placeholder_upload = st.empty() | |
file_names = self.get_existing_file_names('docs/images/') | |
if 'annotation_index' not in st.session_state: | |
st.session_state['annotation_index'] = 0 | |
annotation_index = 0 | |
else: | |
annotation_index = st.session_state['annotation_index'] | |
annotation_selection = placeholder_upload.selectbox(model.annotation_text, file_names, | |
index=annotation_index, | |
help=model.annotation_selection_help) | |
annotation_index = self.get_annotation_index(annotation_selection, file_names) | |
file_extension = self.get_file_extension(annotation_selection, 'docs/images/') | |
model.img_file = f"docs/images/{annotation_selection}" + file_extension | |
model.rects_file = f"docs/json/{annotation_selection}.json" | |
completed_check = st.empty() | |
btn = st.button(model.export_labels_text, disabled=True) | |
if btn: | |
self.export_labels(model) | |
st.write(model.done_text) | |
st.subheader(model.subheader_2) | |
with st.form("upload-form", clear_on_submit=True): | |
uploaded_file = st.file_uploader(model.upload_button_text_desc, accept_multiple_files=False, | |
type=['png', 'jpg', 'jpeg'], | |
help=model.upload_help, disabled=True) | |
submitted = st.form_submit_button(model.upload_button_text, disabled=True) | |
if submitted and uploaded_file is not None: | |
ret = self.upload_file(uploaded_file) | |
if ret is not False: | |
file_names = self.get_existing_file_names('docs/images/') | |
annotation_index = self.get_annotation_index(annotation_selection, file_names) | |
annotation_selection = placeholder_upload.selectbox(model.annotation_text, file_names, | |
index=annotation_index, | |
help=model.annotation_selection_help) | |
st.session_state['annotation_index'] = annotation_index | |
# st.title(model.pageTitle + " - " + annotation_selection) | |
if model.img_file is None: | |
st.caption(model.no_annotation_file) | |
return | |
saved_state = self.fetch_annotations(model.rects_file) | |
# annotation file has been changed | |
if annotation_index != st.session_state['annotation_index']: | |
annotation_v = saved_state['meta']['version'] | |
if annotation_v == "v0.1": | |
st.session_state["annotation_done"] = False | |
else: | |
st.session_state["annotation_done"] = True | |
# store the annotation file index | |
st.session_state['annotation_index'] = annotation_index | |
# first load | |
if "annotation_done" not in st.session_state: | |
annotation_v = saved_state['meta']['version'] | |
if annotation_v == "v0.1": | |
st.session_state["annotation_done"] = False | |
else: | |
st.session_state["annotation_done"] = True | |
with completed_check: | |
annotation_done = st.checkbox(model.completed_text, help=model.completed_help, key="annotation_done", disabled=True) | |
if annotation_done: | |
saved_state['meta']['version'] = "v1.0" | |
else: | |
saved_state['meta']['version'] = "v0.1" | |
with open(model.rects_file, "w") as f: | |
json.dump(saved_state, f, indent=2) | |
st.session_state[model.rects_file] = saved_state | |
assign_labels = st.checkbox(model.assign_labels_text, True, help=model.assign_labels_help) | |
mode = "transform" if assign_labels else "rect" | |
docImg = Image.open(model.img_file) | |
data_processor = DataProcessor() | |
with st.container(): | |
doc_height = saved_state['meta']['image_size']['height'] | |
doc_width = saved_state['meta']['image_size']['width'] | |
canvas_width, number_of_columns = self.canvas_available_width(ui_width, doc_width, device_type, | |
device_width) | |
if number_of_columns > 1: | |
col1, col2 = st.columns([number_of_columns, 10 - number_of_columns]) | |
with col1: | |
result_rects = self.render_doc(model, docImg, saved_state, mode, canvas_width, doc_height, doc_width) | |
with col2: | |
tab = st.radio("Select", ["Mapping", "Grouping", "Ordering"], horizontal=True, | |
label_visibility="collapsed") | |
if tab == "Mapping": | |
self.render_form(model, result_rects, data_processor, annotation_selection) | |
elif tab == "Grouping": | |
self.group_annotations(model, result_rects) | |
elif tab == "Ordering": | |
self.order_annotations(model, model.labels, model.groups, result_rects) | |
else: | |
result_rects = self.render_doc(model, docImg, saved_state, mode, canvas_width, doc_height, doc_width) | |
tab = st.radio("Select", ["Mapping", "Grouping"], horizontal=True, label_visibility="collapsed") | |
if tab == "Mapping": | |
self.render_form(model, result_rects, data_processor, annotation_selection) | |
else: | |
self.group_annotations(model, result_rects) | |
def render_doc(self, model, docImg, saved_state, mode, canvas_width, doc_height, doc_width): | |
with st.container(): | |
height = 1296 | |
width = 864 | |
result_rects = st_sparrow_labeling( | |
fill_color="rgba(0, 151, 255, 0.3)", | |
stroke_width=2, | |
stroke_color="rgba(0, 50, 255, 0.7)", | |
background_image=docImg, | |
initial_rects=saved_state, | |
height=height, | |
width=width, | |
drawing_mode=mode, | |
display_toolbar=True, | |
update_streamlit=True, | |
canvas_width=canvas_width, | |
doc_height=doc_height, | |
doc_width=doc_width, | |
image_rescale=True, | |
key="doc_annotation" + model.img_file | |
) | |
st.caption(model.text_caption_1) | |
st.caption(model.text_caption_2) | |
return result_rects | |
def render_form(self, model, result_rects, data_processor, annotation_selection): | |
with st.container(): | |
if result_rects is not None: | |
with st.form(key="fields_form"): | |
toolbar = st.empty() | |
self.render_form_view(result_rects.rects_data['words'], model.labels, result_rects, | |
data_processor) | |
with toolbar: | |
submit = st.form_submit_button(model.save_text, type="primary", disabled=True) | |
if submit: | |
for word in result_rects.rects_data['words']: | |
if len(word['value']) > 1000: | |
st.error(model.error_text) | |
return | |
with open(model.rects_file, "w") as f: | |
json.dump(result_rects.rects_data, f, indent=2) | |
st.session_state[model.rects_file] = result_rects.rects_data | |
# st.write(model.saved_text) | |
st.experimental_rerun() | |
if len(result_rects.rects_data['words']) == 0: | |
st.caption(model.no_annotation_mapping) | |
return | |
else: | |
with open(model.rects_file, 'rb') as file: | |
st.download_button(label=model.download_text, | |
data=file, | |
file_name=annotation_selection + ".json", | |
mime='application/json', | |
help=model.download_hint) | |
def render_form_view(self, words, labels, result_rects, data_processor): | |
data = [] | |
for i, rect in enumerate(words): | |
group, label = rect['label'].split(":", 1) if ":" in rect['label'] else (None, rect['label']) | |
data.append({'id': i, 'value': rect['value'], 'label': label}) | |
df = pd.DataFrame(data) | |
formatter = { | |
'id': ('ID', {**PINLEFT, 'hide': True}), | |
'value': ('Value', {**PINLEFT, 'editable': True}), | |
'label': ('Label', {**PINLEFT, | |
'width': 80, | |
'editable': True, | |
'cellEditor': 'agSelectCellEditor', | |
'cellEditorParams': { | |
'values': labels | |
}}) | |
} | |
go = { | |
'rowClassRules': { | |
'row-selected': 'data.id === ' + str(result_rects.current_rect_index) | |
} | |
} | |
green_light = "#abf7b1" | |
css = { | |
'.row-selected': { | |
'background-color': f'{green_light} !important' | |
} | |
} | |
response = agstyler.draw_grid( | |
df, | |
formatter=formatter, | |
fit_columns=True, | |
grid_options=go, | |
css=css | |
) | |
data = response['data'].values.tolist() | |
for i, rect in enumerate(words): | |
value = data[i][1] | |
label = data[i][2] | |
data_processor.update_rect_data(result_rects.rects_data, i, value, label) | |
def canvas_available_width(self, ui_width, doc_width, device_type, device_width): | |
doc_width_pct = (doc_width * 100) / ui_width | |
if doc_width_pct < 45: | |
canvas_width_pct = 37 | |
elif doc_width_pct < 55: | |
canvas_width_pct = 49 | |
else: | |
canvas_width_pct = 60 | |
if ui_width > 700 and canvas_width_pct == 37 and device_type == "desktop": | |
return math.floor(canvas_width_pct * ui_width / 100), 4 | |
elif ui_width > 700 and canvas_width_pct == 49 and device_type == "desktop": | |
return math.floor(canvas_width_pct * ui_width / 100), 5 | |
elif ui_width > 700 and canvas_width_pct == 60 and device_type == "desktop": | |
return math.floor(canvas_width_pct * ui_width / 100), 6 | |
else: | |
if device_type == "desktop": | |
ui_width = device_width - math.floor((device_width * 22) / 100) | |
elif device_type == "mobile": | |
ui_width = device_width - math.floor((device_width * 13) / 100) | |
return ui_width, 1 | |
def fetch_annotations(self, rects_file): | |
for key in st.session_state: | |
if key.startswith("docs/json/") and key != rects_file: | |
del st.session_state[key] | |
if rects_file not in st.session_state: | |
with open(rects_file, "r") as f: | |
saved_state = json.load(f) | |
st.session_state[rects_file] = saved_state | |
else: | |
saved_state = st.session_state[rects_file] | |
return saved_state | |
def upload_file(self, uploaded_file): | |
if uploaded_file is not None: | |
if os.path.exists(os.path.join("docs/images/", uploaded_file.name)): | |
st.write("File already exists") | |
return False | |
if len(uploaded_file.name) > 100: | |
st.write("File name too long") | |
return False | |
with open(os.path.join("docs/images/", uploaded_file.name), "wb") as f: | |
f.write(uploaded_file.getbuffer()) | |
img_file = Image.open(os.path.join("docs/images/", uploaded_file.name)) | |
annotations_json = { | |
"meta": { | |
"version": "v0.1", | |
"split": "train", | |
"image_id": len(self.get_existing_file_names("docs/images/")), | |
"image_size": { | |
"width": img_file.width, | |
"height": img_file.height | |
} | |
}, | |
"words": [] | |
} | |
file_name = uploaded_file.name.split(".")[0] | |
with open(os.path.join("docs/json/", file_name + ".json"), "w") as f: | |
json.dump(annotations_json, f, indent=2) | |
st.success("File uploaded successfully") | |
def get_existing_file_names(self, dir_name): | |
# get ordered list of files without file extension, excluding hidden files | |
return natsorted([os.path.splitext(f)[0] for f in os.listdir(dir_name) if not f.startswith('.')]) | |
def get_file_extension(self, file_name, dir_name): | |
# get list of files, excluding hidden files | |
files = [f for f in os.listdir(dir_name) if not f.startswith('.')] | |
for f in files: | |
if file_name is not None and os.path.splitext(f)[0] == file_name: | |
return os.path.splitext(f)[1] | |
def get_annotation_index(self, file, files_list): | |
return files_list.index(file) | |
def group_annotations(self, model, result_rects): | |
with st.form(key="grouping_form"): | |
if result_rects is not None: | |
words = result_rects.rects_data['words'] | |
data = [] | |
for i, rect in enumerate(words): | |
data.append({'id': i, 'value': rect['value']}) | |
df = pd.DataFrame(data) | |
formatter = { | |
'id': ('ID', {**PINLEFT, 'width': 50}), | |
'value': ('Value', PINLEFT) | |
} | |
toolbar = st.empty() | |
response = agstyler.draw_grid( | |
df, | |
formatter=formatter, | |
fit_columns=True, | |
selection='multiple', | |
use_checkbox='True', | |
pagination_size=40 | |
) | |
rows = response['selected_rows'] | |
with toolbar: | |
submit = st.form_submit_button(model.save_text, type="primary", disabled=True) | |
if submit and len(rows) > 0: | |
# check if there are gaps in the selected rows | |
if len(rows) > 1: | |
for i in range(len(rows) - 1): | |
if rows[i]['id'] + 1 != rows[i + 1]['id']: | |
st.error(model.selection_must_be_continuous) | |
return | |
words = result_rects.rects_data['words'] | |
new_words_list = [] | |
coords = [] | |
for row in rows: | |
word_value = words[row['id']]['value'] | |
rect = words[row['id']]['rect'] | |
coords.append(rect) | |
new_words_list.append(word_value) | |
# convert array to string | |
new_word = " ".join(new_words_list) | |
# Get min x1 value from coords array | |
x1_min = min([coord['x1'] for coord in coords]) | |
y1_min = min([coord['y1'] for coord in coords]) | |
x2_max = max([coord['x2'] for coord in coords]) | |
y2_max = max([coord['y2'] for coord in coords]) | |
words[rows[0]['id']]['value'] = new_word | |
words[rows[0]['id']]['rect'] = { | |
"x1": x1_min, | |
"y1": y1_min, | |
"x2": x2_max, | |
"y2": y2_max | |
} | |
# loop array in reverse order and remove selected entries | |
i = 0 | |
for row in rows[::-1]: | |
if i == len(rows) - 1: | |
break | |
del words[row['id']] | |
i += 1 | |
result_rects.rects_data['words'] = words | |
with open(model.rects_file, "w") as f: | |
json.dump(result_rects.rects_data, f, indent=2) | |
st.session_state[model.rects_file] = result_rects.rects_data | |
st.experimental_rerun() | |
def order_annotations(self, model, labels, groups, result_rects): | |
if result_rects is not None: | |
self.action_event = None | |
data = [] | |
idx_list = [""] | |
words = result_rects.rects_data['words'] | |
for i, rect in enumerate(words): | |
if rect['label'] != "": | |
# split string into two variables, assign None to first variable if no split is found | |
group, label = rect['label'].split(":", 1) if ":" in rect['label'] else (None, rect['label']) | |
data.append({'id': i, 'value': rect['value'], 'label': label, 'group': group}) | |
idx_list.append(i) | |
df = pd.DataFrame(data) | |
formatter = { | |
'id': ('ID', {**PINLEFT, 'width': 50}), | |
'value': ('Value', {**PINLEFT}), | |
'label': ('Label', {**PINLEFT, | |
'width': 80, | |
'editable': False, | |
'cellEditor': 'agSelectCellEditor', | |
'cellEditorParams': { | |
'values': labels | |
}}), | |
'group': ('Group', {**PINLEFT, | |
'width': 80, | |
'editable': True, | |
'cellEditor': 'agSelectCellEditor', | |
'cellEditorParams': { | |
'values': groups | |
}}) | |
} | |
go = { | |
'rowClassRules': { | |
'row-selected': 'data.id === ' + str(result_rects.current_rect_index) | |
} | |
} | |
green_light = "#abf7b1" | |
css = { | |
'.row-selected': { | |
'background-color': f'{green_light} !important' | |
} | |
} | |
idx_option = st.selectbox('Select row to move into', idx_list) | |
def run_component(props): | |
value = component_toolbar_main(key='toolbar_main', **props) | |
return value | |
def handle_event(value): | |
if value is not None: | |
if 'action_timestamp' not in st.session_state: | |
self.action_event = value['action'] | |
st.session_state['action_timestamp'] = value['timestamp'] | |
else: | |
if st.session_state['action_timestamp'] != value['timestamp']: | |
self.action_event = value['action'] | |
st.session_state['action_timestamp'] = value['timestamp'] | |
else: | |
self.action_event = None | |
props = { | |
'buttons': { | |
'up': { | |
'disabled': True, | |
'rendered': '' | |
}, | |
'down': { | |
'disabled': True, | |
'rendered': '' | |
}, | |
'save': { | |
'disabled': True, | |
'rendered': '' | |
# 'rendered': 'none', | |
} | |
} | |
} | |
handle_event(run_component(props)) | |
response = agstyler.draw_grid( | |
df, | |
formatter=formatter, | |
fit_columns=True, | |
grid_options=go, | |
css=css | |
) | |
rows = response['selected_rows'] | |
if len(rows) == 0 and result_rects.current_rect_index > -1: | |
for i, row in enumerate(data): | |
if row['id'] == result_rects.current_rect_index: | |
rows = [ | |
{ | |
'_selectedRowNodeInfo': { | |
'nodeRowIndex': i | |
}, | |
'id': row['id'] | |
} | |
] | |
break | |
if str(self.action_event) == 'up': | |
if len(rows) > 0: | |
idx = rows[0]['_selectedRowNodeInfo']['nodeRowIndex'] | |
if idx > 0: | |
row_id = rows[0]['id'] | |
if row_id == idx_option: | |
return | |
# swap row upwards in the array | |
if idx_option == "": | |
words[row_id], words[row_id - 1] = words[row_id - 1], words[row_id] | |
else: | |
for i in range(1000): | |
words[row_id], words[row_id - 1] = words[row_id - 1], words[row_id] | |
row_id -= 1 | |
if row_id == idx_option: | |
break | |
result_rects.rects_data['words'] = words | |
with open(model.rects_file, "w") as f: | |
json.dump(result_rects.rects_data, f, indent=2) | |
st.session_state[model.rects_file] = result_rects.rects_data | |
st.experimental_rerun() | |
elif str(self.action_event) == 'down': | |
if len(rows) > 0: | |
idx = rows[0]['_selectedRowNodeInfo']['nodeRowIndex'] | |
if idx < len(df) - 1: | |
row_id = rows[0]['id'] | |
if row_id == idx_option: | |
return | |
# swap row downwards in the array | |
if idx_option == "": | |
words[row_id], words[row_id + 1] = words[row_id + 1], words[row_id] | |
else: | |
for i in range(1000): | |
words[row_id], words[row_id + 1] = words[row_id + 1], words[row_id] | |
row_id += 1 | |
if row_id == idx_option: | |
break | |
result_rects.rects_data['words'] = words | |
with open(model.rects_file, "w") as f: | |
json.dump(result_rects.rects_data, f, indent=2) | |
st.session_state[model.rects_file] = result_rects.rects_data | |
st.experimental_rerun() | |
elif str(self.action_event) == 'save': | |
data = response['data'].values.tolist() | |
for elem in data: | |
if elem[3] != "None": | |
idx = elem[0] | |
group = elem[3] | |
words[idx]['label'] = f"{group}:{elem[2]}" | |
result_rects.rects_data['words'] = words | |
with open(model.rects_file, "w") as f: | |
json.dump(result_rects.rects_data, f, indent=2) | |
st.session_state[model.rects_file] = result_rects.rects_data | |
st.experimental_rerun() | |
def export_labels(self, model): | |
path_from = os.path.join("docs/json/") | |
path_to = os.path.join("docs/json/key/") | |
files = [f for f in os.listdir(path_from) if not f.startswith('.')] | |
for file in files: | |
path = os.path.join(path_from, file) | |
if os.path.isfile(path): | |
with open(path, "r") as f: | |
data = json.load(f) | |
words = data['words'] | |
keys = {} | |
row_keys = {} | |
for word in words: | |
if word['label'] != '': | |
if ':' in word['label']: | |
group, label = word['label'].split(':', 1) | |
if 'row' not in group: | |
if group not in keys: | |
keys[group] = {} | |
keys[group][label] = word['value'] | |
else: | |
if "items" not in keys: | |
keys["items"] = [] | |
if group not in row_keys: | |
row_keys[group] = {} | |
row_keys[group][label] = word['value'] | |
else: | |
keys[word['label']] = word['value'] | |
if row_keys != {}: | |
for key in row_keys: | |
keys["items"].append(row_keys[key]) | |
if keys != {}: | |
path = os.path.join(path_to, file) | |
with open(path, "w") as f: | |
json.dump(keys, f, indent=2) |