Update app.py

app.py

@@ -18,6 +18,7 @@ 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
 
 # Initialize Hugging Face API
@@ -34,7 +35,7 @@ REPO_ID = "Pavan2k4/Building_area"
 REPO_TYPE = "space"
 
 # Define base directory for Hugging Face Spaces
-BASE_DIR = "DATA/"
+BASE_DIR = "/home/user"
 
 # Define subdirectories
 UPLOAD_DIR = os.path.join(BASE_DIR, "uploaded_images")
@@ -43,48 +44,6 @@ PATCHES_DIR = os.path.join(BASE_DIR, "patches")
 PRED_PATCHES_DIR = os.path.join(BASE_DIR, "pred_patches")
 CSV_LOG_PATH = os.path.join(BASE_DIR, "image_log.csv")
 
-def split(image, destination=PATCHES_DIR, patch_size=650):
-    img = cv2.imread(image)
-    h, w, _ = img.shape
-    for y in range(0, h, patch_size):
-        for x in range(0, w, patch_size):
-            patch = img[y:y+patch_size, x:x+patch_size]
-            patch_filename = f"patch_{y}_{x}.png" 
-            patch_path = os.path.join(destination, patch_filename)
-            cv2.imwrite(patch_path, patch)
-
-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)
-
 # Create directories
 for directory in [UPLOAD_DIR, MASK_DIR, PATCHES_DIR, PRED_PATCHES_DIR]:
     os.makedirs(directory, exist_ok=True)
@@ -230,10 +189,10 @@ def upload_page():
 
                     # Clean up temporary patch files
                     st.info('Cleaning up temporary files...')
-                    for file in os.listdir(PATCHES_DIR):
-                        os.remove(os.path.join(PATCHES_DIR, file))
-                    for file in os.listdir(PRED_PATCHES_DIR):
-                        os.remove(os.path.join(PRED_PATCHES_DIR, file))
+                    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
@@ -320,10 +279,8 @@ def result_page():
         st.error("Image or mask file not found for overlay.")
 
     if st.button('Back to Upload'):
-        for file in os.listdir(PATCHES_DIR):
-            os.remove(os.path.join(PATCHES_DIR, file))
-        for file in os.listdir(PRED_PATCHES_DIR):
-            os.remove(os.path.join(PRED_PATCHES_DIR, file))
+        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
@@ -339,4 +296,7 @@ def main():
     if st.session_state.page == 'upload':
         upload_page()
     elif st.session_state.page == 'result':
-        result_page()
+        result_page()
+
+if __name__ == '__main__':
+    main()