Update app.py

app.py

@@ -1,40 +1,45 @@
 import streamlit as st
 import sys
+import os
+import shutil
+import time
+from datetime import datetime
+import csv
+import cv2
+import numpy as np
+from PIL import Image
+import torch
+
+# Adjust import paths as needed
 sys.path.append('Utils')
 sys.path.append('model')
-import torch
 from model.CBAM.reunet_cbam import reunet_cbam
-import cv2
-from PIL import Image 
 from model.transform import transforms
-import numpy as np
 from model.unet import UNET
 from Utils.area import pixel_to_sqft, process_and_overlay_image
-import matplotlib.pyplot as plt
-import time
-import os
-import csv 
-from datetime import datetime
 from split_merge import split, merge 
-from Utils.convert import  read_pansharpened_rgb
-import shutil
-
+from Utils.convert import read_pansharpened_rgb
 
-# Define directories
-UPLOAD_DIR = "uploaded_images/"
-MASK_DIR = "generated_masks/"
-patches_folder = "Patches/"
-pred_patches = "Patch_pred/"
-CSV_LOG_PATH = "image_log.csv"
+# Define directories for Hugging Face Spaces
+UPLOAD_DIR = "/tmp/uploaded_images/"
+MASK_DIR = "/tmp/generated_masks/"
+patches_folder = "/tmp/Patches/"
+pred_patches = "/tmp/Patch_pred/"
+CSV_LOG_PATH = "outputs/image_log.csv"
 
 # Create directories
-for directory in [UPLOAD_DIR, MASK_DIR, patches_folder, pred_patches]:
+for directory in [UPLOAD_DIR, MASK_DIR, patches_folder, pred_patches, "outputs"]:
     os.makedirs(directory, exist_ok=True)
 
 # Load model
-model = reunet_cbam()
-model.load_state_dict(torch.load('latest.pth', map_location='cpu')['model_state_dict'])
-model.eval()
+@st.cache_resource
+def load_model():
+    model = reunet_cbam()
+    model.load_state_dict(torch.load('latest.pth', map_location='cpu')['model_state_dict'])
+    model.eval()
+    return model
+
+model = load_model()
 
 def predict(image):
     with torch.no_grad():
@@ -63,27 +68,6 @@ def log_image_details(image_id, image_filename, mask_filename):
         
         writer.writerow([sno, date, time, image_id, image_filename, mask_filename])
 
-def overlay_mask(image, mask, alpha=0.5, rgb=[255, 0, 0]):
-    # Ensure image is 3-channel
-    if len(image.shape) == 2:
-        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
-    
-    # Ensure mask is binary and same shape as image
-    mask = mask.astype(bool)
-    if mask.shape[:2] != image.shape[:2]:
-        raise ValueError("Mask and image must have the same dimensions")
-    
-    # Create color overlay
-    color_mask = np.zeros_like(image)
-    color_mask[mask] = rgb
-    
-    # Blend the image and color mask
-    output = cv2.addWeighted(image, 1, color_mask, alpha, 0)
-    
-    return output
-
-import shutil  # Add this import at the top of your file
-
 def upload_page():
     if 'file_uploaded' not in st.session_state:
         st.session_state.file_uploaded = False
@@ -155,8 +139,9 @@ def upload_page():
                         Image.fromarray(mask).save(mask_filepath)
 
                 # Merge predicted patches
-                merged_mask_filename = f"generated_masks/mask_{timestamp}.png"
-                merge(pred_patches, merged_mask_filename, img_array.shape)
+                merged_mask_filename = f"mask_{timestamp}.png"
+                merged_mask_path = os.path.join(MASK_DIR, merged_mask_filename)
+                merge(pred_patches, merged_mask_path, img_array.shape)
 
                 # Save merged mask
                 st.session_state.mask_filename = merged_mask_filename
@@ -176,7 +161,7 @@ def upload_page():
             mask_filename = f"mask_{timestamp}.png"
             mask_filepath = os.path.join(MASK_DIR, mask_filename)
             Image.fromarray(mask).save(mask_filepath)
-            st.session_state.mask_filename = mask_filepath
+            st.session_state.mask_filename = mask_filename
 
         st.session_state.file_uploaded = True
 
@@ -187,6 +172,7 @@ def upload_page():
             st.success('Image analyzed')
             st.session_state.page = 'result'
             st.rerun()
+
 def result_page():
     st.title('Analysis Result')
     
@@ -211,7 +197,7 @@ def result_page():
         col1.error(f"Original image file not found: {original_img_path}")
 
     # Display predicted mask
-    mask_path = st.session_state.mask_filename
+    mask_path = os.path.join(MASK_DIR, st.session_state.mask_filename)
     if os.path.exists(mask_path):
         mask = Image.open(mask_path)
         col2.image(mask, caption='Predicted Mask', use_column_width=True)
@@ -235,9 +221,6 @@ def result_page():
         # Process and overlay image
         overlay_img = process_and_overlay_image(original_np, mask_np, 'output.png')
         
-        # Convert BGR to RGB for displaying with st.image
-       # overlay_rgb = cv2.cvtColor(overlay_img, cv2.COLOR_BGR2RGB)
-        
         st.image(overlay_img, caption='Overlay Image', use_column_width=True)
     else:
         st.error("Image or mask file not found for overlay.")
@@ -263,4 +246,4 @@ def main():
         result_page()
 
 if __name__ == '__main__':
-    main()
+    main()