Update UI

src/streamlit_app.py

@@ -162,142 +162,142 @@ with tab3:
 
 
 
-        # --- Model Training Section ---
-        st.subheader("1. Train OCR Model")
-        st.write("Click the button below to start training the OCR model.")
-
-        # Progress bar and label for training within this tab
-        progress_container = st.empty() # Container for dynamic messages and progress
-        progress_message_placeholder = st.empty()
-        progress_bar_placeholder = st.progress(0)
-        
-        def update_progress_callback(value, text):
-            progress_bar_placeholder.progress(int(value * 100))
-            progress_message_placeholder.info(text) # Use info for dynamic messages
-
-        if st.button("📊 Start Training"):
-            progress_message_placeholder.empty() # Clear previous messages
-            progress_bar_placeholder.progress(0) # Reset progress bar
-
-            if not os.path.exists(TRAIN_CSV_PATH) or not os.path.isdir(TRAIN_IMAGES_DIR):
-                st.error(f"Training CSV '{TRAIN_CSV_PATH}' or Images directory '{TRAIN_IMAGES_DIR}' not found! Please check file paths and ensure data is uploaded correctly.")
-            elif not os.path.exists(TEST_CSV_PATH) or not os.path.isdir(TEST_IMAGES_DIR):
-                st.warning(f"Test CSV '{TEST_CSV_PATH}' or Images directory '{TEST_IMAGES_DIR}' not found. "
-                        "Evaluation might be affected or skipped. Please ensure all data paths are correct and data is uploaded.")
-            else:
-                progress_message_placeholder.info(f"Training a new CRNN model for {NUM_EPOCHS} epochs. This will take significant time...")
+    # --- Model Training Section ---
+    st.subheader("Train OCR Model")
+    st.write("Click the button below to start training the OCR model.")
+
+    # Progress bar and label for training within this tab
+    progress_container = st.empty() # Container for dynamic messages and progress
+    progress_message_placeholder = st.empty()
+    progress_bar_placeholder = st.progress(0)
+    
+    def update_progress_callback(value, text):
+        progress_bar_placeholder.progress(int(value * 100))
+        progress_message_placeholder.info(text) # Use info for dynamic messages
+
+    if st.button("📊 Start Training"):
+        progress_message_placeholder.empty() # Clear previous messages
+        progress_bar_placeholder.progress(0) # Reset progress bar
+
+        if not os.path.exists(TRAIN_CSV_PATH) or not os.path.isdir(TRAIN_IMAGES_DIR):
+            st.error(f"Training CSV '{TRAIN_CSV_PATH}' or Images directory '{TRAIN_IMAGES_DIR}' not found! Please check file paths and ensure data is uploaded correctly.")
+        elif not os.path.exists(TEST_CSV_PATH) or not os.path.isdir(TEST_IMAGES_DIR):
+            st.warning(f"Test CSV '{TEST_CSV_PATH}' or Images directory '{TEST_IMAGES_DIR}' not found. "
+                    "Evaluation might be affected or skipped. Please ensure all data paths are correct and data is uploaded.")
+        else:
+            progress_message_placeholder.info(f"Training a new CRNN model for {NUM_EPOCHS} epochs. This will take significant time...")
+            
+            try:
+                train_df, test_df = load_ocr_dataframes(TRAIN_CSV_PATH, TEST_CSV_PATH)
+                progress_message_placeholder.success("Training and Test DataFrames loaded successfully.")
+                progress_message_placeholder.info(f"Train DataFrame size: {len(train_df)} samples")
+                progress_message_placeholder.info(f"Test DataFrame size: {len(test_df)} samples")
+                if len(test_df) == 0:
+                    progress_message_placeholder.error("ERROR: Test DataFrame is empty! Evaluation cannot proceed. Check TEST_CSV_PATH and TEST_IMAGES_DIR.")
+                if len(train_df) == 0:
+                    progress_message_placeholder.error("ERROR: Train DataFrame is empty! Training cannot proceed. Check TRAIN_CSV_PATH and TRAIN_IMAGES_DIR.")
                 
