ForestAI-TreeExtraction / utils /advanced_extraction.py
DynamicPacific
add area detection
824dd08
import os
import logging
import numpy as np
import rasterio
from rasterio.warp import transform_bounds
import cv2
from scipy import ndimage
def extract_features_from_geotiff(image_path, output_folder, feature_type="trees"):
"""Improved feature extraction with contour-based tree detection."""
try:
logging.info(f"Extracting {feature_type} from {image_path}")
with rasterio.open(image_path) as src:
logging.info(f"Image shape: {src.shape}, bands: {src.count}")
# Enhanced 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
logging.info(f"Red band range: {red.min():.3f} to {red.max():.3f}")
logging.info(f"Green band range: {green.min():.3f} to {green.max():.3f}")
logging.info(f"NIR band range: {nir.min():.3f} to {nir.max():.3f}")
# Improved NDVI calculation
ndvi = np.divide(nir - red, nir + red + 1e-10)
# Adaptive thresholding based on image statistics
ndvi_mean = np.mean(ndvi)
ndvi_std = np.std(ndvi)
threshold = max(0.05, ndvi_mean + 0.1 * ndvi_std) # More lenient threshold
logging.info(f"NDVI mean: {ndvi_mean:.3f}, std: {ndvi_std:.3f}, threshold: {threshold:.3f}")
mask = ndvi > threshold
logging.info(f"Mask pixels above threshold: {np.sum(mask)} out of {mask.size}")
else:
band = src.read(1)
logging.info(f"Single band range: {band.min():.3f} to {band.max():.3f}")
# Adaptive threshold for single-band images
threshold = np.percentile(band, 40) # Reduced from 70 to be more lenient
logging.info(f"Single band threshold: {threshold:.3f}")
mask = band > threshold
logging.info(f"Mask pixels above threshold: {np.sum(mask)} out of {mask.size}")
# 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
# Convert to uint8 for OpenCV processing
mask_uint8 = (mask * 255).astype(np.uint8)
logging.info(f"Mask uint8 range: {mask_uint8.min()} to {mask_uint8.max()}")
# Morphological operations to clean up the mask
kernel = np.ones((3, 3), np.uint8)
mask_cleaned = cv2.morphologyEx(mask_uint8, cv2.MORPH_CLOSE, kernel)
mask_cleaned = cv2.morphologyEx(mask_cleaned, cv2.MORPH_OPEN, kernel)
logging.info(f"After morphological ops: {np.sum(mask_cleaned > 0)} pixels")
# Find contours instead of using grid
contours, _ = cv2.findContours(mask_cleaned, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
logging.info(f"Found {len(contours)} contours")
# Filter contours by area and create features
features = []
height, width = mask.shape
min_area = (height * width) * 0.0001 # Minimum 0.01% of image area (reduced from 0.1%)
max_area = (height * width) * 0.2 # Maximum 20% of image area (increased from 10%)
logging.info(f"Area filter: min={min_area:.0f}, max={max_area:.0f}")
for i, contour in enumerate(contours):
area = cv2.contourArea(contour)
logging.info(f"Contour {i}: area={area:.0f}")
# Filter by area
if area < min_area or area > max_area:
logging.info(f"Contour {i}: filtered out (area {area:.0f} not in range)")
continue
# Simplify contour for better polygon shape
epsilon = 0.005 * cv2.arcLength(contour, True)
approx = cv2.approxPolyDP(contour, epsilon, True)
# Convert contour points to geographic coordinates
polygon_coords = []
for point in approx:
x, y = point[0]
# Convert pixel coordinates to geographic coordinates
x_ratio = x / width
y_ratio = y / height
lon = west + x_ratio * (east - west)
lat = north - y_ratio * (north - south)
polygon_coords.append([lon, lat])
# Close the polygon
if len(polygon_coords) > 2:
polygon_coords.append(polygon_coords[0])
# Calculate confidence based on area and shape
area_ratio = area / (height * width)
confidence = min(0.95, 0.5 + area_ratio * 10) # Adaptive confidence
feature = {
"type": "Feature",
"id": i,
"properties": {
"feature_type": feature_type,
"confidence": round(confidence, 2),
"area_pixels": int(area),
"area_ratio": round(area_ratio, 4)
},
"geometry": {
"type": "Polygon",
"coordinates": [polygon_coords]
}
}
features.append(feature)
logging.info(f"Contour {i}: added as feature with {len(polygon_coords)} points")
logging.info(f"Extracted {len(features)} tree features")
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": []}