terrakio-core 0.3.6__py3-none-any.whl → 0.3.8__py3-none-any.whl

terrakio_core/__init__.py

@@ -5,7 +5,7 @@ Terrakio Core
 Core components for Terrakio API clients.
 """
 
-__version__ = "0.3.4"
+__version__ = "0.3.8"
 
 from .async_client import AsyncClient
 from .sync_client import SyncClient

terrakio_core/async_client.py

@@ -46,65 +46,65 @@ class AsyncClient(BaseClient):
         else:
             return await self._make_request_with_retry(self._session, method, endpoint, **kwargs)
 
-
     async def _make_request_with_retry(self, session: aiohttp.ClientSession, method: str, endpoint: str, **kwargs) -> Dict[Any, Any]:
         url = f"{self.url}/{endpoint.lstrip('/')}"
+        last_exception = None
+        
         for attempt in range(self.retry + 1):
-            try:                
+            try:
                 async with session.request(method, url, **kwargs) as response:
-                    response_text = await response.text()
+                    if not response.ok and self._should_retry(response.status, attempt):
+                        self.logger.info(f"Request failed (attempt {attempt+1}/{self.retry+1}): {response.status}. Retrying...")
+                        continue
                     if not response.ok:
-                        should_retry = False
-                        
-                        if response.status == 400:
-                            should_retry = False
-                        else:
-                            if response.status in [408, 502, 503, 504]:
-                                should_retry = True
-                            elif response.status == 500:
-                                try:
-                                    response_text = await response.text()
-                                    if "Internal server error" not in response_text:
-                                        should_retry = True
-                                except:
-                                    should_retry = True
-                        
-                        if should_retry and attempt < self.retry:
-                            self.logger.info(f"Request failed (attempt {attempt+1}/{self.retry+1}): {response.status} {response.reason}. Retrying...")
-                            continue
-                        else:
-                            error_msg = f"API request failed: {response.status} {response.reason}"
-                            try:
-                                error_data = await response.json()
-                                if "detail" in error_data:
-                                    error_msg += f" - {error_data['detail']}"
-                            except:
-                                pass
-                            raise APIError(error_msg, status_code=response.status)
-                    
-                    content_type = response.headers.get('content-type', '').lower()
-                    content = await response.read()
-                    if 'json' in content_type:
-                        return json.loads(content.decode('utf-8'))
-                    elif 'csv' in content_type:
-                        return pd.read_csv(BytesIO(content))
-                    elif 'image/' in content_type:
-                        return content
-                    elif 'text' in content_type:
-                        return content.decode('utf-8')
-                    else:
+                        error_msg = f"API request failed: {response.status} {response.reason}"
                         try:
-                            return xr.open_dataset(BytesIO(content))
+                            error_data = await response.json()
+                            if "detail" in error_data:
+                                error_msg += f" - {error_data['detail']}"
                         except:
-                            raise APIError(f"Unknown response format. Content-Type: {response.headers.get('content-type', 'unknown')}", status_code=response.status)
+                            pass
+                        raise APIError(error_msg, status_code=response.status)
+                    return await self._parse_response(response)
+                    
             except aiohttp.ClientError as e:
+                last_exception = e
                 if attempt < self.retry:
-                    self.logger.info(f"Request failed (attempt {attempt+1}/{self.retry+1}): {e}. Retrying...")
+                    self.logger.info(f"Networking error (attempt {attempt+1}/{self.retry+1}): {e}. Retrying...")
                     continue
                 else:
-                    raise APIError(f"Request failed after {self.retry+1} attempts: {e}", status_code=None)
-                
+                    break
+        
+        raise APIError(f"Networking error, request failed after {self.retry+1} attempts: {last_exception}", status_code=None)
+    
+    def _should_retry(self, status_code: int, attempt: int) -> bool:
+        """Determine if the request should be retried based on status code."""
+        if attempt >= self.retry:
+            return False
+        elif status_code in [408, 502, 503, 504]:
+            return True
+        else:
+            return False
+
+    async def _parse_response(self, response) -> Any:
+        """Parse response based on content type."""
+        content_type = response.headers.get('content-type', '').lower()
+        content = await response.read()
 
+        if 'json' in content_type:
+            return json.loads(content.decode('utf-8'))
+        elif 'csv' in content_type:
+            return pd.read_csv(BytesIO(content))
+        elif 'image/' in content_type:
+            return content
+        elif 'text' in content_type:
+            return content.decode('utf-8')
+        else:
+            try:
+                return xr.open_dataset(BytesIO(content))
+            except:
+                raise APIError(f"Unknown response format: {content_type}", status_code=response.status)
+            
     async def _regular_request(self, method: str, endpoint: str, **kwargs):
         url = endpoint.lstrip('/')
         if self._session is None:
@@ -134,6 +134,7 @@ class AsyncClient(BaseClient):
         output: str = "csv",
         resolution: int = -1,
         geom_fix: bool = False,
+        validated: bool = True,
         **kwargs
     ):
         """
