Update app.py

app.py

@@ -18,19 +18,13 @@ 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 huggingface_hub import HfApi, login
-import os
-
-# Set up Hugging Face authentication
 
 HF_TOKEN = os.environ.get("HF_TOKEN")
 if not HF_TOKEN:
     raise ValueError("HF_TOKEN environment variable is not set")
 
-login(token=HF_TOKEN, add_to_git_credential=True)
-
+login(token=HF_TOKEN)
 hf_api = HfApi()
 
 REPO_ID = "Pavan2k4/Building_area"
@@ -42,9 +36,36 @@ def load_model():
     model.load_state_dict(torch.load('latest.pth', map_location='cpu', weights_only = True)['model_state_dict'])
     model.eval()
     return model
+# Load model
+model = load_model()
 
 
 
+def refine_mask(mask, blur_kernel=5, threshold_value=127, morph_kernel_size=3, min_object_size=100):
+    """Refine and clean the mask with Gaussian blur, thresholding, morphological operations, and small object removal."""
+    
+    # Ensure mask is grayscale
+    if len(mask.shape) > 2:
+        mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
+    
+    # Apply Gaussian blur to smooth edges
+    mask = cv2.GaussianBlur(mask, (blur_kernel, blur_kernel), 0)
+    
+    # Apply binary threshold
+    _, mask = cv2.threshold(mask, threshold_value, 255, cv2.THRESH_BINARY)
+    
+    # Apply morphological operations (opening and closing)
+    kernel = np.ones((morph_kernel_size, morph_kernel_size), np.uint8)
+    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
+    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
+    
+    # Remove small objects based on area
+    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=8)
+    for i in range(1, num_labels):
+        if stats[i, cv2.CC_STAT_AREA] < min_object_size:
+            mask[labels == i] = 0
+    
+    return mask
 
 
 
@@ -69,20 +90,17 @@ def save_to_hf_repo(local_path, repo_path):
 BASE_DIR = os.getcwd()
 
 # Define subdirectories
-UPLOAD_DIR = os.path.join(BASE_DIR, "uploaded_images")
-MASK_DIR = os.path.join(BASE_DIR, "generated_masks")
-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")
+UPLOAD_DIR =  "uploaded_images"
+MASK_DIR =  "generated_masks"
+PATCHES_DIR = "patches"
+PRED_PATCHES_DIR =  "pred_patches"
+CSV_LOG_PATH =  "image_log.csv"
 
 # Create directories
 for directory in [UPLOAD_DIR, MASK_DIR, PATCHES_DIR, PRED_PATCHES_DIR]:
     os.makedirs(directory, exist_ok=True)
 
-# Load model
-
 
-model = load_model()
 
 def predict(image):
     with torch.no_grad():
@@ -218,7 +236,7 @@ def upload_page():
             #st.success(f"Image saved to {filepath}")
             
             # Save image to Hugging Face repo----------------------------------------------------------------------------------------------------------------------------------
-            # Save image to Hugging Face repo
+          
             try:
                 image_repo_path = f"images/{converted_filename}"
                 save_to_hf_repo(converted_filepath, image_repo_path)
@@ -256,16 +274,16 @@ def upload_page():
                     prediction = predict(img_transformed)
                     full_mask = (prediction > 0.5).astype(np.uint8) * 255
 
-                    
+                  
             
-            # Save the full mask
+            # Save the full mask---------------------------------------------------------------------------------------------------
+            full_mask = refine_mask(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
 
-            # Save mask to Hugging Face repo---------------------------------------------------------------------------------------------
-            # Save mask to Hugging Face repo
+   
             try:
                 mask_repo_path = f"masks/{mask_filename}"
                 save_to_hf_repo(mask_filepath, mask_repo_path)