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()