@@ -147,6 +148,7 @@ class AsyncClient(BaseClient):
             output (str): Output format ('csv' or 'netcdf')
             resolution (int): Resolution parameter
             geom_fix (bool): Whether to fix the geometry (default False)
+            validated (bool): Whether to use validated data (default True)
             **kwargs: Additional parameters to pass to the WCS request
             
         Returns:
@@ -169,6 +171,7 @@ class AsyncClient(BaseClient):
             "resolution": resolution,
             "expr": expr,
             "buffer": geom_fix,
+            "validated": validated,
             **kwargs
         }
         return await self._terrakio_request("POST", "geoquery", json=payload)

terrakio_core/endpoints/dataset_management.py

@@ -42,7 +42,7 @@ class DatasetManagement:
         return self._client._terrakio_request("GET", f"/datasets/{name}", params = params)
 
     @require_api_key
-    def create_dataset(
+    async def create_dataset(
         self, 
         name: str, 
         collection: str = "terrakio-datasets",
@@ -59,7 +59,8 @@ class DatasetManagement:
         proj4: Optional[str] = None,
         abstract: Optional[str] = None,
         geotransform: Optional[List[float]] = None,
-        padding: Optional[Any] = None
+        padding: Optional[Any] = None,
+        input: Optional[str] = None
     ) -> Dict[str, Any]:
         """
         Create a new dataset.
@@ -104,12 +105,13 @@ class DatasetManagement:
             "proj4": proj4,
             "abstract": abstract,
             "geotransform": geotransform,
-            "padding": padding
+            "padding": padding,
+            "input": input
         }
         for param, value in param_mapping.items():
             if value is not None:
                 payload[param] = value
-        return self._client._terrakio_request("POST", "/datasets", params = params, json = payload)
+        return await self._client._terrakio_request("POST", "/datasets", params = params, json = payload)
     
     @require_api_key
     def update_dataset(

terrakio_core/endpoints/mass_stats.py

@@ -12,7 +12,7 @@ class MassStats:
         self._client = client
 
     @require_api_key
-    async def upload_request(
+    async def _upload_request(
         self,
         name: str,
         size: int,
@@ -220,7 +220,7 @@ class MassStats:
         return self._client._terrakio_request("GET", "mass_stats/download", params=params)
     
     @require_api_key
-    async def upload_file(self, file_path: str, url: str, use_gzip: bool = False):
+    async def _upload_file(self, file_path: str, url: str, use_gzip: bool = False):
         """
         Helper method to upload a JSON file to a signed URL.
         
@@ -427,7 +427,7 @@ class MassStats:
             return e
         except json.JSONDecodeError as e:
             return e
-        upload_result = await self.upload_request(name = name, size = size, region = region, output = output, config = config, location = location, force_loc = force_loc, overwrite = overwrite, server = server, skip_existing = skip_existing)
+        upload_result = await self._upload_request(name = name, size = size, region = region, output = output, config = config, location = location, force_loc = force_loc, overwrite = overwrite, server = server, skip_existing = skip_existing)
         requests_url = upload_result.get('requests_url')
         manifest_url = upload_result.get('manifest_url')
         if not requests_url:
@@ -436,7 +436,7 @@ class MassStats:
         try:
             # in this place we are uploading the request json file, we need to check whether the json is in the correct format or not
             self.validate_request(request_json)
-            requests_response = await self.upload_file(request_json, requests_url, use_gzip=True)
+            requests_response = await self._upload_file(request_json, requests_url, use_gzip=True)
             if requests_response.status not in [200, 201, 204]:
                 self._client.logger.error(f"Requests upload error: {requests_response.text()}")
                 raise Exception(f"Failed to upload request JSON: {requests_response.text()}")
@@ -447,7 +447,7 @@ class MassStats:
             raise ValueError("No manifest_url returned from server for manifest JSON upload")
         
         try:
-            manifest_response = await self.upload_file(manifest_json, manifest_url, use_gzip=False)
+            manifest_response = await self._upload_file(manifest_json, manifest_url, use_gzip=False)
             if manifest_response.status not in [200, 201, 204]:
                 self._client.logger.error(f"Manifest upload error: {manifest_response.text()}")
                 raise Exception(f"Failed to upload manifest JSON: {manifest_response.text()}")

terrakio_core/endpoints/model_management.py

@@ -3,9 +3,32 @@ import json
 import time
 import textwrap
 import logging
-from typing import Dict, Any
+from typing import Dict, Any, Union, Tuple
+from io import BytesIO
+import numpy as np
 from google.cloud import storage
 from ..helper.decorators import require_token, require_api_key, require_auth
+TORCH_AVAILABLE = False
+SKL2ONNX_AVAILABLE = False
+
+try:
+    import torch
+    TORCH_AVAILABLE = True
+except ImportError:
+    torch = None
+
+try:
+    from skl2onnx import convert_sklearn
+    from skl2onnx.common.data_types import FloatTensorType
+    from sklearn.base import BaseEstimator
+    SKL2ONNX_AVAILABLE = True
+except ImportError:
+    convert_sklearn = None
+    FloatTensorType = None
+    BaseEstimator = None
+
+from io import BytesIO
+from typing import Tuple
 
 class ModelManagement:
     def __init__(self, client):
@@ -115,14 +138,13 @@ class ModelManagement:
             bar = '█' * filled_length + '░' * (bar_length - filled_length)
             percentage = progress * 100
             
-            # Print status with progress bar
-            print(f"\rJob status: {status} [{bar}] {percentage:.1f}% ({completed}/{total})", end='')
+            self._client.logger.info(f"Job status: {status} [{bar}] {percentage:.1f}% ({completed}/{total})")   
 
             if status == "Completed":
-                print("\nJob completed successfully!")
+                self._client.logger.info("Job completed successfully!")
                 break
             elif status == "Error":
-                print("\n")  # New line before error message
+                self._client.logger.info("Job encountered an error")
                 raise Exception(f"Job {task_id} encountered an error")
             
             # Wait 5 seconds before checking again
@@ -133,43 +155,238 @@ class ModelManagement:
         return task_id
     
     @require_api_key
-    async def upload_model(self, model_path: str):
+    async def upload_model(self, model, model_name: str, input_shape: Tuple[int, ...] = None):
         """
         Upload a model to the bucket so that it can be used for inference.
+        Converts PyTorch and scikit-learn models to ONNX format before uploading.
         
         Args:
-            model_path: Path to the model file
+            model: The model object (PyTorch model or scikit-learn model)
+            model_name: Name for the model (without extension)
+            input_shape: Shape of input data for ONNX conversion (e.g., (1, 10) for batch_size=1, features=10)
+                        Required for PyTorch models, optional for scikit-learn models
         
         Raises:
             APIError: If the API request fails
+            ValueError: If model type is not supported or input_shape is missing for PyTorch models
+            ImportError: If required libraries (torch or skl2onnx) are not installed
         """
         uid = (await self._client.auth.get_user_info())["uid"]
-        model_name = os.path.basename(model_path)
         
         client = storage.Client()
         bucket = client.get_bucket('terrakio-mass-requests')
-        model_file_name = os.path.splitext(model_name)[0]
-        blob = bucket.blob(f'{uid}/{model_file_name}/models/{model_name}')
         
-        blob.upload_from_filename(model_path)
-        self._client.logger.info(f"Model uploaded successfully to {uid}/{model_name}/models/{model_name}")
+        # Convert model to ONNX format
+        onnx_bytes = self._convert_model_to_onnx(model, model_name, input_shape)
+        
+        # Upload ONNX model to bucket
+        blob = bucket.blob(f'{uid}/{model_name}/models/{model_name}.onnx')
+        
+        blob.upload_from_string(onnx_bytes, content_type='application/octet-stream')
+        self._client.logger.info(f"Model uploaded successfully to {uid}/{model_name}/models/{model_name}.onnx")
+
+    def _convert_model_to_onnx(self, model, model_name: str, input_shape: Tuple[int, ...] = None) -> bytes:
+        """
+        Convert a model to ONNX format and return as bytes.
+        
+        Args:
+            model: The model object (PyTorch or scikit-learn)
+            model_name: Name of the model for logging
+            input_shape: Shape of input data
+            
+        Returns:
+            bytes: ONNX model as bytes
+            
+        Raises:
+            ValueError: If model type is not supported
+            ImportError: If required libraries are not installed
+        """
+        # Early check for any conversion capability
+        if not (TORCH_AVAILABLE or SKL2ONNX_AVAILABLE):
+            raise ImportError(
+                "ONNX conversion requires additional dependencies. Install with:\n"
+                "  pip install torch  # For PyTorch models\n"
+                "  pip install skl2onnx  # For scikit-learn models\n"
+                "  pip install torch skl2onnx  # For both"
+            )
+        
+        # Check if it's a PyTorch model using isinstance (preferred) with fallback
+        is_pytorch = False
+        if TORCH_AVAILABLE:
+            is_pytorch = (isinstance(model, torch.nn.Module) or 
+                         hasattr(model, 'state_dict'))
+        
+        # Check if it's a scikit-learn model
+        is_sklearn = False
+        if SKL2ONNX_AVAILABLE:
+            is_sklearn = (isinstance(model, BaseEstimator) or 
+                         (hasattr(model, 'fit') and hasattr(model, 'predict')))
+        
+        if is_pytorch and TORCH_AVAILABLE:
+            return self._convert_pytorch_to_onnx(model, model_name, input_shape)
+        elif is_sklearn and SKL2ONNX_AVAILABLE:
+            return self._convert_sklearn_to_onnx(model, model_name, input_shape)
+        else:
+            # Provide helpful error message
+            model_type = type(model).__name__
+            model_module = type(model).__module__
+            available_types = []
+            missing_deps = []
+            
+            if TORCH_AVAILABLE:
+                available_types.append("PyTorch (torch.nn.Module)")
+            else:
+                missing_deps.append("torch")
+                
+            if SKL2ONNX_AVAILABLE:
+                available_types.append("scikit-learn (BaseEstimator)")
+            else:
+                missing_deps.append("skl2onnx")
+            
+            if missing_deps:
+                raise ImportError(
+                    f"Model type {model_type} from {model_module} detected, but required dependencies missing: {', '.join(missing_deps)}. "
+                    f"Install with: pip install {' '.join(missing_deps)}"
+                )
+            else:
+                raise ValueError(
+                    f"Unsupported model type: {model_type} from {model_module}. "
+                    f"Supported types: {', '.join(available_types)}"
+                )
 
+    def _convert_pytorch_to_onnx(self, model, model_name: str, input_shape: Tuple[int, ...]) -> bytes:
+        """Convert PyTorch model to ONNX format with dynamic input dimensions."""
+        if input_shape is None:
+            raise ValueError("input_shape is required for PyTorch models")
+        
+        self._client.logger.info(f"Converting PyTorch model {model_name} to ONNX...")
+        
+        try:
+            # Set model to evaluation mode
+            model.eval()
+            
+            # Create dummy input
+            dummy_input = torch.randn(input_shape)
+            
+            # Use BytesIO to avoid creating temporary files
+            onnx_buffer = BytesIO()
+            
+            # Determine dynamic axes based on input shape
+            # Common patterns for different input types:
+            if len(input_shape) == 4:  # Convolutional input: (batch, channels, height, width)
+                dynamic_axes = {
+                    'float_input': {
+                        0: 'batch_size',
+                        2: 'height',    # Make height dynamic for variable input sizes
+                        3: 'width'      # Make width dynamic for variable input sizes
+                    },
+                    'output': {0: 'batch_size'}
+                }
+            elif len(input_shape) == 3:  # Could be (batch, sequence, features) or (batch, height, width)
+                dynamic_axes = {
+                    'float_input': {
+                        0: 'batch_size',
+                        1: 'dim1',      # Generic dynamic dimension
+                        2: 'dim2'       # Generic dynamic dimension
+                    },
+                    'output': {0: 'batch_size'}
+                }
+            elif len(input_shape) == 2:  # Likely (batch, features)
+                dynamic_axes = {
+                    'float_input': {
+                        0: 'batch_size'
+                        # Don't make features dynamic as it usually affects model architecture
+                    },
+                    'output': {0: 'batch_size'}
+                }
+            else:
+                # For other shapes, just make batch size dynamic
+                dynamic_axes = {
+                    'float_input': {0: 'batch_size'},
+                    'output': {0: 'batch_size'}
+                }
+            
+            torch.onnx.export(
+                model,
+                dummy_input,
+                onnx_buffer,
+                export_params=True,
+                opset_version=11,
+                do_constant_folding=True,
+                input_names=['float_input'],
+                output_names=['output'],
+                dynamic_axes=dynamic_axes
+            )
+            
+            self._client.logger.info(f"Successfully converted {model_name} with dynamic axes: {dynamic_axes}")
+            return onnx_buffer.getvalue()
+            
+        except Exception as e:
+            raise ValueError(f"Failed to convert PyTorch model {model_name} to ONNX: {str(e)}")
+
+
+    def _convert_sklearn_to_onnx(self, model, model_name: str, input_shape: Tuple[int, ...] = None) -> bytes:
+        """Convert scikit-learn model to ONNX format."""
+        self._client.logger.info(f"Converting scikit-learn model {model_name} to ONNX...")
+        
+        # Try to infer input shape if not provided
+        if input_shape is None:
+            if hasattr(model, 'n_features_in_'):
+                input_shape = (1, model.n_features_in_)
+            else:
+                raise ValueError(
+                    "input_shape is required for scikit-learn models when n_features_in_ is not available. "
+                    "This usually happens with older sklearn versions or models not fitted yet."
+                )
+        
+        try:
+            # Convert scikit-learn model to ONNX
+            initial_type = [('float_input', FloatTensorType(input_shape))]
+            onnx_model = convert_sklearn(model, initial_types=initial_type)
+            return onnx_model.SerializeToString()
+            
+        except Exception as e:
+            raise ValueError(f"Failed to convert scikit-learn model {model_name} to ONNX: {str(e)}")
+        
     @require_api_key
-    def upload_and_deploy_model(self, model_path: str, dataset: str, product: str, input_expression: str, dates_iso8601: list):
+    async def upload_and_deploy_cnn_model(self, model, model_name: str, dataset: str, product: str, input_expression: str, dates_iso8601: list, input_shape: Tuple[int, ...] = None):
+        """
+        Upload a CNN model to the bucket and deploy it.
         
+        Args:
+            model: The model object (PyTorch model or scikit-learn model)
+            model_name: Name for the model (without extension)
+            dataset: Name of the dataset to create
+            product: Product name for the inference
+            input_expression: Input expression for the dataset
+            dates_iso8601: List of dates in ISO8601 format
+            input_shape: Shape of input data for ONNX conversion (required for PyTorch models)
+
+        Raises:
+            APIError: If the API request fails
+            ValueError: If model type is not supported or input_shape is missing for PyTorch models
+            ImportError: If required libraries (torch or skl2onnx) are not installed
+        """
+        await self.upload_model(model=model, model_name=model_name, input_shape=input_shape)
+        # so the uploading process is kinda similar, but the deployment step is kinda different
+        await self.deploy_cnn_model(dataset=dataset, product=product, model_name=model_name, input_expression=input_expression, model_training_job_name=model_name, dates_iso8601=dates_iso8601)
+
+    @require_api_key
+    async def upload_and_deploy_model(self, model, model_name: str, dataset: str, product: str, input_expression: str, dates_iso8601: list, input_shape: Tuple[int, ...] = None):
         """
         Upload a model to the bucket and deploy it.
         
         Args:
-            model_path: Path to the model file
-            dataset Name of the dataset to create
+            model: The model object (PyTorch model or scikit-learn model)
+            model_name: Name for the model (without extension)
+            dataset: Name of the dataset to create
             product: Product name for the inference
             input_expression: Input expression for the dataset
             dates_iso8601: List of dates in ISO8601 format
+            input_shape: Shape of input data for ONNX conversion (required for PyTorch models)
         """
-        self.upload_model(model_path =  model_path)
-        model_name = os.path.basename(model_path)
-        self.deploy_model(dataset = dataset, product = product, model_name = model_name, input_expression = input_expression, model_training_job_name = model_name, dates_iso8601 = dates_iso8601)
+        await self.upload_model(model=model, model_name=model_name, input_shape=input_shape)
+        await self.deploy_model(dataset=dataset, product=product, model_name=model_name, input_expression=input_expression, model_training_job_name=model_name, dates_iso8601=dates_iso8601)
 
     @require_api_key
     def train_model(
@@ -212,7 +429,7 @@ class ModelManagement:
         return self._client._terrakio_request("POST", "/train_model", json=payload)
 
     @require_api_key
-    def deploy_model(
+    async def deploy_model(
         self, 
         dataset: str, 
         product: str, 
@@ -239,7 +456,7 @@ class ModelManagement:
             APIError: If the API request fails
         """
         # Get user info to get UID
-        user_info = self._client.get_user_info()
+        user_info = await self._client.auth.get_user_info()
         uid = user_info["uid"]
         
         # Generate and upload script
@@ -248,7 +465,53 @@ class ModelManagement:
         self._upload_script_to_bucket(script_content, script_name, model_training_job_name, uid)
         
         # Create dataset
-        return self._client.datasets.create_dataset(
+        return await self._client.datasets.create_dataset(
+            name=dataset,
+            collection="terrakio-datasets",
+            products=[product],
+            path=f"gs://terrakio-mass-requests/{uid}/{model_training_job_name}/inference_scripts",
+            input=input_expression,
+            dates_iso8601=dates_iso8601,
+            padding=0
+        )
+    
+    @require_api_key
+    async def deploy_cnn_model(
+        self, 
+        dataset: str, 
+        product: str, 
+        model_name: str, 
+        input_expression: str, 
+        model_training_job_name: str, 
+        dates_iso8601: list
+    ) -> Dict[str, Any]:
+        """
+        Deploy a CNN model by generating inference script and creating dataset.
+        
+        Args:
+            dataset: Name of the dataset to create
+            product: Product name for the inference
+            model_name: Name of the trained model
+            input_expression: Input expression for the dataset
+            model_training_job_name: Name of the training job
+            dates_iso8601: List of dates in ISO8601 format
+
+        Returns:
+            dict: Response from the deployment process
+
+        Raises:
+            APIError: If the API request fails
+        """
+        # Get user info to get UID
+        user_info = await self._client.auth.get_user_info()
+        uid = user_info["uid"]
+        
+        # Generate and upload script
+        script_content = self.generate_cnn_script(model_name, product, model_training_job_name, uid)
+        script_name = f"{product}.py"
+        self._upload_script_to_bucket(script_content, script_name, model_training_job_name, uid)
+        # Create dataset
+        return await self._client.datasets.create_dataset(
             name=dataset,
             collection="terrakio-datasets",
             products=[product],
@@ -374,6 +637,148 @@ class ModelManagement:
                 return result
             ''').strip()
     
+    @require_api_key
+    def generate_cnn_script(self, model_name: str, product: str, model_training_job_name: str, uid: str) -> str:
+        """
+        Generate Python inference script for CNN model with time-stacked bands.
+        
+        Args:
+            model_name: Name of the model
+            product: Product name
+            model_training_job_name: Training job name
+            uid: User ID
+            
+        Returns:
+            str: Generated Python script content
+        """
+        return textwrap.dedent(f'''
+            import logging
+            from io import BytesIO
+
+            import numpy as np
+            import pandas as pd
+            import xarray as xr
+            from google.cloud import storage
+            from onnxruntime import InferenceSession
+
+            logging.basicConfig(
+                level=logging.INFO
+            )
+
+            def get_model():
+                logging.info("Loading CNN model for {model_name}...")
+
+                client = storage.Client()
+                bucket = client.get_bucket('terrakio-mass-requests')
+                blob = bucket.blob('{uid}/{model_training_job_name}/models/{model_name}.onnx')
+
+                model = BytesIO()
+                blob.download_to_file(model)
+                model.seek(0)
+
+                session = InferenceSession(model.read(), providers=["CPUExecutionProvider"])
+                return session
+
+            def {product}(*bands, model):
+                logging.info("Start preparing CNN data with time-stacked bands")
+                
+                data_arrays = list(bands)
+                
+                if not data_arrays:
+                    raise ValueError("No bands provided")
+                    
+                reference_array = data_arrays[0]
+                original_shape = reference_array.shape
+                logging.info(f"Original shape: {{original_shape}}")
+                
+                # Get time coordinates - all bands should have the same time dimension
+                if 'time' not in reference_array.dims:
+                    raise ValueError("Time dimension is required for CNN processing")
+                    
+                time_coords = reference_array.coords['time']
+                num_timestamps = len(time_coords)
+                logging.info(f"Number of timestamps: {{num_timestamps}}")
+                
+                # Get spatial dimensions
+                spatial_dims = [dim for dim in reference_array.dims if dim != 'time']
+                height = reference_array.sizes[spatial_dims[0]]  # assuming first spatial dim is height
+                width = reference_array.sizes[spatial_dims[1]]   # assuming second spatial dim is width
+                logging.info(f"Spatial dimensions: {{height}} x {{width}}")
+                
+                # Stack bands across time dimension
+                # Result will be: (num_bands * num_timestamps, height, width)
+                stacked_channels = []
+                
+                for band_idx, data_array in enumerate(data_arrays):
+                    logging.info(f"Processing band {{band_idx + 1}}/{{len(data_arrays)}}")
+                    
+                    # Ensure consistent time coordinates across bands
+                    if not np.array_equal(data_array.coords['time'].values, time_coords.values):
+                        logging.warning(f"Band {{band_idx}} has different time coordinates, aligning...")
+                        data_array = data_array.sel(time=time_coords, method='nearest')
+                    
+                    # Extract values and ensure proper ordering (time, height, width)
+                    band_values = data_array.values
+                    if band_values.ndim == 3:
+                        # Reorder dimensions if needed to ensure (time, height, width)
+                        time_dim_idx = data_array.dims.index('time')
+                        if time_dim_idx != 0:
+                            axes_order = [time_dim_idx] + [i for i in range(len(data_array.dims)) if i != time_dim_idx]
+                            band_values = np.transpose(band_values, axes_order)
+                    
+                    # Add each timestamp of this band to the channel stack
+                    for t in range(num_timestamps):
+                        stacked_channels.append(band_values[t])
+                        
+                # Stack all channels: (num_bands * num_timestamps, height, width)
+                input_channels = np.stack(stacked_channels, axis=0)
+                total_channels = len(data_arrays) * num_timestamps
+                logging.info(f"Stacked channels shape: {{input_channels.shape}}")
+                logging.info(f"Total channels: {{total_channels}} ({{len(data_arrays)}} bands × {{num_timestamps}} timestamps)")
+                
+                # Add batch dimension: (1, num_channels, height, width)
+                input_data = np.expand_dims(input_channels, axis=0).astype(np.float32)
+                logging.info(f"Final input shape for CNN: {{input_data.shape}}")
+                
+                # Run inference
+                output = model.run(None, {{"float_input": input_data}})[0]
+                logging.info(f"Model output shape: {{output.shape}}")
+                
+                # Process output back to xarray format
+                # Assuming output is (1, height, width) or (1, 1, height, width)
+                if output.ndim == 4 and output.shape[1] == 1:
+                    # Remove channel dimension if it's 1
+                    output_2d = output[0, 0]
+                elif output.ndim == 3:
+                    # Remove batch dimension
+                    output_2d = output[0]
+                else:
+                    # Handle other cases
+                    output_2d = np.squeeze(output)
+                    if output_2d.ndim != 2:
+                        raise ValueError(f"Unexpected output shape after processing: {{output_2d.shape}}")
+                
+                # Determine output timestamp (use the latest timestamp)
+                output_timestamp = time_coords[-1]
+                
+                # Get spatial coordinates from reference array
+                spatial_coords = {{dim: reference_array.coords[dim] for dim in spatial_dims}}
+                
+                # Create output DataArray
+                result = xr.DataArray(
+                    data=np.expand_dims(output_2d.astype(np.float32), axis=0),
+                    dims=['time'] + spatial_dims,
+                    coords={{
+                        'time': [output_timestamp.values],
+                        spatial_dims[0]: spatial_coords[spatial_dims[0]].values,
+                        spatial_dims[1]: spatial_coords[spatial_dims[1]].values
+                    }}
+                )
+                
+                logging.info(f"Final result shape: {{result.shape}}")
+                return result
+            ''').strip()
+    
     @require_api_key
     def _upload_script_to_bucket(self, script_content: str, script_name: str, model_training_job_name: str, uid: str):
         """Upload the generated script to Google Cloud Storage"""

terrakio_core/exceptions.py

@@ -1,7 +1,9 @@
 class APIError(Exception):
     """Exception raised for errors in the API responses."""
-    pass
-
+    
+    def __init__(self, message, status_code=None):
+        super().__init__(message)
+        self.status_code = status_code
 
 class ConfigurationError(Exception):
     """Exception raised for errors in the configuration."""

{terrakio_core-0.3.6.dist-info → terrakio_core-0.3.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: terrakio-core
-Version: 0.3.6
+Version: 0.3.8
 Summary: Core components for Terrakio API clients
 Author-email: Yupeng Chao <yupeng@haizea.com.au>
 Project-URL: Homepage, https://github.com/HaizeaAnalytics/terrakio-python-api
@@ -23,6 +23,12 @@ Requires-Dist: shapely>=2.0.0
 Requires-Dist: geopandas>=0.13.0
 Requires-Dist: google-cloud-storage>=2.0.0
 Requires-Dist: nest_asyncio
+Provides-Extra: ml
+Requires-Dist: torch>=2.7.1; extra == "ml"
+Requires-Dist: scikit-learn>=1.7.0; extra == "ml"
+Requires-Dist: skl2onnx>=1.19.1; extra == "ml"
+Requires-Dist: onnx>=1.18.0; extra == "ml"
+Requires-Dist: onnxruntime>=1.10.0; extra == "ml"
 
 # Terrakio Core
 

{terrakio_core-0.3.6.dist-info → terrakio_core-0.3.8.dist-info}/RECORD

@@ -1,21 +1,21 @@
-terrakio_core/__init__.py,sha256=3MTDxMAQZCgXE5_gREyxTGTNK4WDZHlW91Syr0axr1M,242
-terrakio_core/async_client.py,sha256=0Zz-g5X4B2NZHRJrTuJAwNBUfO596Zvn1Nga3J-PiaE,13092
+terrakio_core/__init__.py,sha256=nUk_Q29ij_1R32AjR8Ygwy0_ID4-zdVkQexvBz7reM4,242
+terrakio_core/async_client.py,sha256=wyuIJGMAzfFwfM5BYgCTMuY0YAKNYJqRNB3IF5NpVow,12730
 terrakio_core/client.py,sha256=h8GW88g6RlGwNFW6MW48c_3BnaeT9nSd19LI1jCn1GU,1008
 terrakio_core/config.py,sha256=r8NARVYOca4AuM88VP_j-8wQxOk1s7VcRdyEdseBlLE,4193
-terrakio_core/exceptions.py,sha256=9S-I20-QiDRj1qgjFyYUwYM7BLic_bxurcDOIm2Fu_0,410
+terrakio_core/exceptions.py,sha256=4qnpOM1gOxsNIXDXY4qwY1d3I4Myhp7HBh7b2D0SVrU,529
 terrakio_core/sync_client.py,sha256=v1mcBtUaKWACqZgw8dTTVPMxUfKfiY0kjtBKzDwtGTU,13634
 terrakio_core/convenience_functions/convenience_functions.py,sha256=U7bLGwfBF-FUYc0nv49pAViPsBQ6LgPlV6c6b-zeKo8,10616
 terrakio_core/endpoints/auth.py,sha256=e_hdNE6JOGhRVlQMFdEoOmoMHp5EzK6CclOEnc_AmZw,5863
-terrakio_core/endpoints/dataset_management.py,sha256=8uf6cxlSSevqnQWcldtA9Cd24D5VrmWyxkE7Ngx3IEw,13084
+terrakio_core/endpoints/dataset_management.py,sha256=BUm8IIlW_Q45vDiQp16CiJGeSLheI8uWRVRQtMdhaNk,13161
 terrakio_core/endpoints/group_management.py,sha256=VFl3jakjQa9OPi351D3DZvLU9M7fHdfjCzGhmyJsx3U,6309
-terrakio_core/endpoints/mass_stats.py,sha256=KDmIlMYy4nkehPU5Ejtb_WN9Cz5mkt_rIsyDZkTWOLA,21351
-terrakio_core/endpoints/model_management.py,sha256=1ZYymaTQ7IY191sLSS7MWvhrHLmy2VeAM2A1Ty5NhU0,15346
+terrakio_core/endpoints/mass_stats.py,sha256=y1w3QLkDD0sKP1tBcFDqgLYLNxX94I-LYbNotaKhLYM,21356
+terrakio_core/endpoints/model_management.py,sha256=Q2bqsVfBILu-hZVw1tr5WjOR68qoYF6m326YJXgAOeo,33886
 terrakio_core/endpoints/space_management.py,sha256=YWb55nkJnFJGlALJ520DvurxDqVqwYtsvqQPWzxzhDs,2266
 terrakio_core/endpoints/user_management.py,sha256=x0JW6VET7eokngmkhZPukegxoJNR1X09BVehJt2nIdI,3781
 terrakio_core/helper/bounded_taskgroup.py,sha256=wiTH10jhKZgrsgrFUNG6gig8bFkUEPHkGRT2XY7Rgmo,677
 terrakio_core/helper/decorators.py,sha256=L6om7wmWNgCei3Wy5U0aZ-70OzsCwclkjIf7SfQuhCg,2289
 terrakio_core/helper/tiles.py,sha256=xNtp3oDD912PN_FQV5fb6uQYhwfHANuXyIcxoVCCfZU,2632
-terrakio_core-0.3.6.dist-info/METADATA,sha256=b0a6IvGiQAjaN-iBUu7gmE4-oxaZLIhfw8KQ_xV8tOs,1476
-terrakio_core-0.3.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-terrakio_core-0.3.6.dist-info/top_level.txt,sha256=5cBj6O7rNWyn97ND4YuvvXm0Crv4RxttT4JZvNdOG6Q,14
-terrakio_core-0.3.6.dist-info/RECORD,,
+terrakio_core-0.3.8.dist-info/METADATA,sha256=oZlZhEda5qq8myogGbxvlV0ZJExcEf5kMaYWEVES0BE,1728
+terrakio_core-0.3.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+terrakio_core-0.3.8.dist-info/top_level.txt,sha256=5cBj6O7rNWyn97ND4YuvvXm0Crv4RxttT4JZvNdOG6Q,14
+terrakio_core-0.3.8.dist-info/RECORD,,