361 lines
15 KiB
Python
361 lines
15 KiB
Python
import os
|
|
import io
|
|
import json
|
|
import time
|
|
import copy
|
|
import joblib
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.data import Dataset, DataLoader
|
|
import pandas as pd
|
|
import numpy as np
|
|
from sklearn.neighbors import KNeighborsRegressor
|
|
from catboost import CatBoostClassifier
|
|
|
|
# Digital Earth Africa STAC specific imports
|
|
try:
|
|
from pystac_client import Client
|
|
import odc.stac
|
|
import xarray as xr
|
|
import rioxarray
|
|
except ImportError:
|
|
Client = None
|
|
odc = None
|
|
xr = None
|
|
rioxarray = None
|
|
|
|
# ==========================================
|
|
# 1. CPU-OPTIMIZED ARCHITECTURES
|
|
# ==========================================
|
|
|
|
class TemporalFCN(nn.Module):
|
|
def __init__(self, num_bands, num_classes):
|
|
super().__init__()
|
|
self.conv_block1 = nn.Sequential(
|
|
nn.Conv1d(num_bands, 64, kernel_size=5, padding=2),
|
|
nn.BatchNorm1d(64),
|
|
nn.ReLU()
|
|
)
|
|
self.conv_block2 = nn.Sequential(
|
|
nn.Conv1d(64, 128, kernel_size=3, padding=1),
|
|
nn.BatchNorm1d(128),
|
|
nn.ReLU()
|
|
)
|
|
self.global_avg_pool = nn.AdaptiveAvgPool1d(1)
|
|
self.fc = nn.Linear(128, num_classes)
|
|
|
|
def forward(self, x, return_features=False):
|
|
x = self.conv_block1(x)
|
|
x = self.conv_block2(x)
|
|
features = self.global_avg_pool(x).squeeze(-1)
|
|
out = self.fc(features)
|
|
if return_features:
|
|
return out, features
|
|
return out
|
|
|
|
class SmallGRU(nn.Module):
|
|
def __init__(self, num_bands, num_classes, hidden_size=64):
|
|
super().__init__()
|
|
self.gru = nn.GRU(input_size=num_bands, hidden_size=hidden_size, num_layers=1, batch_first=True)
|
|
self.fc = nn.Linear(hidden_size, num_classes)
|
|
|
|
def forward(self, x, return_features=False):
|
|
x = x.transpose(1, 2)
|
|
out, _ = self.gru(x)
|
|
features = out[:, -1, :]
|
|
final_out = self.fc(features)
|
|
if return_features:
|
|
return final_out, features
|
|
return final_out
|
|
|
|
# ==========================================
|
|
# 2. DATA PREPARATION & PYTORCH UTILS
|
|
# ==========================================
|
|
|
|
class CropDataset(Dataset):
|
|
def __init__(self, X, y, augment=False):
|
|
self.X = torch.FloatTensor(X)
|
|
self.y = torch.LongTensor(y)
|
|
self.augment = augment
|
|
|
|
def __len__(self):
|
|
return len(self.y)
|
|
|
|
def __getitem__(self, idx):
|
|
x = self.X[idx].clone()
|
|
if self.augment:
|
|
if torch.rand(1).item() > 0.5:
|
|
noise = torch.randn_like(x) * 0.03
|
|
x = x + noise
|
|
if torch.rand(1).item() > 0.7:
|
|
seq_len = x.shape[1]
|
|
t_idx = torch.randint(0, seq_len, (1,)).item()
|
|
x[:, t_idx] = 0.0
|
|
return x, self.y[idx]
|
|
|
|
def prepare_tensors(df, bands, dates):
|
|
num_samples = len(df)
|
|
X_3d = np.zeros((num_samples, len(bands), len(dates)), dtype=np.float32)
|
|
for b_idx, band in enumerate(bands):
|
|
for d_idx, date in enumerate(dates):
|
|
col = f"{date}_{band}"
|
|
if col in df.columns:
|
|
X_3d[:, b_idx, d_idx] = df[col].values
|
|
|
|
means = X_3d.mean(axis=2, keepdims=True)
|
|
stds = X_3d.std(axis=2, keepdims=True) + 1e-8
|
|
X_3d = (X_3d - means) / stds
|
|
return X_3d
|
|
|
|
# ==========================================
|
|
# 3. DIGITAL EARTH AFRICA STAC INTEGRATION
|
|
# ==========================================
|
|
|
|
class DEAfricaSTACWrapper:
|
|
def __init__(self, stac_url="https://explorer.digitalearth.africa/stac"):
|
|
if Client is None or odc is None or xr is None:
|
|
raise ImportError("Missing required libraries: pystac-client, odc-stac, xarray")
|
|
print(f"Connecting to Digital Earth Africa STAC Catalog at {stac_url}...")
|
|
self.catalog = Client.open(stac_url)
|
|
|
|
@staticmethod
|
|
def _patch_s3_url(url: str) -> str:
|
|
if url.startswith("s3://deafrica-sentinel-2"):
|
|
return url.replace(
|
|
"s3://deafrica-sentinel-2",
|
|
"/vsicurl/https://deafrica-sentinel-2.s3.af-south-1.amazonaws.com"
|
|
)
|
|
return url
|
|
|
|
def fetch_and_format_data(self, lat_range, lon_range, time_range, resolution=20):
|
|
bbox = [lon_range[0], lat_range[0], lon_range[1], lat_range[1]]
|
|
print(f"Searching STAC for Bounding Box: {bbox} over {time_range}...")
|
|
search = self.catalog.search(
|
|
collections=["s2_l2a"],
|
|
bbox=bbox,
|
|
datetime=f"{time_range[0]}/{time_range[1]}"
|
|
)
|
|
items = list(search.items())
|
|
if not items:
|
|
raise ValueError("No STAC items found for this bounding box and time range.")
|
|
|
|
print(f"Found {len(items)} STAC items. Loading into xarray...")
|
|
|
|
band_map = {
|
|
'B04': 'red',
|
|
'B03': 'green',
|
|
'B02': 'blue',
|
|
'B08': 'nir',
|
|
'B05': 'red_edge_1',
|
|
'SCL': 'scl'
|
|
}
|
|
|
|
os.environ["GDAL_DISABLE_READDIR_ON_OPEN"] = "EMPTY_DIR"
|
|
|
|
ds = odc.stac.load(
|
|
items,
|
|
measurements=list(band_map.keys()),
|
|
bbox=bbox,
|
|
crs="EPSG:6933",
|
|
resolution=resolution,
|
|
groupby="solar_day",
|
|
patch_url=self._patch_s3_url
|
|
)
|
|
|
|
# Rename bands to expected names
|
|
ds = ds.rename(band_map)
|
|
|
|
print("Masking clouds and shadows...")
|
|
valid_mask = (ds.scl == 4) | (ds.scl == 5) | (ds.scl == 6) | (ds.scl == 2) | (ds.scl == 7)
|
|
ds = ds.where(valid_mask)
|
|
ds = ds / 10000.0
|
|
|
|
print("Computing Spectral Indices (NDVI, NDRE, SAVI, EVI)...")
|
|
ds['ndvi'] = (ds.nir - ds.red) / (ds.nir + ds.red + 1e-8)
|
|
ds['ndre'] = (ds.nir - ds.red_edge_1) / (ds.nir + ds.red_edge_1 + 1e-8)
|
|
ds['savi'] = ((ds.nir - ds.red) / (ds.nir + ds.red + 0.5)) * 1.5
|
|
ds['evi'] = 2.5 * ((ds.nir - ds.red) / (ds.nir + 6 * ds.red - 7.5 * ds.blue + 1))
|
|
|
|
ds_indices = ds[['ndvi', 'ndre', 'evi', 'savi']]
|
|
|
|
print("Reshaping multi-dimensional xarray into flat Tabular DataFrame...")
|
|
df = ds_indices.compute().to_dataframe().reset_index()
|
|
df['date_str'] = df['time'].dt.strftime('%Y%m%d')
|
|
|
|
df_pivot = df.pivot(index=['y', 'x'], columns='date_str', values=['ndvi', 'ndre', 'evi', 'savi'])
|
|
df_pivot.columns = [f"{date}_{band}" for band, date in df_pivot.columns]
|
|
|
|
df_final = df_pivot.reset_index().rename(columns={'y': 'lat', 'x': 'lon'})
|
|
print(f"✅ Data Ready! {df_final.shape[0]} spatial pixels generated.")
|
|
return df_final
|
|
|
|
# ==========================================
|
|
# 4. INFERENCE PIPELINE
|
|
# ==========================================
|
|
|
|
class CropInferencePipeline:
|
|
def __init__(self, model_dir="/tmp/geocrop-cache"):
|
|
print(f"Loading Crop Inference Pipeline from {model_dir}...")
|
|
meta_path = os.path.join(model_dir, "pipeline_meta.pkl")
|
|
if not os.path.exists(meta_path):
|
|
raise FileNotFoundError(f"Pipeline metadata not found at {meta_path}")
|
|
|
|
self.meta = joblib.load(meta_path)
|
|
self.le = self.meta["le"]
|
|
self.bands = self.meta["bands"]
|
|
self.dates = self.meta["dates"]
|
|
self.w_fcn = self.meta["weights"]["w_fcn"]
|
|
self.w_cb = self.meta["weights"]["w_cb"]
|
|
|
|
self.fcn = TemporalFCN(len(self.bands), self.meta["num_classes"])
|
|
fcn_path = os.path.join(model_dir, "Temporal_FCN.pth")
|
|
self.fcn.load_state_dict(torch.load(fcn_path, map_location=torch.device('cpu')))
|
|
self.fcn.eval()
|
|
|
|
cb_path = os.path.join(model_dir, "calibrated_hybrid_cb.pkl")
|
|
self.calibrated_cb = joblib.load(cb_path)
|
|
print("Models loaded successfully.")
|
|
|
|
def _impute_inference_data(self, df):
|
|
print("Imputing cloudy/missing timesteps via temporal interpolation...")
|
|
from feature_computation import handle_temporal_gaps, spatial_fill_nan
|
|
|
|
df = df.copy()
|
|
missing_mask = {}
|
|
|
|
# Track original NaNs before any imputation
|
|
for band in self.bands:
|
|
band_cols = [f"{date}_{band}" for date in self.dates if f"{date}_{band}" in df.columns]
|
|
if band_cols:
|
|
missing_mask[band] = df[band_cols].isna().astype(float)
|
|
|
|
# Process each band: apply handle_temporal_gaps per pixel for each band
|
|
for band in self.bands:
|
|
band_cols = [f"{date}_{band}" for date in self.dates if f"{date}_{band}" in df.columns]
|
|
if band_cols:
|
|
print(f" Processing band {band} with gap handling...")
|
|
|
|
# For each pixel, apply handle_temporal_gaps to the time series
|
|
for idx in range(len(df)):
|
|
time_series = df[band_cols].iloc[idx].values.astype(np.float64)
|
|
|
|
# Apply handle_temporal_gaps: gaps >= 3 will result in NaNs for those timesteps
|
|
time_series = handle_temporal_gaps(time_series, gap_threshold=3)
|
|
df.loc[df.index[idx], band_cols] = time_series
|
|
|
|
# After gap handling, fill remaining NaNs with linear interpolation
|
|
df[band_cols] = df[band_cols].interpolate(method='linear', axis=1, limit_direction='both')
|
|
df[band_cols] = df[band_cols].ffill(axis=1).bfill(axis=1).fillna(0)
|
|
|
|
# Apply spatial fill to each band using spatial_fill_nan
|
|
# Reshape to (num_dates, num_pixels) for each band, apply spatial fill
|
|
for band in self.bands:
|
|
band_cols = [f"{date}_{band}" for date in self.dates if f"{date}_{band}" in df.columns]
|
|
if band_cols:
|
|
print(f" Applying spatial fill for band {band}...")
|
|
|
|
# Transpose to (T, H*W) for spatial filling
|
|
band_data = df[band_cols].values.T # Shape: (num_dates, num_pixels)
|
|
|
|
# Apply spatial_fill_nan per time step
|
|
for t_idx in range(band_data.shape[0]):
|
|
band_data[t_idx] = spatial_fill_nan(band_data[t_idx].reshape(-1, 1)).squeeze()
|
|
|
|
# Put back into dataframe
|
|
df[band_cols] = band_data.T
|
|
|
|
return df, missing_mask
|
|
|
|
def predict(self, raw_df, apply_spatial_smoothing=False, coord_cols=['lat', 'lon']):
|
|
df, missing_mask = self._impute_inference_data(raw_df)
|
|
X_infer = prepare_tensors(df, self.bands, self.dates)
|
|
infer_loader = DataLoader(CropDataset(X_infer, np.zeros(len(df)), augment=False), batch_size=128, shuffle=False)
|
|
|
|
fcn_probs = []
|
|
fcn_feats = []
|
|
with torch.no_grad():
|
|
for X_batch, _ in infer_loader:
|
|
out, feats = self.fcn(X_batch, return_features=True)
|
|
fcn_probs.extend(torch.softmax(out, dim=1).numpy())
|
|
fcn_feats.append(feats.numpy())
|
|
|
|
fcn_probs = np.array(fcn_probs)
|
|
fcn_feats = np.vstack(fcn_feats)
|
|
|
|
X_infer_flat = X_infer.reshape(X_infer.shape[0], -1)
|
|
X_stack = np.hstack([X_infer_flat, fcn_feats])
|
|
cb_probs = self.calibrated_cb.predict_proba(X_stack)
|
|
|
|
final_probs = (fcn_probs * self.w_fcn) + (cb_probs * self.w_cb)
|
|
final_preds = np.argmax(final_probs, axis=1)
|
|
|
|
# Identify No Data pixels: those with all NaNs or zeros after imputation
|
|
no_data_mask = np.zeros(len(df), dtype=bool)
|
|
for band in self.bands:
|
|
band_cols = [f"{date}_{band}" for date in self.dates if f"{date}_{band}" in df.columns]
|
|
if band_cols:
|
|
band_data = df[band_cols].values
|
|
# Check if pixel is all zeros or all NaN for this band
|
|
all_zeros = np.all(band_data == 0, axis=1)
|
|
all_nan = np.all(np.isnan(band_data), axis=1)
|
|
no_data_mask = no_data_mask | all_zeros | all_nan
|
|
|
|
# Override predictions for No Data pixels to class 0 (Background/No Data)
|
|
final_preds[no_data_mask] = 0
|
|
final_probs[no_data_mask] = 0.0
|
|
final_probs[no_data_mask, 0] = 1.0 # Set probability to 1.0 for class 0
|
|
|
|
if apply_spatial_smoothing and all(col in df.columns for col in coord_cols):
|
|
print(f"Applying spatial probability smoothing using {coord_cols}...")
|
|
coords = df[coord_cols].values
|
|
knn = KNeighborsRegressor(n_neighbors=9, weights='distance')
|
|
knn.fit(coords, final_probs)
|
|
smoothed_probs = knn.predict(coords)
|
|
final_preds = np.argmax(smoothed_probs, axis=1)
|
|
final_probs = smoothed_probs
|
|
|
|
# Re-apply No Data override after smoothing
|
|
final_preds[no_data_mask] = 0
|
|
final_probs[no_data_mask, 0] = 1.0
|
|
|
|
df['class_id'] = final_preds
|
|
df['predicted_crop'] = self.le.inverse_transform(final_preds)
|
|
df['confidence'] = np.max(final_probs, axis=1)
|
|
|
|
# Track missing data ratio for quality flag
|
|
missing_ratio = np.mean([m.mean(axis=1) for m in missing_mask.values()], axis=0)
|
|
df['high_missing'] = missing_ratio > 0.4
|
|
df['low_quality'] = (df['confidence'] < 0.5) | df['high_missing'] | no_data_mask
|
|
|
|
# Set NoData (0) for low quality pixels
|
|
df.loc[df['low_quality'], 'class_id'] = 0
|
|
df.loc[df['low_quality'], 'predicted_crop'] = 'Unknown/NoData'
|
|
return df
|
|
|
|
def export_to_geotiff(self, df, output_path="lulc_map.tif", crs="EPSG:6933"):
|
|
if xr is None or rioxarray is None:
|
|
raise ImportError("Missing required libraries: xarray, rioxarray")
|
|
print(f"Exporting LULC masks to {output_path}...")
|
|
|
|
ds_out = df.set_index(['lat', 'lon'])[['class_id', 'confidence', 'low_quality']].to_xarray()
|
|
ds_out = ds_out.rename({'lat': 'y', 'lon': 'x'})
|
|
ds_out = ds_out.sortby('y', ascending=False)
|
|
ds_out = ds_out.rio.set_spatial_dims(x_dim='x', y_dim='y')
|
|
ds_out.rio.write_crs(crs, inplace=True)
|
|
|
|
ds_out['class_id'].astype('uint16').rio.to_raster(output_path)
|
|
|
|
conf_path = output_path.replace('.tif', '_confidence.tif')
|
|
ds_out['confidence'].astype('float32').rio.to_raster(conf_path)
|
|
|
|
mask_path = output_path.replace('.tif', '_cloud_mask.tif')
|
|
ds_out['low_quality'].astype('uint8').rio.to_raster(mask_path)
|
|
|
|
legend_path = output_path.replace('.tif', '_legend.json')
|
|
legend_dict = {int(i): str(c) for i, c in enumerate(self.le.classes_)}
|
|
if 0 not in legend_dict:
|
|
legend_dict[0] = 'Unknown/NoData'
|
|
with open(legend_path, 'w') as f:
|
|
json.dump(legend_dict, f, indent=4)
|
|
print(f"✅ Successfully exported GeoTIFFs and class legend!")
|