Merge pull request #310 from hotosm/feat/image-processing-failed

feat: add image processing failure handling and retry mechanism for drone task image uploads & re-run process

src/backend/app/models/enums.py

@@ -168,6 +168,7 @@ class EventType(str, Enum):
     - ``assign`` -- For a requester user to assign a task to another user. Set the state *locked for mapping* passing in the required user id.
     - ``comment`` -- Keep the state the same, but simply add a comment.
     - ``unlock`` -- Unlock a task state by unlocking it if it's locked.
+    - ``image_upload`` -- Set the state to *image uploaded* when the task image is uploaded.
 
     Note that ``task_id`` must be specified in the endpoint too.
     """
@@ -183,3 +184,4 @@ class EventType(str, Enum):
     ASSIGN = "assign"
     COMMENT = "comment"
     UNLOCK = "unlock"
+    IMAGE_UPLOAD = "image_upload"

src/backend/app/projects/image_processing.py

@@ -6,14 +6,14 @@
 from app.tasks import task_logic
 from app.models.enums import State
 from app.utils import timestamp
+from app.db import database
 from pyodm import Node
 from app.s3 import get_file_from_bucket, list_objects_from_bucket, add_file_to_bucket
 from loguru import logger as log
 from concurrent.futures import ThreadPoolExecutor
 from psycopg import Connection
 from asgiref.sync import async_to_sync
 from app.config import settings
-from app.db import database
 
 
 class DroneImageProcessor:
@@ -199,6 +199,8 @@ async def download_and_upload_assets_from_odm_to_s3(
     dtm_project_id: uuid.UUID,
     dtm_task_id: uuid.UUID,
     user_id: str,
+    current_state: State,
+    comment: str,
 ):
     """
     Downloads results from ODM and uploads them to S3 (Minio).