-                try:
-                    train_df, test_df = load_ocr_dataframes(TRAIN_CSV_PATH, TEST_CSV_PATH)
-                    progress_message_placeholder.success("Training and Test DataFrames loaded successfully.")
-                    progress_message_placeholder.info(f"Train DataFrame size: {len(train_df)} samples")
-                    progress_message_placeholder.info(f"Test DataFrame size: {len(test_df)} samples")
-                    if len(test_df) == 0:
-                        progress_message_placeholder.error("ERROR: Test DataFrame is empty! Evaluation cannot proceed. Check TEST_CSV_PATH and TEST_IMAGES_DIR.")
-                    if len(train_df) == 0:
-                        progress_message_placeholder.error("ERROR: Train DataFrame is empty! Training cannot proceed. Check TRAIN_CSV_PATH and TRAIN_IMAGES_DIR.")
-                    
-                    if len(train_df) == 0 or len(test_df) == 0: # Stop if critical data is missing
-                        st.stop() # Added st.stop for critical data missing scenario
+                if len(train_df) == 0 or len(test_df) == 0: # Stop if critical data is missing
+                    st.stop() # Added st.stop for critical data missing scenario
 
-                    char_indexer_for_training = CharIndexer(vocabulary_string=VOCABULARY, blank_token_symbol=BLANK_TOKEN_SYMBOL)
-                    progress_message_placeholder.success(f"CharIndexer initialized with {char_indexer_for_training.num_classes} classes.")
+                char_indexer_for_training = CharIndexer(vocabulary_string=VOCABULARY, blank_token_symbol=BLANK_TOKEN_SYMBOL)
+                progress_message_placeholder.success(f"CharIndexer initialized with {char_indexer_for_training.num_classes} classes.")
 
-                    train_loader, test_loader = create_ocr_dataloaders(train_df, test_df, char_indexer_for_training, BATCH_SIZE)
-                    progress_message_placeholder.success("DataLoaders created successfully.")
-                    
-                    # Re-initialize the model to train from scratch if the button is pressed
-                    # This ensures we don't continue training a potentially already trained model if it was loaded.
-                    ocr_model_for_training = CRNN(num_classes=char_indexer_for_training.num_classes, cnn_output_channels=512, rnn_hidden_size=256, rnn_num_layers=2)
-                    ocr_model_for_training.to(device)
-                    ocr_model_for_training.train() 
+                train_loader, test_loader = create_ocr_dataloaders(train_df, test_df, char_indexer_for_training, BATCH_SIZE)
+                progress_message_placeholder.success("DataLoaders created successfully.")
+                
+                # Re-initialize the model to train from scratch if the button is pressed
+                # This ensures we don't continue training a potentially already trained model if it was loaded.
+                ocr_model_for_training = CRNN(num_classes=char_indexer_for_training.num_classes, cnn_output_channels=512, rnn_hidden_size=256, rnn_num_layers=2)
+                ocr_model_for_training.to(device)
+                ocr_model_for_training.train() 
 
-                    progress_message_placeholder.write("Training in progress... This may take a while.")
-                    
-                    # Capture the model and history
-                    ocr_model_for_training, history_result = train_ocr_model(
-                        model=ocr_model_for_training,
-                        train_loader=train_loader,
-                        test_loader=test_loader,
-                        char_indexer=char_indexer_for_training,
-                        epochs=NUM_EPOCHS,
-                        device=device,
-                        progress_callback=update_progress_callback
-                    )
-                    
-                    st.session_state.training_history = history_result # Save history to session state
-                    
-                    progress_message_placeholder.success("OCR model training finished!")
-                    update_progress_callback(1.0, "Training complete!")
+                progress_message_placeholder.write("Training in progress... This may take a while.")
+                
+                # Capture the model and history
+                ocr_model_for_training, history_result = train_ocr_model(
+                    model=ocr_model_for_training,
+                    train_loader=train_loader,
+                    test_loader=test_loader,
+                    char_indexer=char_indexer_for_training,
+                    epochs=NUM_EPOCHS,
+                    device=device,
+                    progress_callback=update_progress_callback
+                )
+                
+                st.session_state.training_history = history_result # Save history to session state
+                
+                progress_message_placeholder.success("OCR model training finished!")
+                update_progress_callback(1.0, "Training complete!")
 
