File size: 3,994 Bytes
80e070e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
import os
import glob
import yaml
import ee
import geemap
import geopandas as gpd
import rasterio
from shapely.geometry import box

def load_config():
    config_path = os.path.join(os.path.dirname(__file__), 'config.yaml')
    with open(config_path, 'r') as f:
        return yaml.safe_load(f)

def get_s2_image(roi, date_start, date_end):
    collection = ee.ImageCollection('COPERNICUS/S2_SR_HARMONIZED') \
        .filterBounds(roi) \
        .filterDate(date_start, date_end)
    
    def mask_s2_clouds(image):
        qa = image.select('QA60')
        cloudBitMask = 1 << 10
        cirrusBitMask = 1 << 11
        mask = qa.bitwiseAnd(cloudBitMask).eq(0).And(qa.bitwiseAnd(cirrusBitMask).eq(0))
        return image.updateMask(mask).divide(10000)
    
    s2_median = collection.map(mask_s2_clouds).median().clip(roi)
    
    ndvi = s2_median.normalizedDifference(['B8', 'B4']).rename('NDVI')
    ndwi = s2_median.normalizedDifference(['B3', 'B8']).rename('NDWI')
    ndsi = s2_median.normalizedDifference(['B3', 'B11']).rename('NDSI')
    
    return s2_median.select(['B2', 'B3', 'B4', 'B8', 'B11', 'B12']).addBands([ndvi, ndwi, ndsi])

def get_terrain_image(roi):
    def add_terrain(img):
        slope = ee.Terrain.slope(img).rename('slope')
        aspect = ee.Terrain.aspect(img).rename('aspect')
        return img.addBands([slope, aspect])
        
    terrain = ee.ImageCollection('COPERNICUS/DEM/GLO30').select('DEM').map(add_terrain).mosaic().clip(roi)
    elevation = terrain.select('DEM').rename('elevation')
    slope = terrain.select('slope')
    aspect = terrain.select('aspect')
    return elevation.addBands([slope, aspect])

def main():
    config = load_config()
    project_id = config['gee']['project_id']
    
    try:
        ee.Initialize(project=project_id)
    except Exception:
        print("\n--- Google Earth Engine Authentication Required ---")
        ee.Authenticate()
        ee.Initialize(project=project_id)
        print("Authentication successful!\n")

    mask_dir = os.path.join(os.path.dirname(__file__), config['paths']['output_dir'], 'masks')
    output_dir = os.path.join(os.path.dirname(__file__), config['paths']['output_dir'], 'features')
    os.makedirs(output_dir, exist_ok=True)
    
    target_crs = config['preprocessing']['target_crs']
    date_start = config['gee']['date_start']
    date_end = config['gee']['date_end']

    mask_files = sorted(glob.glob(os.path.join(mask_dir, 'mask_*.tif')))
    print(f"Found {len(mask_files)} masks. Extracting features for these locations...")
    
    for filepath in mask_files:
        filename = os.path.basename(filepath)
        idx_str = filename.split('_')[1].split('.')[0]
        idx = int(idx_str)
        
        out_tif = os.path.join(output_dir, f"feature_{idx:04d}.tif")
        if os.path.exists(out_tif):
            print(f"Skipping {idx}, feature stack already exists.")
            continue
            
        with rasterio.open(filepath) as src:
            bounds = src.bounds
            
        minx, miny, maxx, maxy = bounds
        padded_box_utm = gpd.GeoDataFrame(geometry=[box(minx, miny, maxx, maxy)], crs=target_crs)
        padded_box_4326 = padded_box_utm.to_crs('EPSG:4326')
        bounds_4326 = padded_box_4326.total_bounds
        
        roi = ee.Geometry.BBox(bounds_4326[0], bounds_4326[1], bounds_4326[2], bounds_4326[3])
        
        try:
            s2_img = get_s2_image(roi, date_start, date_end)
            terrain_img = get_terrain_image(roi)
            final_img = s2_img.addBands(terrain_img)
            
            print(f"Downloading feature stack {idx:04d}...")
            geemap.ee_export_image(
                final_img, 
                filename=out_tif, 
                scale=10, 
                region=roi, 
                crs=target_crs,
                file_per_band=False
            )
        except Exception as e:
            print(f"Failed to process {idx:04d}: {e}")

if __name__ == '__main__':
    main()