Update app.py

app.py

@@ -5,15 +5,15 @@ import torch
 import numpy as np
 import urllib.request
 
-# Load model without caching due to serialization issue with PretrainedConfig
+# Load the model without caching to avoid serialization issues
 def load_model():
     model = AutoModel.from_pretrained("ragavsachdeva/magi", trust_remote_code=True)
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model.to(device)
     return model
 
-# Load the model outside caching
-model = load_model()  # Place this line here, right after defining load_model
+# Initialize the model once at the top level, outside any caching functions
+model = load_model()
 
 @st.cache_data
 def read_image_as_np_array(image_path):
@@ -45,7 +45,6 @@ def predict_detections_and_associations(
             )[0]
     return result
 
-
 @st.cache_data
 def predict_ocr(
     image_path,
@@ -59,7 +58,7 @@ def predict_ocr(
         return
     image = read_image_as_np_array(image_path)
     result = predict_detections_and_associations(
-        image_path,
+        path_to_image,
         character_detection_threshold,
         panel_detection_threshold,
         text_detection_threshold,
@@ -71,23 +70,25 @@ def predict_ocr(
         ocr_results = model.predict_ocr([image], text_bboxes_for_all_images)
     return ocr_results
 
-def clear_memory():
-    st.session_state.memory = {"characters": {}, "transcript": ""}
-    st.write("Memory cleared.")
+# Streamlit UI elements
+st.markdown("""    
+<style> .title-container { background-color: #0d1117; padding: 20px; border-radius: 10px; margin: 20px; } 
+.title { font-size: 2em; text-align: center; color: #fff; font-family: 'Comic Sans MS', cursive; text-transform: uppercase; 
+letter-spacing: 0.1em; padding: 0.5em 0 0.2em; background: 0 0; } .title span { background: -webkit-linear-gradient(45deg, 
+#6495ed, #4169e1); -webkit-background-clip: text; -webkit-text-fill-color: transparent; } .subheading { font-size: 1.5em; 
+text-align: center; color: #ddd; font-family: 'Comic Sans MS', cursive; } .affil, .authors { font-size: 1em; text-align: center; 
+color: #ddd; font-family: 'Comic Sans MS', cursive; } .authors { padding-top: 1em; } </style> 
+<div class='title-container'> <div class='title'> The <span>Ma</span>n<span>g</span>a Wh<span>i</span>sperer </div> 
+<div class='subheading'> Automatically Generating Transcriptions for Comics </div> <div class='authors'> Ragav Sachdeva and 
+Andrew Zisserman </div> <div class='affil'> University of Oxford </div> </div>""", unsafe_allow_html=True)
 
-model = load_model()
-
-# Display header and UI components
-st.markdown("""    <style> ... styles here ... </style> """, unsafe_allow_html=True)
 path_to_image = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
 
-# Memory control button
-st.button("Clear Memory", on_click=clear_memory)
-
 st.sidebar.markdown("**Mode**")
-generate_detections_and_associations = st.sidebar.toggle("Generate detections and associations", True)
-generate_transcript = st.sidebar.toggle("Generate transcript (slower)", False)
+generate_detections_and_associations = st.sidebar.checkbox("Generate detections and associations", True)
+generate_transcript = st.sidebar.checkbox("Generate transcript (slower)", False)
 
+# Hyperparameter Sliders
 st.sidebar.markdown("**Hyperparameters**")
 input_character_detection_threshold = st.sidebar.slider('Character detection threshold', 0.0, 1.0, 0.30, step=0.01)
 input_panel_detection_threshold = st.sidebar.slider('Panel detection threshold', 0.0, 1.0, 0.2, step=0.01)
@@ -95,11 +96,13 @@ input_text_detection_threshold = st.sidebar.slider('Text detection threshold', 0
 input_character_character_matching_threshold = st.sidebar.slider('Character-character matching threshold', 0.0, 1.0, 0.7, step=0.01)
 input_text_character_matching_threshold = st.sidebar.slider('Text-character matching threshold', 0.0, 1.0, 0.4, step=0.01)
 
+# Main processing based on image input
 if path_to_image is not None:
     image = read_image_as_np_array(path_to_image)
     st.markdown("**Prediction**")
-    
-    if generate_detections_and_associations or generate_transcript:
+
+    # Run predictions based on checkbox selections
+    if generate_detections_and_associations:
         result = predict_detections_and_associations(
             path_to_image,
             input_character_detection_threshold,
@@ -108,6 +111,8 @@ if path_to_image is not None:
             input_character_character_matching_threshold,
             input_text_character_matching_threshold,
         )
+        output = model.visualise_single_image_prediction(image, result)
+        st.image(output)
 
     if generate_transcript:
         ocr_results = predict_ocr(
@@ -118,26 +123,5 @@ if path_to_image is not None:
             input_character_character_matching_threshold,
             input_text_character_matching_threshold,
         )
-        
-        # Append new characters and transcript to memory
-        if generate_detections_and_associations:
-            output = model.visualise_single_image_prediction(image, result)
-            st.image(output)
-            # Update character memory based on detected characters
-            detected_characters = result.get("characters", {})
-            st.session_state.memory["characters"].update(detected_characters)
-        
-        # Append the current transcript to the ongoing transcript in memory
         transcript = model.generate_transcript_for_single_image(result, ocr_results[0])
-        st.session_state.memory["transcript"] += transcript + "\n"
-
-        # Display the cumulative transcript from memory
-        st.text(st.session_state.memory["transcript"])
-
-    elif generate_detections_and_associations:
-        output = model.visualise_single_image_prediction(image, result)
-        st.image(output)
-
-    elif generate_transcript:
-        # Display the cumulative transcript
-        st.text(st.session_state.memory["transcript"])
+        st.text(transcript)