Worker implementation for processing images in the Antenna task queue (PSv2)

.env.example

@@ -8,3 +8,9 @@ AMI_CLASSIFICATION_THRESHOLD=0.6
 AMI_LOCALIZATION_BATCH_SIZE=2
 AMI_CLASSIFICATION_BATCH_SIZE=20
 AMI_NUM_WORKERS=1
+
+# Antenna API Worker Settings (for processing jobs from Antenna platform)
+# See: https://github.com/RolnickLab/antenna
+AMI_ANTENNA_API_BASE_URL=http://localhost:8000/api/v2
+AMI_ANTENNA_API_AUTH_TOKEN=your_antenna_auth_token_here
+AMI_ANTENNA_API_BATCH_SIZE=4

.gitignore

@@ -142,3 +142,6 @@ db_data/
 # Test files
 sample_images
 bak
+
+# Local scratch for moving untracked files
+scratch/

.pre-commit-config.yaml

@@ -38,7 +38,7 @@ repos:
         types: [pyi]
 
   - repo: https://github.com/pycqa/flake8
-    rev: 3.8.3
+    rev: 4.0.0
     hooks:
       - id: flake8
         files: .

.vscode/launch.json

@@ -0,0 +1,29 @@
+{
+  // Use IntelliSense to learn about possible attributes.
+  // Hover to view descriptions of existing attributes.
+  // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+  "version": "0.2.0",
+  "configurations": [
+    {
+      "name": "Python Debugger: Current File",
+      "type": "debugpy",
+      "request": "launch",
+      "program": "${file}",
+      "console": "integratedTerminal"
+    },
+    {
+      "name": "Run worker",
+      "type": "debugpy",
+      "request": "launch",
+      "module": "trapdata.cli.base",
+      "args": ["worker"]
+    },
+    {
+      "name": "Run api",
+      "type": "debugpy",
+      "request": "launch",
+      "module": "trapdata.cli.base",
+      "args": ["api"]
+    }
+  ]
+}

CLAUDE.md

@@ -13,8 +13,9 @@ This file helps AI agents (like Claude) work efficiently with the AMI Data Compa
 3. **Always prefer command line tools** to avoid expensive API requests (e.g., use git and jq instead of reading whole files)
 4. **Use bulk operations and prefetch patterns** to minimize database queries
 5. **Commit often** - Small, focused commits make debugging easier
-6. **Use TDD whenever possible** - Tests prevent regressions and document expected behavior
-7. **Keep it simple** - Always think hard and evaluate more complex approaches and alternative approaches before moving forward
+6. **Use `git add -p` for staging** - Interactive staging to add only relevant changes, creating logical commits
+7. **Use TDD whenever possible** - Tests prevent regressions and document expected behavior
+8. **Keep it simple** - Always think hard and evaluate more complex approaches and alternative approaches before moving forward
 
 ### Think Holistically
 

README.md

@@ -234,6 +234,60 @@ ami api
 
 View the interactive API docs at http://localhost:2000/
 