-                    os.makedirs(os.path.dirname(MODEL_SAVE_PATH), exist_ok=True)
-                    save_ocr_model(ocr_model_for_training, MODEL_SAVE_PATH)
-                    progress_message_placeholder.success(f"Trained model saved to `{MODEL_SAVE_PATH}`")
+                os.makedirs(os.path.dirname(MODEL_SAVE_PATH), exist_ok=True)
+                save_ocr_model(ocr_model_for_training, MODEL_SAVE_PATH)
+                progress_message_placeholder.success(f"Trained model saved to `{MODEL_SAVE_PATH}`")
 
-                    ocr_model = ocr_model_for_training 
-                    ocr_model.eval() # Set to eval mode for subsequent predictions
+                ocr_model = ocr_model_for_training 
+                ocr_model.eval() # Set to eval mode for subsequent predictions
 
-                except Exception as e:
-                    progress_message_placeholder.error(f"An error occurred during training: {e}")
-                    st.exception(e) # This will print a detailed traceback in the Streamlit UI
-                    update_progress_callback(0.0, "Training failed!")
-            
-        st.write("---")
+            except Exception as e:
+                progress_message_placeholder.error(f"An error occurred during training: {e}")
+                st.exception(e) # This will print a detailed traceback in the Streamlit UI
+                update_progress_callback(0.0, "Training failed!")
+        
+    st.write("---")
 
-        # --- Model Loading Section ---
-        st.subheader("2. Load Pre-trained Model")
-        st.write("If you have a saved model, you can load it here instead of training.")
+    # --- Model Loading Section ---
+    st.subheader("Load Pre-trained Model")
+    st.write("If you have a saved model, you can load it here instead of training.")
 
