diff --git a/streamlit_app.py b/streamlit_app.py index 433afaa..739b026 100644 --- a/streamlit_app.py +++ b/streamlit_app.py @@ -347,9 +347,22 @@ def play_old_messages(): # timestamp = datetime.utcnow() -with left_column: - if st.session_state['annotations']: - pdf_viewer(input=st.session_state['binary'], annotations=st.session_state['annotations']) +def rgb_to_hex(rgb): + return "#{:02x}{:02x}{:02x}".format(*rgb) + +def generate_color_gradient(num_elements): + # Define warm and cold colors in RGB format + warm_color = (255, 165, 0) # Orange + cold_color = (0, 0, 255) # Blue + + # Generate a linear gradient of colors + color_gradient = [ + rgb_to_hex(tuple(int(warm * (1 - i/num_elements) + cold * (i/num_elements)) for warm, cold in zip(warm_color, cold_color))) + for i in range(num_elements) + ] + + return color_gradient + with right_column: # css = ''' @@ -396,9 +409,13 @@ def play_old_messages(): _, text_response, coordinates = st.session_state['rqa'][model].query_document(question, st.session_state.doc_id, context_size=context_size) - st.session_state['annotations'] = [ - {"page": coo[0], "x": coo[1], "y": coo[2], "width": coo[3], "height": coo[4], "color": "blue"} for coo in [c.split(",") for coord in + annotations = [ + {"page": coo[0], "x": coo[1], "y": coo[2], "width": coo[3], "height": coo[4], "color": "grey"} for coo in [c.split(",") for coord in coordinates for c in coord]] + gradients = generate_color_gradient(len(annotations)) + for i, color in enumerate(gradients): + annotations[i]['color'] = color + st.session_state['annotations'] = annotations # with left_column: # pdf_viewer(input=st.session_state['binary'], annotations=st.session_state['annotations'], key=1) @@ -427,3 +444,7 @@ def play_old_messages(): elif st.session_state.loaded_embeddings and st.session_state.doc_id: play_old_messages() + +with left_column: + if st.session_state['binary']: + pdf_viewer(input=st.session_state['binary'], width=600, height=800, annotations=st.session_state['annotations'])