Export Annotations Task

bats_ai/core/tasks/export_task.py

@@ -4,10 +4,14 @@
 from datetime import timedelta
 from io import BytesIO, StringIO
 import json
+from urllib.parse import urljoin
 import zipfile
 
+from django.conf import settings
 from django.contrib.auth.models import User
 from django.core.files import File
+from django.core.files.storage import default_storage
+from django.db.models import Prefetch
 from django.utils.timezone import now
 
 from bats_ai.celery import app
@@ -18,6 +22,7 @@
     RecordingTag,
     SequenceAnnotations,
 )
+from bats_ai.core.models.recording_annotation import RecordingAnnotationSpecies
 
 
 def build_filters(filters, *, has_confidence=False):
@@ -179,6 +184,89 @@ def export_tag_annotation_summary_task(self, export_id: int):
         raise
 
 
+@app.task(bind=True)
+def export_recording_annotation_hierarchy_task(self, export_id: int):
+    export_record = ExportedAnnotationFile.objects.get(pk=export_id)
+    try:
+        recordings_payload = _build_recording_annotations_payload()
+
+        buffer = BytesIO()
+        with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as zipf:
+            zipf.writestr(
+                "recording_annotations.json",
+                json.dumps({"recordings": recordings_payload}, indent=2),
+            )
+
+        buffer.seek(0)
+        filename = f"recording-annotations-{export_id}.zip"
+        export_record.file.save(filename, File(buffer), save=False)
+        export_record.download_url = export_record.file.url
+        export_record.status = "complete"
+        export_record.expires_at = now() + timedelta(hours=24)
+        export_record.save()
+    except Exception:
+        export_record.status = "failed"
+        export_record.save()
+        raise
+
+
+def _build_recording_annotations_payload():
+    species_links_prefetch = Prefetch(
+        "recordingannotationspecies_set",
+        queryset=RecordingAnnotationSpecies.objects.select_related("species").order_by("order"),
+        to_attr="ordered_species_links",
+    )
+    annotations = (
+        RecordingAnnotation.objects.select_related("recording", "owner")
+        .prefetch_related(species_links_prefetch)
+        .order_by("recording_id", "id")
+    )
+
+    recordings_by_id = {}
+    for annotation in annotations:
+        recording = annotation.recording
+        recording_entry = recordings_by_id.get(recording.id)
+        if recording_entry is None:
+            recording_entry = {
+                "recording_id": recording.id,
+                "filename": recording.name,
+                "grts_cell_id": recording.grts_cell_id,
+                "sample_frame_id": recording.sample_frame_id,
+                "submitted_annotations": 0,
+                "unsubmitted_annotations": 0,
+                "spectrogram_url": urljoin(
+                    settings.BATAI_WEB_URL, f"/recording/{recording.id}/spectrogram"
+                ),
+                "wav_download_url": (
+                    default_storage.url(recording.audio_file.name) if recording.audio_file else None
+                ),
+                "annotations": [],
+            }
+            recordings_by_id[recording.id] = recording_entry
+
+        if annotation.submitted:
+            recording_entry["submitted_annotations"] += 1
+        else:
+            recording_entry["unsubmitted_annotations"] += 1
+
+        species_codes = [
+            species_link.species.species_code for species_link in annotation.ordered_species_links
+        ]
+        recording_entry["annotations"].append(
+            {
+                "annotation_id": annotation.id,
+                "user": annotation.owner.username,
+                "species_codes": species_codes,
+                "confidence": annotation.confidence,
+                "additional_data": annotation.additional_data,
+                "comment": annotation.comments,
+                "submitted": annotation.submitted,
+            }
+        )
+
+    return list(recordings_by_id.values())
+
+
 def _collect_tag_summary_rows():
     tag_rows = []
     tag_user_rows = []

bats_ai/core/views/configuration.py

@@ -9,7 +9,10 @@
 from ninja.pagination import RouterPaginated
 
 from bats_ai.core.models import Configuration, ExportedAnnotationFile