@@ -210,6 +212,7 @@ async def download_and_upload_assets_from_odm_to_s3(
 
     """
     log.info(f"Starting download for task {task_id}")
+
     # Replace with actual ODM node details and URL
     node = Node.from_url(node_odm_url)
 
@@ -241,8 +244,8 @@ async def download_and_upload_assets_from_odm_to_s3(
                 project_id=dtm_project_id,
                 task_id=dtm_task_id,
                 user_id=user_id,
-                comment="Task completed.",
-                initial_state=State.IMAGE_UPLOADED,
+                comment=comment,
+                initial_state=current_state,
                 final_state=State.IMAGE_PROCESSED,
                 updated_at=timestamp(),
             )

src/backend/app/projects/project_routes.py

@@ -34,6 +34,7 @@
 from app.tasks import task_schemas
 from app.utils import geojson_to_kml, timestamp
 from app.users import user_schemas
+from minio.deleteobjects import DeleteObject
 
 
 router = APIRouter(
@@ -287,8 +288,9 @@ async def preview_split_by_square(
 
 @router.post("/generate-presigned-url/", tags=["Image Upload"])
 async def generate_presigned_url(
-    data: project_schemas.PresignedUrlRequest,
     user: Annotated[AuthUser, Depends(login_required)],
+    data: project_schemas.PresignedUrlRequest,
+    replace_existing: bool = False,
 ):
     """
     Generate a pre-signed URL for uploading an image to S3 Bucket.
@@ -297,21 +299,54 @@ async def generate_presigned_url(
     an S3 bucket. The URL expires after a specified duration.
 
     Args:
-
-        image_name: The name of the image you want to upload
-        expiry : Expiry time in hours
+        image_name: The name of the image(s) you want to upload.
+        expiry : Expiry time in hours.
+        replace_existing: A boolean flag to indicate if the image should be replaced.
 
     Returns:
-
-        str: The pre-signed URL to upload the image
+        list: A list of dictionaries with the image name and the pre-signed URL to upload.
     """
     try:
-        # Generate a pre-signed URL for an object
+        # Initialize the S3 client
         client = s3_client()
         urls = []
+
+        # Process each image in the request
         for image in data.image_name:
             image_path = f"projects/{data.project_id}/{data.task_id}/images/{image}"
 
+            # If replace_existing is True, delete the image first
+            if replace_existing:
+                image_dir = f"projects/{data.project_id}/{data.task_id}/images/"
+                try:
+                    # Prepare the list of objects to delete (recursively if necessary)
+                    delete_object_list = map(
+                        lambda x: DeleteObject(x.object_name),
+                        client.list_objects(
+                            settings.S3_BUCKET_NAME, image_dir, recursive=True
+                        ),
+                    )
+
+                    # Remove the objects (images)
+                    errors = client.remove_objects(
+                        settings.S3_BUCKET_NAME, delete_object_list
+                    )
+
+                    # Handle deletion errors, if any
+                    for error in errors:
+                        log.error("Error occurred when deleting object", error)
+                        raise HTTPException(
+                            status_code=HTTPStatus.BAD_REQUEST,
+                            detail=f"Failed to delete existing image: {error}",
+                        )
+
+                except Exception as e:
+                    raise HTTPException(
+                        status_code=HTTPStatus.BAD_REQUEST,
+                        detail=f"Failed to delete existing image. {e}",
+                    )
+
+            # Generate a new pre-signed URL for the image upload
             url = client.get_presigned_url(
                 "PUT",
                 settings.S3_BUCKET_NAME,
@@ -321,6 +356,7 @@ async def generate_presigned_url(
             urls.append({"image_name": image, "url": url})
 
         return urls
+
     except Exception as e:
         raise HTTPException(
             status_code=HTTPStatus.BAD_REQUEST,
@@ -395,17 +431,6 @@ async def process_imagery(
     db: Annotated[Connection, Depends(database.get_db)],
 ):
     user_id = user_data.id
-    # TODO: Update task state to reflect completion of image uploads.
-    await task_logic.update_task_state(
-        db,
-        project.id,
-        task_id,
-        user_id,
-        "Task images upload completed.",
-        State.LOCKED_FOR_MAPPING,
-        State.IMAGE_UPLOADED,
-        timestamp(),
-    )
     background_tasks.add_task(
         project_logic.process_drone_images, project.id, task_id, user_id, db
     )
@@ -443,7 +468,10 @@ async def get_assets_info(
 
         return results
     else:
-        return project_logic.get_project_info_from_s3(project.id, task_id)
+        current_state = await task_logic.get_task_state(db, project.id, task_id)
+        project_info = project_logic.get_project_info_from_s3(project.id, task_id)
+        project_info.state = current_state.get("state")
+        return project_info
 
 
 @router.post(
@@ -479,16 +507,74 @@ async def odm_webhook(
     # If status is 'success', download and upload assets to S3.
     # 40 is the status code for success in odm
     if status["code"] == 40:
-        # Call function to download assets from ODM and upload to S3
-        background_tasks.add_task(
-            image_processing.download_and_upload_assets_from_odm_to_s3,
-            settings.NODE_ODM_URL,
-            task_id,
-            dtm_project_id,
-            dtm_task_id,
-            dtm_user_id,
-        )
+        log.info(f"Task ID: {task_id}, Status: going for download......")
+
+        current_state = await task_logic.get_task_state(db, dtm_project_id, dtm_task_id)
+        current_state_value = State[current_state.get("state")]
+        match current_state_value:
+            case State.IMAGE_UPLOADED:
+                log.info(
+                    f"Task ID: {task_id}, Status: already IMAGE_UPLOADED - no update needed."
+                )
+                # Call function to download assets from ODM and upload to S3
+                background_tasks.add_task(
+                    image_processing.download_and_upload_assets_from_odm_to_s3,
+                    settings.NODE_ODM_URL,
+                    task_id,
+                    dtm_project_id,
+                    dtm_task_id,
+                    dtm_user_id,
+                    State.IMAGE_UPLOADED,
+                    "Task completed.",
+                )
+
+            case State.IMAGE_PROCESSING_FAILED:
+                log.warning(
+                    f"Task ID: {task_id}, Status: previously failed, updating to IMAGE_UPLOADED"
+                )
+                # Call function to download assets from ODM and upload to S3
+                background_tasks.add_task(
+                    image_processing.download_and_upload_assets_from_odm_to_s3,
+                    settings.NODE_ODM_URL,
+                    task_id,
+                    dtm_project_id,
+                    dtm_task_id,
+                    dtm_user_id,
+                    State.IMAGE_UPLOADED,
+                    "Task completed.",
+                )
+
+            case _:
+                log.info(
+                    f"Task ID: {task_id}, Status: updating to IMAGE_UPLOADED from {current_state}"
+                )
+
     elif status["code"] == 30:
-        # failed task
-        log.error(f'ODM task {task_id} failed: {status["errorMessage"]}')
+        current_state = await task_logic.get_task_state(db, dtm_project_id, dtm_task_id)
+        # If the current state is not already IMAGE_PROCESSING_FAILED, update it
+        if current_state != State.IMAGE_PROCESSING_FAILED:
+            await task_logic.update_task_state(
+                db,
+                dtm_project_id,
+                dtm_task_id,
+                dtm_user_id,
+                "Image processing failed.",
+                State.IMAGE_UPLOADED,
+                State.IMAGE_PROCESSING_FAILED,
+                timestamp(),
+            )
+
+            background_tasks.add_task(
+                image_processing.download_and_upload_assets_from_odm_to_s3,
+                settings.NODE_ODM_URL,
+                task_id,
+                dtm_project_id,
+                dtm_task_id,
+                dtm_user_id,
+                State.IMAGE_PROCESSING_FAILED,
+                "Image processing failed.",
+            )
+
+    log.info(f"Task ID: {task_id}, Status: Webhook received")
+
     return {"message": "Webhook received", "task_id": task_id}

src/backend/app/projects/project_schemas.py

@@ -12,7 +12,7 @@
 from psycopg import Connection
 from psycopg.rows import class_row
 from slugify import slugify
-from app.models.enums import FinalOutput, ProjectVisibility
+from app.models.enums import FinalOutput, ProjectVisibility, UserRole
 from app.models.enums import (
     IntEnum,
     ProjectStatus,
@@ -58,6 +58,7 @@ class AssetsInfo(BaseModel):
     task_id: str
     image_count: int
     assets_url: Optional[str]
+    state: Optional[UserRole] = None
 
 
 def validate_geojson(
@@ -197,6 +198,7 @@ class DbProject(BaseModel):
     altitude_from_ground: Optional[float] = None
     is_terrain_follow: bool = False
     image_url: Optional[str] = None
+    created_at: datetime
 
     async def one(db: Connection, project_id: uuid.UUID):
         """Get a single project &  all associated tasks by ID."""
@@ -258,15 +260,7 @@ async def one(db: Connection, project_id: uuid.UUID):
                     SELECT DISTINCT ON (te.task_id)
                         te.task_id,
                         te.user_id,
-                        CASE
-                            WHEN te.state = 'REQUEST_FOR_MAPPING' THEN 'request logs'
-                            WHEN te.state = 'LOCKED_FOR_MAPPING' OR te.state = 'IMAGE_UPLOADED' THEN 'ongoing'
-                            WHEN te.state = 'IMAGE_PROCESSED' THEN 'completed'
-                            WHEN te.state = 'UNFLYABLE_TASK' THEN 'unflyable task'
-                            WHEN te.state = 'IMAGE_PROCESSING_FAILED' THEN 'task failed'
-
-                            ELSE ''
-                        END AS calculated_state
+                        te.state
                     FROM
                         task_events te
                     ORDER BY
@@ -283,7 +277,7 @@ async def one(db: Connection, project_id: uuid.UUID):
                         ST_YMin(ST_Envelope(t.outline)) AS ymin,
                         ST_XMax(ST_Envelope(t.outline)) AS xmax,
                         ST_YMax(ST_Envelope(t.outline)) AS ymax,
-                        COALESCE(tsc.calculated_state) AS state,
+                        tsc.state AS state,
                         tsc.user_id,
                         u.name,
                         ST_Area(ST_Transform(t.outline, 3857)) / 1000000 AS task_area
@@ -343,7 +337,7 @@ async def all(
             await cur.execute(
                 """
                 SELECT
-                    p.id, p.slug, p.name, p.description, p.per_task_instructions,
+                    p.id, p.slug, p.name, p.description, p.per_task_instructions, p.created_at,
                     ST_AsGeoJSON(p.outline)::jsonb AS outline,
                     p.requires_approval_from_manager_for_locking,
 
@@ -541,6 +535,7 @@ class ProjectInfo(BaseModel):
     ongoing_task_count: Optional[int] = 0
     completed_task_count: Optional[int] = 0
     status: Optional[str] = "not-started"
+    created_at: datetime
 
     @model_validator(mode="after")
     def set_image_url(cls, values):

src/backend/app/tasks/task_logic.py

@@ -7,26 +7,6 @@
 from datetime import datetime
 
 
-async def get_current_state(db, project_id, task_id):
-    try:
-        async with db.cursor() as cur:
-            await cur.execute(
-                """
-                SELECT DISTINCT ON (state) state
-                FROM task_events
-                WHERE task_id = %(task_id)s AND project_id = %(project_id)s
-                ORDER BY state, created_at DESC
-                """,
-                {"task_id": task_id, "project_id": project_id},
-            )
-            return await cur.fetchone()
-
-    except Exception as err:
-        raise HTTPException(
-            status_code=HTTPStatus.INTERNAL_SERVER_ERROR, detail=str(err)
-        )
-
-
 async def update_take_off_point_in_db(
     db: Connection, task_id: uuid.UUID, take_off_point: str
 ):

src/backend/app/tasks/task_routes.py

@@ -401,4 +401,43 @@ async def new_event(
                 detail.updated_at,
             )
 
+        case EventType.IMAGE_UPLOAD:
+            current_task_state = await task_logic.get_task_state(
+                db, project_id, task_id
+            )
+            if not current_task_state:
+                raise HTTPException(
+                    status_code=400, detail="Task is not ready for image upload."
+                )
+            state = current_task_state.get("state")
+            locked_user_id = current_task_state.get("user_id")
+
+            # Determine error conditions: Current State must be IMAGE_UPLOADED or IMAGE_PROCESSING_FAILED or lokec for mapping.
+            if state not in (
+                State.IMAGE_UPLOADED.name,
+                State.IMAGE_PROCESSING_FAILED.name,
+                State.LOCKED_FOR_MAPPING.name,
+            ):
+                raise HTTPException(
+                    status_code=400,
+                    detail="Task state does not match expected state for image upload.",
+                )
+
+            if user_id != locked_user_id:
+                raise HTTPException(
+                    status_code=403,
+                    detail="You cannot upload an image for this task as it is locked by another user.",
+                )
+
+            return await task_logic.update_task_state(
+                db,
+                project_id,
+                task_id,
+                user_id,
+                f"Task image uploaded by user {user_data.name}.",
+                State[state],
+                State.IMAGE_UPLOADED,
+                detail.updated_at,
+            )
+
     return True