Update app.py

app.py

@@ -18,7 +18,7 @@ from model.transform import transforms
 from model.unet import UNET
 from Utils.area import pixel_to_sqft, process_and_overlay_image
 from Utils.convert import read_pansharpened_rgb
-
+from clean_refine import clean_mask, refine_mask
 
 from huggingface_hub import HfApi, login
 import os
@@ -66,7 +66,7 @@ def save_to_hf_repo(local_path, repo_path):
         st.exception(e)
 
 
-BASE_DIR = os.getcwd()
+BASE_DIR = "/"
 
 # Define subdirectories
 UPLOAD_DIR = os.path.join(BASE_DIR, "uploaded_images")
@@ -255,7 +255,10 @@ def upload_page():
                     img_transformed = transforms(img)
                     prediction = predict(img_transformed)
                     full_mask = (prediction > 0.5).astype(np.uint8) * 255
-            
+
+                    
+            full_mask = clean_mask(full_mask, morph_kernel_size=3, min_object_size=50)
+            full_mask = refine_mask(full_mask, blur_kernel=5, edge_kernel=3, threshold_value=127)
             # Save the full mask
             mask_filename = f"mask_{timestamp}.png"
             mask_filepath = os.path.join(MASK_DIR, mask_filename)