Merge pull request #6 from GeoCompas/gpu-info

Update gpu info endpoint

.env.example

@@ -1 +1 @@
-BASE_URL=https://samgeo-api.geocompas.ai/
+BASE_URL=https://samgeo-api.geocompas.ai

app/routes/predictions.py

@@ -4,7 +4,7 @@
 from fastapi import APIRouter, HTTPException
 from fastapi.responses import JSONResponse
 from fastapi import UploadFile, File, Depends
-from utils.utils import date_minute_str
+from utils.utils import get_timestamp
 from schemas.geojson import JSONDataBase
 
 router = APIRouter()
@@ -43,7 +43,7 @@ def list_files_in_project(project_id: str = ""):
                     "tif_url": None,
                 }
             detections[base_name]["geojson_files"] = [
-                f"{BASE_URL}/files/{project_id}/{f.name}"
+                f"{BASE_URL}files/{project_id}/{f.name}"
                 for f in public_dir.iterdir()
                 if f.suffix == ".geojson" and base_name in f.stem
             ]
@@ -75,7 +75,7 @@ async def upload_geojson(request: JSONDataBase):
     if not public_dir.exists():
         public_dir.mkdir(parents=True)
 
-    geojson_file_name = f"{request.id}_{date_minute_str()}_fixed.geojson"
+    geojson_file_name = f"{request.id}_{get_timestamp()}_fixed.geojson"
     geojson_file_path = public_dir / geojson_file_name
 
     try:

app/schemas/segment.py

@@ -1,5 +1,5 @@
 from pydantic import BaseModel, Field, field_validator
-from typing import List, Tuple, Optional, Any, Dict
+from typing import List, Tuple, Optional, Any, Dict, Literal
 
 
 class SegmentRequestBase(BaseModel):
@@ -25,6 +25,9 @@ class SegmentRequestBase(BaseModel):
         None,
         description="Type of action requested (e.g., single_point or multi_point segmentation), default is None",
     )
+    return_format: Literal["geojson", "url"] = Field(
+        "geojson", description="Return format, must be 'geojson' or 'url'. Default is 'geojson'."
+    )
 
     @field_validator("bbox", mode="before")
     def validate_bbox(cls, bbox):

app/utils/sam2.py

@@ -4,7 +4,7 @@
 from utils.utils import (
     generate_geojson,
     save_geojson,
-    date_minute_str,
+    get_timestamp,
 )
 from schemas.segment import SegmentRequestBase, SegmentResponseBase
 from utils.logger_config import log
@@ -39,7 +39,7 @@ def detect_automatic_sam2(bbox, zoom, id, project) -> SegmentResponseBase:
     """
     Detect objects automatically using SAM2 model based on the provided bounding box.
     """
-    date_time = date_minute_str()
+    date_time = get_timestamp()
     tif_file_name = f"{id}.tif"
     mask_file_name = f"tmp/mask_{id}.tif"
     geojson_file_name = f"{id}_{date_time}.geojson"
@@ -74,12 +74,13 @@ def detect_predictor_sam2(request: SegmentRequestBase) -> SegmentResponseBase:
     id = request.id
     project = request.project
     action_type = request.action_type
+    return_format = return_format.return_format
 
     log.info(
         f"Processing segmentation for bbox: {bbox}, zoom: {zoom}, id: {id}, project: {project}, action_type: {action_type}"
     )
 
-    date_time = date_minute_str()
+    date_time = get_timestamp()
     tif_file_name = f"{id}.tif"
     mask_file_name = f"tmp/mask_{id}.tif"
     geojson_file_name = f"{id}_{date_time}.geojson"

app/utils/utils.py

@@ -6,7 +6,7 @@
 import json
 import psutil
 from datetime import datetime
-from samgeo import tms_to_geotiff
+from samgeo import tms_to_geotiff, choose_device
 from shapely.geometry import Polygon, MultiPolygon
 from utils.logger_config import log
 
@@ -55,17 +55,45 @@ def group_files_by_base_name(public_dir: str, base_url: str) -> list:
     return sorted_grouped_files
 
 
+def format_memory(size_in_bytes):
+    """Formats memory size: returns in GB if size > 1024 MB, otherwise in MB."""
+    size_in_mb = size_in_bytes / (1024**2)
+    if size_in_mb > 1024:
+        size_in_gb = round(size_in_bytes / (1024**3), 2)
+        return f"{size_in_gb} GB"
+    else:
+        size_in_mb = round(size_in_mb, 2)
+        return f"{size_in_mb} MB"
+
+
 def check_gpu():
     """
-    Checks if a GPU is available on the system and returns CPU and memory information.
+    Checks if a GPU is available on the system and returns detailed GPU, CPU, and memory information.
 
     Returns:
         dict: A dictionary containing GPU, CPU, and memory information.
     """
+    device = choose_device()
+    gpu_info = {}
+
     if torch.cuda.is_available():
-        gpu_info = {"gpu": True, "device": torch.cuda.get_device_name(0)}
+        num_gpus = torch.cuda.device_count()
+        gpu_memory_total = format_memory(torch.cuda.get_device_properties(0).total_memory)
+        gpu_memory_allocated = format_memory(torch.cuda.memory_allocated(0))
+        gpu_memory_cached = format_memory(torch.cuda.memory_reserved(0))
+
+        gpu_info = {
+            "device_name": torch.cuda.get_device_name(0),
+            "num_gpus": num_gpus,
+            "gpu_memory_total": gpu_memory_total,
+            "gpu_memory_allocated": gpu_memory_allocated,
+            "gpu_memory_cached": gpu_memory_cached,
+            "gpu_memory_free": format_memory(
+                torch.cuda.get_device_properties(0).total_memory - torch.cuda.memory_allocated(0)
+            ),
+        }
     else:
-        gpu_info = {"gpu": False, "message": "No GPU available, using CPU"}
+        gpu_info = {}
 
     cpu_info = {
         "cpu_percent": psutil.cpu_percent(interval=1),
@@ -75,13 +103,13 @@ def check_gpu():
 
     memory_info = psutil.virtual_memory()
     memory_usage = {
-        "total_memory": memory_info.total / (1024**3),
-        "used_memory": memory_info.used / (1024**3),
-        "free_memory": memory_info.available / (1024**3),
+        "total_memory": format_memory(memory_info.total),
+        "used_memory": format_memory(memory_info.used),
+        "free_memory": format_memory(memory_info.available),
         "memory_percent": memory_info.percent,
     }
 
-    return {"gpu": gpu_info, "cpu": cpu_info, "memory": memory_usage}
+    return {"device": str(device), "gpu": gpu_info, "cpu": cpu_info, "memory": memory_usage}
 
 
 def save_geojson(json_data, output_geojson_path):
@@ -173,3 +201,13 @@ def date_minute_str():
         str: The current date and time formatted as YYYYMMDD_HHMM.
     """
     return datetime.now().strftime("%Y%m%d_%H%M")
+
+
+def get_timestamp():
+    """
+    Returns the current Unix timestamp as an integer (no decimals).
+
+    Returns:
+        int: The current Unix timestamp.
+    """
+    return int(datetime.now().timestamp())