Update app.py

app.py

@@ -158,25 +158,39 @@ def get_transform_from_tif(tif_path):
         crs = src.crs
     return transform, crs
     
-def mask_to_polygons(mask, transform, mining_class_id=1):
-    # Keep only mining class
-    mining_mask = (mask == mining_class_id).astype(np.uint8)
+def mask_to_polygons(mask, transform, target_class_id):
+    """
+    Convert a single-class mask to polygons with proper geospatial reference.
+    
+    Args:
+        mask (np.ndarray): Predicted mask (H, W) with class indices.
+        transform (Affine): Georeferencing transform from the input TIF.
+        target_class_id (int): The class ID to convert into polygons.
+        
+    Returns:
+        gpd.GeoDataFrame: Geodataframe containing polygons for the target class.
+    """
+    # Binary mask for the target class
+    binary_mask = (mask == target_class_id).astype(np.uint8)
 
-    # Extract shapes for mining
+    # Extract polygon shapes
     results = (
-        {'properties': {'class': 'mining_land'}, 'geometry': s}
-        for s, v in shapes(mining_mask, mask=None, transform=transform)
-        if v == 1
+        {'properties': {'class': f'class_{target_class_id}'}, 'geometry': geom}
+        for geom, value in shapes(binary_mask, mask=None, transform=transform)
+        if value == 1
     )
 
     geoms = list(results)
+
+    # Return empty GeoDataFrame if no shapes found
     if not geoms:
         return gpd.GeoDataFrame(columns=['geometry'], geometry='geometry', crs='EPSG:4326')
-    
+
     gdf = gpd.GeoDataFrame.from_features(geoms)
     gdf.set_crs(epsg=4326, inplace=True)
     return gdf
 
+
 def calculate_illegal_area(mining_gdf, wiup_gdf):
     # Ensure both are in same CRS
     mining_gdf = mining_gdf.to_crs(wiup_gdf.crs)
@@ -351,12 +365,17 @@ if st.button("Run Inference"):
             transform, crs = get_transform_from_tif(optic_path)
             
             # Convert mask to mining polygons
-            non_mining_gdf = mask_to_polygons(pred_label_mask, transform, mining_class_id=0)
-            mining_gdf = mask_to_polygons(pred_label_mask, transform, mining_class_id=1)
-            beach_gdf = mask_to_polygons(pred_label_mask, transform, mining_class_id=2)
+            non_mining_gdf = mask_to_polygons(pred_label_mask, transform, target_class_id=0)
+            mining_gdf = mask_to_polygons(pred_label_mask, transform, target_class_id=1)
+            beach_gdf = mask_to_polygons(pred_label_mask, transform, target_class_id=2)
+
+            st.success(f"Non-mining polygons: {len(non_mining_gdf)}")
+            st.success(f"Mining polygons: {len(mining_gdf)}")
+            st.success(f"Beach polygons: {len(beach_gdf)}")
             
             # Make sure WIUP and prediction are in same CRS
             wiup_gdf = wiup_gdf.to_crs(mining_gdf.crs)
+            st.success(f"WIUP CRS: {wiup_gdf.crs}")
             
             # Find area
             non_mining_area, non_mining_area_gdf = calculate_illegal_area(non_mining_gdf, wiup_gdf)