+## Running the Antenna Worker
+
+The worker polls the Antenna platform API for queued image processing jobs, downloads images, runs detection and classification, and posts results back to Antenna.
+
+**Setup:**
+
+1. Get your Antenna auth token from your Antenna project settings
+2. Configure the worker in `.env`:
+
+```sh
+AMI_ANTENNA_API_BASE_URL=https://antenna.insectai.org/api/v2  # Or your Antenna instance
+AMI_ANTENNA_API_AUTH_TOKEN=your_token_here
+AMI_ANTENNA_API_BATCH_SIZE=4
+AMI_NUM_WORKERS=2  # Safe for REST API (atomic task dequeue)
+```
+
+**Register pipelines (optional):**
+
+Register available ML pipelines with your Antenna projects:
+
+```sh
+ami worker register "My Worker Name" --project 1 --project 2
+# Or register for all accessible projects:
+ami worker register "My Worker Name"
+```
+
+**Run the worker:**
+
+```sh
+# Process all pipelines:
+ami worker
+
+# Or specify specific pipeline(s):
+ami worker --pipeline moth_binary
+ami worker --pipeline moth_binary --pipeline panama_moths_2024
+```
+
+The worker will:
+
+1. Poll Antenna for jobs matching the specified pipeline(s)
+2. Download images from the job queue
+3. Run detection and classification
+4. Post results back to Antenna
+5. Repeat until queue is empty, then sleep and poll again
+
+**Notes:**
+
+- Multiple workers can run in parallel (they won't duplicate work)
+- Auth token ties results to your Antenna project
+- Worker continues running until interrupted (Ctrl+C)
+- Safe to run multiple workers on different machines
+
+For more information, see the [Antenna platform documentation](https://github.com/RolnickLab/antenna).
+
 ## Web UI demo (Gradio)
 
 A simple web UI is also available to test the inference pipeline. This is a quick way to test models on a remote server via a web browser.

pyproject.toml

@@ -53,6 +53,7 @@ gradio = "^4.41.0"
 
 [tool.pytest.ini_options]
 asyncio_mode = 'auto'
+testpaths = ["trapdata/tests", "trapdata/**/tests"]
 
 [tool.isort]
 profile = "black"

trapdata/antenna/__init__.py

@@ -0,0 +1,20 @@
+"""Antenna platform integration module.
+
+This module provides integration with the Antenna platform for remote image processing.
+It includes:
+- API client for fetching jobs and posting results
+- Worker loop for continuous job processing
+- Pipeline registration with Antenna projects
+- Schemas for Antenna API requests/responses
+- Dataset classes for streaming tasks from the API
+"""
+
+from trapdata.antenna import client, datasets, registration, schemas, worker
+
+__all__ = [
+    "client",
+    "datasets",
+    "registration",
+    "schemas",
+    "worker",
+]

trapdata/antenna/client.py

@@ -0,0 +1,110 @@
+"""Antenna API client for fetching jobs and posting results."""
+
+import requests
+
+from trapdata.antenna.schemas import AntennaJobsListResponse, AntennaTaskResult
+from trapdata.api.utils import get_http_session
+from trapdata.common.logs import logger
+
+
+def get_jobs(
+    base_url: str,
+    auth_token: str,
+    pipeline_slug: str,
+) -> list[int]:
+    """Fetch job ids from the API for the given pipeline.
+
+    Calls: GET {base_url}/jobs?pipeline__slug=<pipeline>&ids_only=1
+
+    Args:
+        base_url: Antenna API base URL (e.g., "http://localhost:8000/api/v2")
+        auth_token: API authentication token
+        pipeline_slug: Pipeline slug to filter jobs
+
+    Returns:
+        List of job ids (possibly empty) on success or error.
+    """
+    with get_http_session(auth_token) as session:
+        try:
+            url = f"{base_url.rstrip('/')}/jobs"
+            params = {
+                "pipeline__slug": pipeline_slug,
+                "ids_only": 1,
+                "incomplete_only": 1,
+            }
+
+            resp = session.get(url, params=params, timeout=30)
+            resp.raise_for_status()
+
+            # Parse and validate response with Pydantic
+            jobs_response = AntennaJobsListResponse.model_validate(resp.json())
+            return [job.id for job in jobs_response.results]
+        except requests.RequestException as e:
+            logger.error(f"Failed to fetch jobs from {base_url}: {e}")
+            return []
+        except Exception as e:
+            logger.error(f"Failed to parse jobs response: {e}")
+            return []
+
+
+def post_batch_results(
+    base_url: str,
+    auth_token: str,
+    job_id: int,
+    results: list[AntennaTaskResult],
+) -> bool:
+    """
+    Post batch results back to the API.
+
+    Args:
+        base_url: Antenna API base URL (e.g., "http://localhost:8000/api/v2")
+        auth_token: API authentication token
+        job_id: Job ID
+        results: List of AntennaTaskResult objects
+
+    Returns:
+        True if successful, False otherwise
+    """
+    url = f"{base_url.rstrip('/')}/jobs/{job_id}/result/"
+    payload = [r.model_dump(mode="json") for r in results]
+
+    with get_http_session(auth_token) as session:
+        try:
+            response = session.post(url, json=payload, timeout=60)
+            response.raise_for_status()
+            logger.info(f"Successfully posted {len(results)} results to {url}")
+            return True
+        except requests.RequestException as e:
+            logger.error(f"Failed to post results to {url}: {e}")
+            return False
+
+
+def get_user_projects(base_url: str, auth_token: str) -> list[dict]:
+    """
+    Fetch all projects the user has access to.
+
+    Args:
+        base_url: Base URL for the API (should NOT include /api/v2)
+        auth_token: API authentication token
+
+    Returns:
+        List of project dictionaries with 'id' and 'name' fields
+    """
+    with get_http_session(auth_token) as session:
+        try:
+            url = f"{base_url.rstrip('/')}/projects/"
+            response = session.get(url, timeout=30)
+            response.raise_for_status()
+            data = response.json()
+
+            projects = data.get("results", [])
+            if isinstance(projects, list):
+                return projects
+            else:
+                logger.warning(
+                    f"Unexpected projects format from {url}: {type(projects)}"
+                )
+                return []
+        except requests.RequestException as e:
+            logger.error(f"Failed to fetch projects from {base_url}: {e}")
+            return []