-        if st.button("💾 Load Model"):
-            if os.path.exists(MODEL_SAVE_PATH):
-                try:
-                    loaded_model_instance = CRNN(num_classes=char_indexer.num_classes, cnn_output_channels=512, rnn_hidden_size=256, rnn_num_layers=2)
-                    load_ocr_model(loaded_model_instance, MODEL_SAVE_PATH)
-                    loaded_model_instance.to(device)
-                    ocr_model = loaded_model_instance
-                    ocr_model.eval() 
-                    st.success(f"Model loaded successfully from `{MODEL_SAVE_PATH}`")
-                    
-                    # If a model is loaded, we can try to re-evaluate it to get history,
-                    # but typically history is stored from a training run.
-                    # For simplicity, we'll assume training history is only stored after a training run.
-                    
-                except Exception as e:
-                    st.error(f"Error loading model: {e}")
-                    st.exception(e)
-            else:
-                st.warning(f"No model found at `{MODEL_SAVE_PATH}`. Please train a model first or check the path.")
-
-        st.write("---")
-
-        # --- Training History Plots Section ---
-        st.subheader("3. Training History Plots")
-        if st.session_state.training_history: # Check if history exists in session state
-            history_df = pd.DataFrame({
-                'Epoch': range(1, len(st.session_state.training_history['train_loss']) + 1),
-                'Train Loss': st.session_state.training_history['train_loss'],
-                'Test Loss': st.session_state.training_history['test_loss'],
-                'Test CER (%)': [cer * 100 for cer in st.session_state.training_history['test_cer']],
-                'Test Exact Match Accuracy (%)': [acc * 100 for acc in st.session_state.training_history['test_exact_match_accuracy']]
-            })
-
-            st.markdown("**Loss over Epochs**")
-            st.line_chart(history_df.set_index('Epoch')[['Train Loss', 'Test Loss']])
-            st.caption("Lower loss indicates better model performance.")
-
-            st.markdown("**Character Error Rate (CER) over Epochs**")
-            st.line_chart(history_df.set_index('Epoch')[['Test CER (%)']])
-            st.caption("Lower CER indicates fewer character errors (0% is perfect).")
-
-            st.markdown("**Exact Match Accuracy over Epochs**")
-            st.line_chart(history_df.set_index('Epoch')[['Test Exact Match Accuracy (%)']])
-            st.caption("Higher exact match accuracy indicates more perfectly recognized names.")
-
-            st.markdown("**Performance Metrics over Epochs (CER vs. Exact Match Accuracy)**")
-            st.line_chart(history_df.set_index('Epoch')[['Test CER (%)', 'Test Exact Match Accuracy (%)']])
-            st.caption("CER should decrease, Accuracy should increase.")
+    if st.button("💾 Load Model"):
+        if os.path.exists(MODEL_SAVE_PATH):
+            try:
+                loaded_model_instance = CRNN(num_classes=char_indexer.num_classes, cnn_output_channels=512, rnn_hidden_size=256, rnn_num_layers=2)
+                load_ocr_model(loaded_model_instance, MODEL_SAVE_PATH)
+                loaded_model_instance.to(device)
+                ocr_model = loaded_model_instance
+                ocr_model.eval() 
+                st.success(f"Model loaded successfully from `{MODEL_SAVE_PATH}`")
+                
+                # If a model is loaded, we can try to re-evaluate it to get history,
+                # but typically history is stored from a training run.
+                # For simplicity, we'll assume training history is only stored after a training run.
+                
+            except Exception as e:
+                st.error(f"Error loading model: {e}")
+                st.exception(e)
         else:
-            st.info("Train the model first to see training history plots here.")
+            st.warning(f"No model found at `{MODEL_SAVE_PATH}`. Please train a model first or check the path.")
+
+    st.write("---")
+
+    # --- Training History Plots Section ---
+    st.subheader("Training History Plots")
+    if st.session_state.training_history: # Check if history exists in session state
+        history_df = pd.DataFrame({
+            'Epoch': range(1, len(st.session_state.training_history['train_loss']) + 1),
+            'Train Loss': st.session_state.training_history['train_loss'],
+            'Test Loss': st.session_state.training_history['test_loss'],
+            'Test CER (%)': [cer * 100 for cer in st.session_state.training_history['test_cer']],
+            'Test Exact Match Accuracy (%)': [acc * 100 for acc in st.session_state.training_history['test_exact_match_accuracy']]
+        })
+
+        st.markdown("**Loss over Epochs**")
+        st.line_chart(history_df.set_index('Epoch')[['Train Loss', 'Test Loss']])
+        st.caption("Lower loss indicates better model performance.")
+
+        st.markdown("**Character Error Rate (CER) over Epochs**")
+        st.line_chart(history_df.set_index('Epoch')[['Test CER (%)']])
+        st.caption("Lower CER indicates fewer character errors (0% is perfect).")
+
+        st.markdown("**Exact Match Accuracy over Epochs**")
+        st.line_chart(history_df.set_index('Epoch')[['Test Exact Match Accuracy (%)']])
+        st.caption("Higher exact match accuracy indicates more perfectly recognized names.")
+
+        st.markdown("**Performance Metrics over Epochs (CER vs. Exact Match Accuracy)**")
+        st.line_chart(history_df.set_index('Epoch')[['Test CER (%)', 'Test Exact Match Accuracy (%)']])
+        st.caption("CER should decrease, Accuracy should increase.")
+    else:
+        st.info("Train the model first to see training history plots here.")
 
 
 # --- Final Footer ---