Update app.py

app.py

@@ -23,7 +23,7 @@ SAR_SHAPE = (HEIGHT, WIDTH, 1)
 OPTIC_SHAPE = (HEIGHT, WIDTH, 3)
 MASK_SHAPE = (HEIGHT, WIDTH, 4)  # One-hot encoded masks with 4 classes
 
-# Class colors: non-mining (green), illegal mining (red), beach (black)
+# Class colors: non-mining (green), mining (red), beach (black)
 CLASS_COLORS = [
     [115/255, 178/255, 115/255],
     [1, 0, 0],
@@ -349,23 +349,23 @@ 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)
             
             # 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)
+            # Find area
+            non_mining_area, non_mining_area_gdf = calculate_illegal_area(non_mining_gdf, wiup_gdf)
+            mining_area, mining_area_gdf = calculate_illegal_area(mining_gdf, wiup_gdf)
+            beach_area, beach_area_gdf = calculate_illegal_area(beach_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")
-
+            st.success(f" Non Mining Area : {non_mining_area/1e6:.2f} sq. km")
+            st.success(f" Mining Area : {mining_area/1e6:.2f} sq. km")
+            st.success(f" Beach Area : {beach_area/1e6:.2f} sq. km")
+                       
             
         else:
             st.warning("Please upload all three .tiff files to proceed.")