hoàn thành server api dự đoán ra file tiff

api_server.py

@@ -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")