Update app.py

app.py

@@ -19,20 +19,18 @@ from Utils.split_merge import split, merge
 from Utils.convert_raster import convert_gtiff_to_8bit
 import shutil
 
-patches_folder = 'data/Patches'
-pred_patches = 'data/Patch_pred'
-os.makedirs(patches_folder, exist_ok=True)
-os.makedirs(pred_patches, exist_ok=True)
-
-# Define the upload directories
+# Define directories
 UPLOAD_DIR = "data/uploaded_images"
 MASK_DIR = "data/generated_masks"
+PATCHES_DIR = 'data/Patches'
+PRED_PATCHES_DIR = 'data/Patch_pred'
 CSV_LOG_PATH = "image_log.csv"
 
-# Create the directories if they don't exist
-os.makedirs(UPLOAD_DIR, exist_ok=True)
-os.makedirs(MASK_DIR, exist_ok=True)
+# Create directories
+for directory in [UPLOAD_DIR, MASK_DIR, PATCHES_DIR, PRED_PATCHES_DIR]:
+    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()
@@ -44,7 +42,6 @@ def predict(image):
 
 def log_image_details(image_id, image_filename, mask_filename):
     file_exists = os.path.exists(CSV_LOG_PATH)
-    
     current_time = datetime.now()
     date = current_time.strftime('%Y-%m-%d')
     time = current_time.strftime('%H:%M:%S')
@@ -54,35 +51,9 @@ def log_image_details(image_id, image_filename, mask_filename):
         if not file_exists:
             writer.writerow(['S.No', 'Date', 'Time', 'Image ID', 'Image Filename', 'Mask Filename'])
         
-        # Get the next S.No
-        if file_exists:
-            with open(CSV_LOG_PATH, mode='r') as f:
-                reader = csv.reader(f)
-                sno = sum(1 for row in reader)
-        else:
-            sno = 1
-        
+        sno = sum(1 for row in open(CSV_LOG_PATH)) if file_exists else 1
         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
-
 def reset_state():
     st.session_state.file_uploaded = False
     st.session_state.filename = None
@@ -91,95 +62,78 @@ def reset_state():
     if 'page' in st.session_state:
         del st.session_state.page
 
+def process_image(image, timestamp):
+    filename = f"image_{timestamp}{os.path.splitext(image.name)[1]}"
+    filepath = os.path.join(UPLOAD_DIR, filename)
+    
+    with open(filepath, "wb") as f:
+        f.write(image.getvalue())
+    
+    if filename.lower().endswith(('.tiff', '.tif')):
+        st.info('Processing GeoTIFF image...')
+        convert_gtiff_to_8bit(filepath)
+        st.success('GeoTIFF converted to 8-bit image')
+    
+    return filename, filepath
+
+def predict_image(img_array, filename, timestamp):
+    if img_array.shape[0] > 650 or img_array.shape[1] > 650:
+        split(os.path.join(UPLOAD_DIR, filename), patch_size=256)
+        
+        with st.spinner('Analyzing...'):
+            for patch_filename in os.listdir(PATCHES_DIR):
+                if patch_filename.endswith(".png"):
+                    patch_path = os.path.join(PATCHES_DIR, patch_filename)
+                    patch_img = Image.open(patch_path)
+                    patch_tr_img = transforms(patch_img)
+                    prediction = predict(patch_tr_img)
+                    mask = (prediction > 0.5).astype(np.uint8) * 255
+                    mask_filename = f"mask_{patch_filename}"
+                    mask_filepath = os.path.join(PRED_PATCHES_DIR, mask_filename)
+                    Image.fromarray(mask).save(mask_filepath)
+
+            merged_mask_filename = f"mask_{timestamp}.png"
+            merged_mask_filepath = os.path.join(MASK_DIR, merged_mask_filename)
+            merge(PRED_PATCHES_DIR, merged_mask_filepath, img_array.shape)
+            
+            st.info('Cleaning up temporary files...')
+            for dir in [PATCHES_DIR, PRED_PATCHES_DIR]:
+                shutil.rmtree(dir)
+                os.makedirs(dir)
+            st.success('Temporary files cleaned up')
+    else:
+        tr_img = transforms(Image.open(os.path.join(UPLOAD_DIR, filename)))
+        prediction = predict(tr_img)
+        mask = (prediction > 0.5).astype(np.uint8) * 255
+        merged_mask_filename = f"mask_{timestamp}.png"
+        merged_mask_filepath = os.path.join(MASK_DIR, merged_mask_filename)
+        Image.fromarray(mask).save(merged_mask_filepath)
+    
+    return merged_mask_filepath
+
 def upload_page():
     if 'file_uploaded' not in st.session_state:
         st.session_state.file_uploaded = False
     
-    if 'filename' not in st.session_state:
-        st.session_state.filename = None
-    
-    if 'mask_filename' not in st.session_state:
-        st.session_state.mask_filename = None
-
     image = st.file_uploader('Choose a satellite image', type=['jpg', 'png', 'jpeg', 'tiff', 'tif'])
 
     if image is not None:
