Create app.py

app.py

@@ -0,0 +1,197 @@
+import gradio as gr
+import geopandas as gpd
+import leafmap.foliumap as leafmap
+from PIL import Image
+import rasterio
+from rasterio.windows import Window
+from tqdm import tqdm
+import io
+import zipfile
+import os
+import albumentations as albu
+import segmentation_models_pytorch as smp
+from albumentations.pytorch.transforms import ToTensorV2
+from shapely.geometry import shape
+from shapely.ops import unary_union
+from rasterio.features import shapes
+import torch
+import numpy as np
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ENCODER = 'se_resnext50_32x4d'
+ENCODER_WEIGHTS = 'imagenet'
+
+# Load and prepare the model
+def load_model():
+    model = torch.load('deeplabv3+ v15.pth', map_location=DEVICE)
+    model.eval().float()
+    return model
+
+best_model = load_model()
+
+def to_tensor(x, **kwargs):
+    return x.astype('float32')
+
+# Preprocessing
+preprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)
+
+def get_preprocessing(tile_size):
+    _transform = [
+        albu.PadIfNeeded(min_height=tile_size, min_width=tile_size, always_apply=True),
+        albu.Lambda(image=preprocessing_fn),
+        albu.Lambda(image=to_tensor, mask=to_tensor),
+        ToTensorV2(),
+    ]
+    return albu.Compose(_transform)
+
+def extract_tiles(map_file, model, tile_size=512, overlap=0, batch_size=4, threshold=0.6):
+    preprocess = get_preprocessing(tile_size)
+    tiles = []
+
+    with rasterio.open(map_file) as src:
+        height = src.height
+        width = src.width
+        effective_tile_size = tile_size - overlap
+
+        for y in stqdm(range(0, height, effective_tile_size)):
+            for x in range(0, width, effective_tile_size):
+                batch_images = []
+                batch_metas = []
+
+                for i in range(batch_size):
+                    curr_y = y + (i * effective_tile_size)
+                    if curr_y >= height:
+                        break
+
+                    window = Window(x, curr_y, tile_size, tile_size)
+                    out_image = src.read(window=window)
+
+                    if out_image.shape[0] == 1:
+                        out_image = np.repeat(out_image, 3, axis=0)
+                    elif out_image.shape[0] != 3:
+                        raise ValueError("The number of channels in the image is not supported")
+
+                    out_image = np.transpose(out_image, (1, 2, 0))
+                    tile_image = Image.fromarray(out_image.astype(np.uint8))
+
+                    out_meta = src.meta.copy()
+                    out_meta.update({
+                        "driver": "GTiff",
+                        "height": tile_size,
+                        "width": tile_size,
+                        "transform": rasterio.windows.transform(window, src.transform)
+                    })
+
+                    tile_image = np.array(tile_image)
+                    preprocessed_tile = preprocess(image=tile_image)['image']
+                    batch_images.append(preprocessed_tile)
+                    batch_metas.append(out_meta)
+
+                if not batch_images:
+                    break
+
+                batch_tensor = torch.cat([img.unsqueeze(0).to(DEVICE) for img in batch_images], dim=0)
+                with torch.no_grad():
+                    batch_masks = model(batch_tensor)
+
+                batch_masks = torch.sigmoid(batch_masks)
+                batch_masks = (batch_masks > threshold).float()
+
+                for j, mask_tensor in enumerate(batch_masks):
+                    mask_resized = torch.nn.functional.interpolate(mask_tensor.unsqueeze(0),
+                                                                   size=(tile_size, tile_size), mode='bilinear',
+                                                                   align_corners=False).squeeze(0)
+
+                    mask_array = mask_resized.squeeze().cpu().numpy()
+
+                    if mask_array.any() == 1:
+                        tiles.append([mask_array, batch_metas[j]])
+
+    return tiles
+
+def create_vector_mask(tiles, output_path):
+    all_polygons = []
+    for mask_array, meta in tiles:
+        # Ensure mask is binary
+        mask_array = (mask_array > 0).astype(np.uint8)
+
+        # Get shapes from the mask
+        mask_shapes = list(shapes(mask_array, mask=mask_array, transform=meta['transform']))
+
+        # Convert shapes to Shapely polygons
+        polygons = [shape(geom) for geom, value in mask_shapes if value == 1]
+
+        all_polygons.extend(polygons)
+    # Perform union of all polygons
+    union_polygon = unary_union(all_polygons)
+    # Create a GeoDataFrame
+    gdf = gpd.GeoDataFrame({'geometry': [union_polygon]}, crs=meta['crs'])
+    # Save to file
+    gdf.to_file(output_path)
+
+    # Calculate area in square meters
+    area_m2 = gdf.to_crs(epsg=3857).area.sum()
+
+    return gdf, area_m2
+
+def display_map(shapefile_path, tif_path):
+    # Create a leafmap centered on the shapefile bounds
+    mask = gpd.read_file(shapefile_path)
+    if mask.crs is None or mask.crs.to_string() != 'EPSG:3857':
+        mask = mask.to_crs('EPSG:3857')
+    bounds = mask.total_bounds
+    center = [(bounds[1] + bounds[3]) / 2, (bounds[0] + bounds[2]) / 2]
+    m = leafmap.Map(center=[center[1], center[0]], zoom=10, crs='EPSG3857')
+    m.add_gdf(mask, layer_name="Shapefile Mask")
+    m.add_raster(tif_path, layer_name="Satellite Image", rgb=True, opacity=0.9)
+    return m
+
+def process_file(tif_file, resolution, overlap, threshold):
+    with open("temp.tif", "wb") as f:
+        f.write(tif_file.read())
+
+    best_model.float()
+    tiles = extract_tiles("temp.tif", best_model, tile_size=resolution, overlap=overlap, batch_size=4, threshold=threshold)
+
+    output_path = "output_mask.shp"
+    result_gdf, area_m2 = create_vector_mask(tiles, output_path)
+
+    # Create zip file for shapefile
+    shp_files = [f for f in os.listdir() if f.startswith("output_mask") and f.endswith((".shp", ".shx", ".dbf", ".prj"))]
+    with io.BytesIO() as zip_buffer:
+        with zipfile.ZipFile(zip_buffer, 'a', zipfile.ZIP_DEFLATED, False) as zip_file:
+            for file in shp_files:
+                zip_file.write(file)
+        zip_buffer.seek(0)
+        with open("output_mask.zip", "wb") as f:
+            f.write(zip_buffer.getvalue())
+
+    # Display map
+    map_html = display_map("output_mask.shp", "temp.tif").to_html()
+
+    # Clean up temporary files
+    os.remove("temp.tif")
+    for file in shp_files:
+        os.remove(file)
+
+    return f"Total area occupied by PV panels: {area_m2:.4f} m^2", "output_mask.zip", map_html
+
+iface = gr.Interface(
+    fn=process_file,
+    inputs=[
+        gr.File(label="Upload TIF file"),
+        gr.Radio([512, 1024], label="Processing resolution", value=512),
+        gr.Slider(50, 150, value=100, step=25, label="Overlap"),
+        gr.Slider(0.1, 0.9, value=0.6, step=0.01, label="Threshold")
+    ],
+    outputs=[
+        gr.Textbox(label="Result"),
+        gr.File(label="Download Shapefile"),
+        gr.HTML(label="Map")
+    ],
+    title="PV Segmentor",
+    description="Upload a TIF file to process and segment PV panels."
+)
+
+if __name__ == "__main__":
+    iface.launch()