hoàn thành server api dự đoán ra file tiff
This commit is contained in:
Victor Phan
+243
-8
@@ -34,6 +34,18 @@ training_status = {
|
||||
"error": None,
|
||||
"result": None,
|
||||
"start_time": None,
|
||||
"end_time": None,
|
||||
"cancel_requested": False
|
||||
}
|
||||
|
||||
# Global prediction status
|
||||
prediction_status = {
|
||||
"is_predicting": False,
|
||||
"progress": "",
|
||||
"error": None,
|
||||
"result": None,
|
||||
"output_file": None,
|
||||
"start_time": None,
|
||||
"end_time": None
|
||||
}
|
||||
|
||||
@@ -65,6 +77,27 @@ class TrainingConfig(BaseModel):
|
||||
training_shapefile: str = "train/ST_training data_updated_1130points_new.shp"
|
||||
|
||||
|
||||
class PredictionConfig(BaseModel):
|
||||
"""Cấu hình dự đoán"""
|
||||
# Model to use
|
||||
model_filename: str
|
||||
|
||||
# Khu vực (bbox)
|
||||
min_lon: float = 105.6
|
||||
min_lat: float = 9.3
|
||||
max_lon: float = 106.2
|
||||
max_lat: float = 9.8
|
||||
|
||||
# Thời gian
|
||||
start_date: str = "2023-03-01"
|
||||
end_date: str = "2023-05-31"
|
||||
|
||||
# Dữ liệu
|
||||
max_scenes: int = 12
|
||||
cloud_cover: int = 30
|
||||
resolution: int = 20
|
||||
|
||||
|
||||
class TrainingStatus(BaseModel):
|
||||
"""Trạng thái training"""
|
||||
is_training: bool
|
||||
@@ -161,10 +194,15 @@ async def start_training(config: TrainingConfig, background_tasks: BackgroundTas
|
||||
async def stop_training():
|
||||
"""Dừng training (nếu đang chạy)"""
|
||||
global training_status
|
||||
training_status["is_training"] = False
|
||||
training_status["error"] = "Đã dừng bởi người dùng"
|
||||
training_status["end_time"] = datetime.now().isoformat()
|
||||
return {"message": "Training đã dừng"}
|
||||
|
||||
if not training_status["is_training"]:
|
||||
return {"message": "Không có training nào đang chạy"}
|
||||
|
||||
# Set cancel flag - the training will check this and stop
|
||||
training_status["cancel_requested"] = True
|
||||
training_status["progress"] = "Đang hủy training..."
|
||||
|
||||
return {"message": "Đang dừng training..."}
|
||||
|
||||
|
||||
@app.get("/api/models/list")
|
||||
@@ -194,11 +232,43 @@ async def list_models():
|
||||
return {"models": models}
|
||||
|
||||
|
||||
@app.post("/api/prediction/start")
|
||||
async def start_prediction(config: PredictionConfig, background_tasks: BackgroundTasks):
|
||||
"""Bắt đầu dự đoán"""
|
||||
global prediction_status
|
||||
|
||||
if prediction_status["is_predicting"]:
|
||||
raise HTTPException(status_code=400, detail="Đang có dự đoán khác đang chạy")
|
||||
|
||||
# Reset status
|
||||
prediction_status = {
|
||||
"is_predicting": True,
|
||||
"progress": "Đang khởi động...",
|
||||
"error": None,
|
||||
"result": None,
|
||||
"output_file": None,
|
||||
"start_time": datetime.now().isoformat(),
|
||||
"end_time": None
|
||||
}
|
||||
|
||||
# Run prediction in background
|
||||
background_tasks.add_task(run_prediction, config)
|
||||
|
||||
return {"message": "Đã bắt đầu dự đoán", "status": prediction_status}
|
||||
|
||||
|
||||
@app.get("/api/prediction/status")
|
||||
async def get_prediction_status():
|
||||
"""Kiểm tra trạng thái dự đoán"""
|
||||
return prediction_status
|
||||
|
||||
|
||||
async def run_training(config: TrainingConfig):
|
||||
"""Chạy training process"""
|
||||
global training_status
|
||||
|
||||
try:
|
||||
training_status["cancel_requested"] = False
|
||||
training_status["progress"] = "Đang import thư viện..."
|
||||
|
||||
# Import training module
|
||||
@@ -206,6 +276,10 @@ async def run_training(config: TrainingConfig):
|
||||
|
||||
training_status["progress"] = "Đang load dữ liệu Sentinel-2..."
|
||||
|
||||
# Function to check if training should be cancelled
|
||||
def should_cancel():
|
||||
return training_status.get("cancel_requested", False)
|
||||
|
||||
# Run training
|
||||
result = train_model(
|
||||
bbox=[config.min_lon, config.min_lat, config.max_lon, config.max_lat],
|
||||
@@ -218,12 +292,18 @@ async def run_training(config: TrainingConfig):
|
||||
max_depth=config.max_depth,
|
||||
learning_rate=config.learning_rate,
|
||||
use_gpu=config.use_gpu,
|
||||
status_callback=lambda msg: update_progress(msg)
|
||||
status_callback=lambda msg: update_progress(msg),
|
||||
cancel_check=should_cancel
|
||||
)
|
||||
|
||||
training_status["is_training"] = False
|
||||
training_status["progress"] = "Hoàn thành!"
|
||||
training_status["result"] = result
|
||||
if training_status.get("cancel_requested", False):
|
||||
training_status["is_training"] = False
|
||||
training_status["progress"] = "Đã hủy training"
|
||||
training_status["error"] = "Training cancelled by user"
|
||||
else:
|
||||
training_status["is_training"] = False
|
||||
training_status["progress"] = "Hoàn thành!"
|
||||
training_status["result"] = result
|
||||
training_status["end_time"] = datetime.now().isoformat()
|
||||
|
||||
except Exception as e:
|
||||
@@ -242,6 +322,161 @@ def update_progress(message: str):
|
||||
print(f"[PROGRESS] {message}")
|
||||
|
||||
|
||||
def update_prediction_progress(message: str):
|
||||
"""Cập nhật prediction progress message"""
|
||||
global prediction_status
|
||||
prediction_status["progress"] = message
|
||||
print(f"[PREDICTION PROGRESS] {message}")
|
||||
|
||||
|
||||
async def run_prediction(config: PredictionConfig):
|
||||
"""Chạy prediction process"""
|
||||
global prediction_status
|
||||
|
||||
try:
|
||||
prediction_status["progress"] = "Đang import thư viện..."
|
||||
|
||||
# Import required libraries
|
||||
import xarray as xr
|
||||
import numpy as np
|
||||
from datetime import datetime as dt
|
||||
import rioxarray
|
||||
|
||||
prediction_status["progress"] = "Đang load model..."
|
||||
|
||||
# Load model
|
||||
model_path = Path("model_train") / config.model_filename
|
||||
if not model_path.exists():
|
||||
raise FileNotFoundError(f"Model không tồn tại: {config.model_filename}")
|
||||
|
||||
model_data = joblib.load(model_path)
|
||||
|
||||
# Extract model from dict (models are saved as {'model': xgb_model, 'label_encoder': encoder})
|
||||
if isinstance(model_data, dict):
|
||||
model = model_data.get('model')
|
||||
label_encoder = model_data.get('label_encoder')
|
||||
else:
|
||||
model = model_data
|
||||
label_encoder = None
|
||||
|
||||
prediction_status["progress"] = "Đang kết nối Microsoft Planetary Computer..."
|
||||
|
||||
# Import and use Microsoft Planetary Computer STAC API
|
||||
import pystac_client
|
||||
import planetary_computer
|
||||
|
||||
catalog = pystac_client.Client.open(
|
||||
"https://planetarycomputer.microsoft.com/api/stac/v1",
|
||||
modifier=planetary_computer.sign_inplace,
|
||||
)
|
||||
|
||||
bbox = [config.min_lon, config.min_lat, config.max_lon, config.max_lat]
|
||||
time_range = f"{config.start_date}/{config.end_date}"
|
||||
|
||||
prediction_status["progress"] = "Đang tải dữ liệu Sentinel-2..."
|
||||
|
||||
# Search Sentinel-2 data
|
||||
search = catalog.search(
|
||||
collections=["sentinel-2-l2a"],
|
||||
bbox=bbox,
|
||||
datetime=time_range,
|
||||
query={"eo:cloud_cover": {"lt": config.cloud_cover}}
|
||||
)
|
||||
|
||||
items = list(search.items()) # Changed from items_as_dicts() to items()
|
||||
if not items:
|
||||
raise ValueError("Không tìm thấy dữ liệu Sentinel-2 cho khu vực và thời gian này")
|
||||
|
||||
items = items[:config.max_scenes]
|
||||
|
||||
prediction_status["progress"] = f"Đang xử lý {len(items)} scenes Sentinel-2..."
|
||||
|
||||
# Load and process Sentinel-2 data (simplified)
|
||||
# Note: This is a simplified version. Full implementation would need more processing
|
||||
from odc.stac import load
|
||||
|
||||
s2_data = load(
|
||||
items,
|
||||
bbox=bbox,
|
||||
chunks={"time": 1, "x": 2048, "y": 2048},
|
||||
groupby="solar_day",
|
||||
resolution=config.resolution
|
||||
)
|
||||
|
||||
prediction_status["progress"] = "Đang tính toán các chỉ số..."
|
||||
|
||||
# Calculate NDVI using Sentinel-2 band names
|
||||
# B08 = NIR, B04 = Red
|
||||
nir = s2_data["B08"] # NIR band
|
||||
red = s2_data["B04"] # Red band
|
||||
ndvi = (nir - red) / (nir + red + 1e-8) # Add small value to avoid division by zero
|
||||
|
||||
# Resample to monthly
|
||||
ndvi_monthly = ndvi.resample(time="1M").mean()
|
||||
|
||||
prediction_status["progress"] = "Đang dự đoán..."
|
||||
|
||||
# Prepare features for prediction
|
||||
features_list = []
|
||||
for t in range(len(ndvi_monthly.time)):
|
||||
ndvi_t = ndvi_monthly.isel(time=t).values
|
||||
features_list.append(ndvi_t.flatten())
|
||||
|
||||
# Stack features
|
||||
features = np.column_stack(features_list)
|
||||
|
||||
# Make prediction
|
||||
predictions = model.predict(features)
|
||||
|
||||
# Reshape to original shape
|
||||
pred_shape = ndvi_monthly.isel(time=0).shape
|
||||
predictions_2d = predictions.reshape(pred_shape)
|
||||
|
||||
# Create output xarray
|
||||
prediction_da = xr.DataArray(
|
||||
predictions_2d,
|
||||
coords={
|
||||
"y": ndvi_monthly.y,
|
||||
"x": ndvi_monthly.x
|
||||
},
|
||||
dims=["y", "x"],
|
||||
name="classification"
|
||||
)
|
||||
|
||||
# Save output
|
||||
output_dir = Path("predictions")
|
||||
output_dir.mkdir(exist_ok=True)
|
||||
|
||||
timestamp = dt.now().strftime("%Y%m%d_%H%M%S")
|
||||
output_file = output_dir / f"prediction_{timestamp}.tif"
|
||||
|
||||
prediction_status["progress"] = "Đang lưu kết quả..."
|
||||
|
||||
# Save as GeoTIFF
|
||||
prediction_da.rio.write_crs(s2_data.rio.crs, inplace=True)
|
||||
prediction_da.rio.to_raster(output_file, driver="GTiff")
|
||||
|
||||
prediction_status["is_predicting"] = False
|
||||
prediction_status["progress"] = "Hoàn thành!"
|
||||
prediction_status["output_file"] = str(output_file)
|
||||
prediction_status["result"] = {
|
||||
"output_file": str(output_file),
|
||||
"shape": pred_shape,
|
||||
"unique_classes": np.unique(predictions).tolist(),
|
||||
"bbox": bbox,
|
||||
"time_range": time_range
|
||||
}
|
||||
prediction_status["end_time"] = dt.now().isoformat()
|
||||
|
||||
except Exception as e:
|
||||
prediction_status["is_predicting"] = False
|
||||
prediction_status["error"] = str(e)
|
||||
prediction_status["progress"] = f"Lỗi: {str(e)}"
|
||||
prediction_status["end_time"] = dt.now().isoformat()
|
||||
import traceback
|
||||
print(traceback.format_exc())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("=" * 70)
|
||||
print("🚀 LAND CLASSIFICATION TRAINING API SERVER")
|
||||
|
||||
Reference in New Issue
Block a user