-        reset_state()  # Reset the state when a new image is uploaded
-        bytes_data = image.getvalue()
-        
+        reset_state()
         timestamp = int(time.time())
-        original_filename = image.name
-        file_extension = os.path.splitext(original_filename)[1].lower()
-        
-        if file_extension in ['.tiff', '.tif']:
-            filename = f"image_{timestamp}.tif"
-        else:
-            filename = f"image_{timestamp}.png"
-        
-        filepath = os.path.join(UPLOAD_DIR, filename)
-        
-        with open(filepath, "wb") as f:
-            f.write(bytes_data)
-        
-        # Check if the uploaded file is a GeoTIFF
-        if file_extension in ['.tiff', '.tif']:
-            st.info('Processing GeoTIFF image...')
-            convert_gtiff_to_8bit(filepath)
-            st.success('GeoTIFF converted to 8-bit image')
+        filename, filepath = process_image(image, timestamp)
         
         img = Image.open(filepath)
         st.image(img, caption='Uploaded Image', use_column_width=True)
         st.success(f'Image saved as {filename}')
 
-        # Store the full path of the uploaded image
         st.session_state.filename = filename
-
-        # Convert image to numpy array
         img_array = np.array(img)
+        
+        mask_filepath = predict_image(img_array, filename, timestamp)
+        st.session_state.mask_filename = mask_filepath
 
-        # Check if image shape is more than 650x650
-        if img_array.shape[0] > 650 or img_array.shape[1] > 650:
-            # Split image into patches
-            split(filepath, patch_size=256)
-
-            # Display buffer while analyzing
-            with st.spinner('Analyzing...'):
-                # Predict on each patch
-                for patch_filename in os.listdir(patches_folder):
-                    if patch_filename.endswith(".png"):
-                        patch_path = os.path.join(patches_folder, patch_filename)
-                        patch_img = Image.open(patch_path)
-                        patch_tr_img = transforms(patch_img)
-                        prediction = predict(patch_tr_img)
-                        mask = (prediction > 0.5).astype(np.uint8) * 255
-                        mask_filename = f"mask_{patch_filename}"
-                        mask_filepath = os.path.join(pred_patches, mask_filename)
-                        Image.fromarray(mask).save(mask_filepath)
-
-                # Merge predicted patches
-                merged_mask_filename = f"data/generated_masks/mask_{timestamp}.png"
-                merge(pred_patches, merged_mask_filename, img_array.shape)
-
-                # Save merged mask
-                st.session_state.mask_filename = merged_mask_filename
-
-                # Clean up temporary patch files
-                st.info('Cleaning up temporary files...')
-                shutil.rmtree(patches_folder)
-                shutil.rmtree(pred_patches)
-                os.makedirs(patches_folder)  # Recreate empty folders
-                os.makedirs(pred_patches)
-                st.success('Temporary files cleaned up')
-        else:
-            # Predict on whole image
-            st.session_state.tr_img = transforms(img)
-            prediction = predict(st.session_state.tr_img)
-            mask = (prediction > 0.5).astype(np.uint8) * 255
-            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
-
+        log_image_details(timestamp, filename, os.path.basename(mask_filepath))
+        
         st.session_state.file_uploaded = True
 
     if st.session_state.file_uploaded and st.button('View result'):
@@ -202,16 +156,15 @@ def result_page():
 
     col1, col2 = st.columns(2)
 
-    # Display original image
     original_img_path = os.path.join(UPLOAD_DIR, st.session_state.filename)
+    mask_path = st.session_state.mask_filename
+
     if os.path.exists(original_img_path):
         original_img = Image.open(original_img_path)
         col1.image(original_img, caption='Original Image', use_column_width=True)
     else:
         col1.error(f"Original image file not found: {original_img_path}")
 
-    # Display predicted mask
-    mask_path = 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)
@@ -220,19 +173,15 @@ def result_page():
 
     st.subheader("Overlay with Area of Buildings (sqft)")
 
-    # Display overlayed image
     if os.path.exists(original_img_path) and os.path.exists(mask_path):
         original_np = cv2.imread(original_img_path)
         mask_np = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
         
-        # Ensure mask is binary
         _, mask_np = cv2.threshold(mask_np, 127, 255, cv2.THRESH_BINARY)
 
-        # Resize mask to match original image size if necessary
         if original_np.shape[:2] != mask_np.shape[:2]:
             mask_np = cv2.resize(mask_np, (original_np.shape[1], original_np.shape[0]))
 
-        # Process and overlay image
         overlay_img = process_and_overlay_image(original_np, mask_np, 'output.png')
         
         st.image(overlay_img, caption='Overlay Image', use_column_width=True)
@@ -255,4 +204,4 @@ def main():
         result_page()
 
 if __name__ == '__main__':
-    main()
+    main()