Update app.py

app.py

@@ -141,6 +141,40 @@ def save_uploaded_file(uploaded_file, suffix=".tif"):
         tmp.write(uploaded_file.read())
         return tmp.name  
 
+def get_transform_from_tif(tif_path):
+    """
+    Extracts the affine transform and CRS from a GeoTIFF image.
+
+    Args:
+        tif_path (str): Path to the GeoTIFF file.
+
+    Returns:
+        transform (Affine): Affine transformation for pixel to coordinate mapping.
+        crs (CRS): Coordinate Reference System of the image.
+    """
+    with rasterio.open(tif_path) as src:
+        transform = src.transform
+        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)
+
+    # Extract shapes for mining
+    results = (
+        {'properties': {'class': 'mining_land'}, 'geometry': s}
+        for s, v in shapes(mining_mask, mask=None, transform=transform)
+        if v == 1
+    )
+
+    geoms = list(results)
+    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
 
 # Streamlit App Title
 st.title("Satellite Mining Segmentation: SAR + Optic Image Inference")
@@ -289,27 +323,32 @@ if st.button("Run Inference"):
             plt.tight_layout()
             st.pyplot(fig)
 
-            wiup_mask = features.rasterize(
-                [(geom, 1) for geom in wiup_gdf.geometry],
-                out_shape=(HEIGHT, WIDTH),
-                transform=transform,
-                fill=0,
-                dtype=np.uint8
-            )
-            # Binary mask of predicted illegal mining
-            pred_illegal_mask = illegal_mask.astype(np.uint8)
-
-            # Mining outside WIUP
-            outside_mask = (pred_illegal_mask == 1) & (wiup_mask == 0)
-            outside_percentage = 100 * np.sum(outside_mask) / np.sum(pred_illegal_mask)
-
-            st.markdown(f"### 🚨 Illegal Mining Outside WIUP: `{outside_percentage:.2f}%`")
+            transform, crs = get_transform_from_tif(optic_path)
+            
+            # After getting prediction from model
+            pred_label_mask = np.argmax(pred_masks[0], axis=-1)  # shape: (H, W)
+            
+            # Get georeference transform from TIF image
+            transform, crs = get_transform_from_tif(optic_path)
+            
+            # Convert mask to mining polygons
+            mining_gdf = mask_to_polygons(pred_label_mask, transform, mining_class_id=1)
+            
+            # Make sure WIUP and prediction are in same CRS
+            wiup_gdf = wiup_gdf.to_crs(mining_gdf.crs)
+            
+            # Find area outside WIUP
+            illegal_area, illegal_area_gdf = calculate_illegal_area(mining_gdf, wiup_gdf)
+            
+            # Display in Streamlit
+            st.success(f"🛑 Illegal Mining Area Outside WIUP: {illegal_area/1e6:.2f} sq. km")
+            
+            # Optional: Visualize and allow download
+            st.map(illegal_area_gdf.to_crs(epsg=4326))
+            
+            geojson = illegal_area_gdf.to_crs(epsg=4326).to_json()
+            st.download_button("Download GeoJSON", geojson, file_name="illegal_area.geojson", mime="application/json")
 
-            fig2, ax2 = plt.subplots(figsize=(10, 6))
-            ax2.imshow(outside_mask, cmap='Reds')
-            ax2.set_title("Illegal Mining Outside WIUP Boundary")
-            ax2.axis('off')
-            st.pyplot(fig2)
             
         else:
             st.warning("Please upload all three .tiff files to proceed.")