ForestAI-TreeExtraction / utils /advanced_extraction.py
DynamicPacific
Resolve merge conflicts and add multi-format image support
3c0c6f0
raw
history blame
3.24 kB
import os
import logging
import numpy as np
import rasterio
from rasterio.warp import transform_bounds
def extract_features_from_geotiff(image_path, output_folder, feature_type="trees"):
"""Simple feature extraction for HF Spaces."""
try:
logging.info(f"Extracting {feature_type} from {image_path}")
with rasterio.open(image_path) as src:
# Simple NDVI calculation
if src.count >= 3:
red = src.read(1).astype(float)
green = src.read(2).astype(float)
nir = src.read(4).astype(float) if src.count >= 4 else green
ndvi = np.divide(nir - red, nir + red + 1e-10)
mask = ndvi > 0.2
else:
band = src.read(1)
mask = band > np.percentile(band, 60)
# Get bounds
bounds = src.bounds
if src.crs:
west, south, east, north = transform_bounds(
src.crs, 'EPSG:4326',
bounds.left, bounds.bottom, bounds.right, bounds.top
)
else:
west, south, east, north = -74.1, 40.6, -73.9, 40.8
# Create simple features
features = []
height, width = mask.shape
grid_size = max(10, min(height, width) // 50)
feature_id = 0
for y in range(0, height, grid_size):
for x in range(0, width, grid_size):
cell = mask[y:y+grid_size, x:x+grid_size]
if np.sum(cell) > grid_size * grid_size * 0.3:
x_ratio = x / width
y_ratio = y / height
lon1 = west + x_ratio * (east - west)
lat1 = north - y_ratio * (north - south)
x2_ratio = min((x + grid_size) / width, 1.0)
y2_ratio = min((y + grid_size) / height, 1.0)
lon2 = west + x2_ratio * (east - west)
lat2 = north - y2_ratio * (north - south)
polygon_coords = [
[lon1, lat1], [lon2, lat1], [lon2, lat2], [lon1, lat2], [lon1, lat1]
]
feature = {
"type": "Feature",
"id": feature_id,
"properties": {
"feature_type": feature_type,
"confidence": 0.8
},
"geometry": {
"type": "Polygon",
"coordinates": [polygon_coords]
}
}
features.append(feature)
feature_id += 1
return {
"type": "FeatureCollection",
"features": features,
"feature_type": feature_type
}
except Exception as e:
logging.error(f"Error extracting features: {str(e)}")
return {"type": "FeatureCollection", "features": []}