Update app.py

app.py

@@ -17,11 +17,10 @@ from model.CBAM.reunet_cbam import reunet_cbam
 from model.transform import transforms
 from model.unet import UNET
 from Utils.area import pixel_to_sqft, process_and_overlay_image
-from split_merge import split, merge 
 from Utils.convert import read_pansharpened_rgb
 
 # Define base directory for Hugging Face Spaces
-BASE_DIR = "/home/user"
+BASE_DIR = "/Data"
 
 # Define subdirectories
 UPLOAD_DIR = os.path.join(BASE_DIR, "uploaded_images")
@@ -49,6 +48,76 @@ def predict(image):
         output = model(image.unsqueeze(0))
     return output.squeeze().cpu().numpy()
 
+def split_image(image, patch_size=512):
+    h, w, _ = image.shape
+    patches = []
+    for y in range(0, h, patch_size):
+        for x in range(0, w, patch_size):
+            patch = image[y:min(y+patch_size, h), x:min(x+patch_size, w)]
+            patches.append((f"patch_{y}_{x}.png", patch))
+    return patches
+
+def merge(patch_folder, dest_image='out.png', image_shape=None):
+    merged = np.zeros(image_shape[:-1] + (3,), dtype=np.uint8)
+    for filename in os.listdir(patch_folder):
+        if filename.endswith(".png"):
+            patch_path = os.path.join(patch_folder, filename)
+            patch = cv2.imread(patch_path)
+            patch_height, patch_width, _ = patch.shape
+
+            # Extract patch coordinates from filename
+            parts = filename.split("_")
+            x, y = None, None
+            for part in parts:
+                if part.endswith(".png"):
+                    x = int(part.split(".")[0])
+                elif part.isdigit():
+                    y = int(part)
+            if x is None or y is None:
+                raise ValueError(f"Invalid filename: {filename}")
+
+            # Check if patch fits within image boundaries
+            if x + patch_width > image_shape[1] or y + patch_height > image_shape[0]:
+                # Adjust patch position to fit within image boundaries
+                if x + patch_width > image_shape[1]:
+                    x = image_shape[1] - patch_width
+                if y + patch_height > image_shape[0]:
+                    y = image_shape[0] - patch_height
+
+            # Merge patch into the main image
+            merged[y:y+patch_height, x:x+patch_width, :] = patch
+
+    cv2.imwrite(dest_image, merged)
+    return merged
+
+def process_large_image(model, image_path, patch_size=512):
+    # Read the image
+    img = cv2.imread(image_path)
+    if img is None:
+        raise ValueError(f"Failed to read image from {image_path}")
+
+    h, w, _ = img.shape
+    st.write(f"Processing image of size {w}x{h}")
+
+    # Split the image into patches
+    patches = split_image(img, patch_size)
+
+    # Process each patch
+    for filename, patch in patches:
+        patch_pil = Image.fromarray(cv2.cvtColor(patch, cv2.COLOR_BGR2RGB))
+        patch_transformed = transforms(patch_pil)
+        prediction = predict(patch_transformed)
+        mask = (prediction > 0.5).astype(np.uint8) * 255
+        
+        # Save the mask patch
+        mask_filepath = os.path.join(PRED_PATCHES_DIR, filename)
+        cv2.imwrite(mask_filepath, mask)
+
+    # Merge the predicted patches
+    merged_mask = merge(PRED_PATCHES_DIR, dest_image='merged_mask.png', image_shape=img.shape)
+
+    return merged_mask
+
 def log_image_details(image_id, image_filename, mask_filename):
     file_exists = os.path.exists(CSV_LOG_PATH)
     
@@ -122,55 +191,30 @@ def upload_page():
             # Store the full path of the converted image
             st.session_state.filename = converted_filename
 
-            # Convert image to numpy array
-            img_array = np.array(img)
-
-            # 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(converted_filepath, patch_size=512)
-
-                # Display buffer while analyzing
-                with st.spinner('Analyzing...'):
-                    # Predict on each patch
-                    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)
-
-                    # Merge predicted patches
-                    merged_mask_filename = f"mask_{timestamp}.png"
-                    merged_mask_path = os.path.join(MASK_DIR, merged_mask_filename)
-                    merge(PRED_PATCHES_DIR, merged_mask_path, 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_DIR)
-                    shutil.rmtree(PRED_PATCHES_DIR)
-                    os.makedirs(PATCHES_DIR)  # Recreate empty folders
-                    os.makedirs(PRED_PATCHES_DIR)
-                    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_filename
+            # Process the image
+            st.write("Processing image...")
+            with st.spinner('Analyzing...'):
+                full_mask = process_large_image(model, converted_filepath)
+            
+            # Save the full mask
+            mask_filename = f"mask_{timestamp}.png"
+            mask_filepath = os.path.join(MASK_DIR, mask_filename)
+            cv2.imwrite(mask_filepath, full_mask)
+            st.session_state.mask_filename = mask_filename
+
+            st.success("Image processed successfully")
+
+            # Log image details
+            log_image_details(timestamp, converted_filename, mask_filename)
 
             st.session_state.file_uploaded = True
             
+            # Clean up temporary patch files
+            st.info('Cleaning up temporary files...')
+            for file in os.listdir(PRED_PATCHES_DIR):
+                os.remove(os.path.join(PRED_PATCHES_DIR, file))
+            st.success('Temporary files cleaned up')
+            
         except Exception as e:
             st.error(f"An error occurred: {str(e)}")
             st.error("Please check the logs for more details.")
@@ -237,8 +281,6 @@ def result_page():
         st.error("Image or mask file not found for overlay.")
 
     if st.button('Back to Upload'):
-        shutil.rmtree(PATCHES_DIR)
-        shutil.rmtree(PRED_PATCHES_DIR)
         st.session_state.page = 'upload'
         st.session_state.file_uploaded = False
         st.session_state.filename = None