-from bats_ai.core.tasks.export_task import export_tag_annotation_summary_task
+from bats_ai.core.tasks.export_task import (
+    export_recording_annotation_hierarchy_task,
+    export_tag_annotation_summary_task,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -99,3 +102,20 @@ def export_tag_summary(request):
     )
     export_tag_annotation_summary_task.delay(export.id)
     return {"exportId": export.id}
+
+
+@router.post(
+    "/export-recording-annotations",
+    response=ExportTagSummaryResponse,
+)
+def export_recording_annotations(request):
+    if not request.user.is_authenticated or not request.user.is_superuser:
+        return JsonResponse({"error": "Permission denied"}, status=403)
+
+    export = ExportedAnnotationFile.objects.create(
+        filters_applied={"type": "recording_annotation_hierarchy"},
+        status="pending",
+        expires_at=now() + timedelta(hours=24),
+    )
+    export_recording_annotation_hierarchy_task.delay(export.id)
+    return {"exportId": export.id}

bats_ai/core/views/export_annotation.py

@@ -23,10 +23,22 @@ def _is_tag_annotation_summary_export(export: ExportedAnnotationFile) -> bool:
     )
 
 
+def _is_recording_annotation_hierarchy_export(
+    export: ExportedAnnotationFile,
+) -> bool:
+    filters_applied = export.filters_applied
+    return (
+        isinstance(filters_applied, dict)
+        and filters_applied.get("type") == "recording_annotation_hierarchy"
+    )
+
+
 def _can_access_export(request, export: ExportedAnnotationFile) -> bool:
     # Tag annotation summary exports include user-level aggregate stats,
     # so only admins can access them.
-    if _is_tag_annotation_summary_export(export):
+    if _is_tag_annotation_summary_export(export) or _is_recording_annotation_hierarchy_export(
+        export
+    ):
         return request.user.is_authenticated and request.user.is_superuser
     return True
 

client/src/api/api.ts

@@ -787,6 +787,13 @@ async function exportTagSummary(): Promise<ExportTagSummaryResponse> {
   return result.data;
 }
 
+async function exportRecordingAnnotations(): Promise<ExportTagSummaryResponse> {
+  const result = await axiosInstance.post<ExportTagSummaryResponse>(
+    "/configuration/export-recording-annotations",
+  );
+  return result.data;
+}
+
 export interface VettingDetails {
   id: number;
   user_id: number;
@@ -911,6 +918,7 @@ export {
   getFileAnnotationDetails,
   getExportStatus,
   exportTagSummary,
+  exportRecordingAnnotations,
   getRecordingTags,
   getUnsubmittedNeighbors,
   getComputedPulseContour,

client/src/views/Admin.vue

@@ -1,7 +1,11 @@
 <script lang="ts">
 import { reactive, defineComponent, watch, ref, type Ref } from "vue";
 import useState from "@use/useState";
-import { exportTagSummary, patchConfiguration } from "../api/api";
+import {
+  exportRecordingAnnotations,
+  exportTagSummary,
+  patchConfiguration,
+} from "../api/api";
 import NABatAdmin from "./NABat/NABatAdmin.vue";
 import ColorPickerMenu from "@components/ColorPickerMenu.vue";
 import ColorSchemeSelect from "@components/ColorSchemeSelect.vue";
@@ -86,6 +90,11 @@ export default defineComponent({
       exportId.value = result.exportId;
     };
 
+    const runRecordingAnnotationsExport = async () => {
+      const result = await exportRecordingAnnotations();
+      exportId.value = result.exportId;
+    };
+
     const clearExport = () => {
       exportId.value = null;
     };
@@ -99,6 +108,7 @@ export default defineComponent({
       defaultColorScheme,
       exportId,
       runTagSummaryExport,
+      runRecordingAnnotationsExport,
       clearExport,
     };
   },
@@ -208,6 +218,14 @@ export default defineComponent({
             >
               Export Tag Annotation Summary
             </v-btn>
+            <v-btn
+              color="secondary"
+              variant="outlined"
+              class="mx-2"
+              @click="runRecordingAnnotationsExport"
+            >
+              Export Recording Annotations
+            </v-btn>
           </v-row>
         </v-card-actions>
         <Exporting