conda create -n geoai-sam2 python=3.10 -y
conda activate geoai-sam2

pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
pip install transformers accelerate datasets
pip install rasterio geopandas shapely fiona scikit-image matplotlib

import torch
print(f"CUDA Available: {torch.cuda.is_available()}")
print(f"Device Name: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU'}")

import rasterio
from rasterio.windows import Window
import numpy as np
import torch
from PIL import Image
from transformers import Sam2Model, Sam2Processor

# Load the georeferenced aerial orthophoto
raster_path = "data/nairobi_ortho_0.2m.tif"

with rasterio.open(raster_path) as src:
    meta = src.meta.copy()
    crs = src.crs
    transform = src.transform
    # Read RGB bands (assuming bands 1, 2, 3 correspond to R, G, B)
    img_array = src.read([1, 2, 3])
    # Reshape from (Bands, Height, Width) to (Height, Width, Bands) for PIL/Transformers
    img_array = np.moveaxis(img_array, 0, -1)
    
# Initialize SAM 2 Model and Processor from Hugging Face
model_id = "facebook/sam2.1-hiera-large"
device = "cuda" if torch.cuda.is_available() else "cpu"

print(f"Loading SAM 2 model ({model_id}) onto {device}...")
processor = Sam2Processor.from_pretrained(model_id)
model = Sam2Model.from_pretrained(model_id).to(device)

# Convert numpy array to PIL Image for the processor
pil_image = Image.fromarray(img_array.astype('uint8'))

# Define bounding boxes [x_min, y_min, x_max, y_max] for target buildings in pixel coordinates
input_boxes = [
    [450, 320, 680, 510],
    [720, 610, 950, 800]
]

# Format inputs for SAM 2 processor
inputs = processor(
    images=pil_image,
    input_boxes=[input_boxes],
    return_tensors="pt"
).to(model.device)

# Perform inference without gradient tracking
with torch.no_grad():
    outputs = model(**inputs)

# Post-process masks back to original image size
masks = processor.post_process_masks(
    outputs.pred_masks.cpu(),
    inputs["original_sizes"].cpu()
)

# Extract the highest scoring binary mask tensor
# masks[0] shape: (num_boxes, num_masks_per_box, H, W)
pred_masks = masks[0][:, 0, :, :].numpy()
binary_mask = (pred_masks.sum(axis=0) > 0).astype(np.uint8)

from rasterio.features import shapes
import geopandas as gpd
from shapely.geometry import shape

# Extract shapes (polygons and values) from the binary mask array
# We apply the raster's affine transform to map pixels to geospatial coordinates
shape_generator = shapes(
    binary_mask,
    transform=transform
)

# Convert generator output into a list of Shapely geometries and attributes
records = []
for geom, value in shape_generator:
    if value == 1:  # Keep foreground building features
        polygon = shape(geom)
        records.append({"geometry": polygon})

# Construct initial GeoDataFrame in native CRS
gdf = gpd.GeoDataFrame.from_features(records, crs=crs)
print(f"Extracted {len(gdf)} raw building polygons.")

from shapely.geometry import Polygon

def smooth_and_clean_polygon(poly, tolerance=0.5, buffer_dist=0.6):
    """
    Applies negative/positive buffering (morphological closing/opening equivalent)
    to fill small gaps, followed by Douglas-Peucker simplification.
    """
    if not poly.is_valid:
        poly = poly.buffer(0)
    # Morphological smoothing via buffering
    smoothed = poly.buffer(buffer_dist).buffer(-buffer_dist)
    # Simplify geometry
    simplified = smoothed.simplify(tolerance, preserve_topology=True)
    return simplified

# Apply cleaning function to all geometries
gdf["geometry"] = gdf["geometry"].apply(lambda p: smooth_and_clean_polygon(p, tolerance=0.4, buffer_dist=0.5))

# Filter out sliver polygons falling below minimum area threshold (e.g., 5 square meters)
gdf = gdf[gdf.geometry.area >= 5.0].reset_index(drop=True)

# Export to production GeoPackage layer for QGIS consumption
output_gpkg = "outputs/extracted_buildings_sam2.gpkg"
gdf.to_file(output_gpkg, layer="sam2_buildings", driver="GPKG")
print(f"Successfully exported cleaned vector footprints to {output_gpkg}")