|
4 | 4 | from datetime import timedelta |
5 | 5 | from io import BytesIO, StringIO |
6 | 6 | import json |
| 7 | +from urllib.parse import urljoin |
7 | 8 | import zipfile |
8 | 9 |
|
| 10 | +from django.conf import settings |
9 | 11 | from django.contrib.auth.models import User |
10 | 12 | from django.core.files import File |
| 13 | +from django.core.files.storage import default_storage |
| 14 | +from django.db.models import Prefetch |
11 | 15 | from django.utils.timezone import now |
12 | 16 |
|
13 | 17 | from bats_ai.celery import app |
14 | 18 | from bats_ai.core.models import ( |
15 | 19 | Annotations, |
16 | 20 | ExportedAnnotationFile, |
17 | 21 | RecordingAnnotation, |
| 22 | + RecordingAnnotationSpecies, |
18 | 23 | RecordingTag, |
19 | 24 | SequenceAnnotations, |
20 | 25 | ) |
@@ -179,6 +184,89 @@ def export_tag_annotation_summary_task(self, export_id: int): |
179 | 184 | raise |
180 | 185 |
|
181 | 186 |
|
| 187 | +@app.task(bind=True) |
| 188 | +def export_recording_annotation_hierarchy_task(self, export_id: int): |
| 189 | + export_record = ExportedAnnotationFile.objects.get(pk=export_id) |
| 190 | + try: |
| 191 | + recordings_payload = _build_recording_annotations_payload() |
| 192 | + |
| 193 | + buffer = BytesIO() |
| 194 | + with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as zipf: |
| 195 | + zipf.writestr( |
| 196 | + "recording_annotations.json", |
| 197 | + json.dumps({"recordings": recordings_payload}, indent=2), |
| 198 | + ) |
| 199 | + |
| 200 | + buffer.seek(0) |
| 201 | + filename = f"recording-annotations-{export_id}.zip" |
| 202 | + export_record.file.save(filename, File(buffer), save=False) |
| 203 | + export_record.download_url = export_record.file.url |
| 204 | + export_record.status = "complete" |
| 205 | + export_record.expires_at = now() + timedelta(hours=24) |
| 206 | + export_record.save() |
| 207 | + except Exception: |
| 208 | + export_record.status = "failed" |
| 209 | + export_record.save() |
| 210 | + raise |
| 211 | + |
| 212 | + |
| 213 | +def _build_recording_annotations_payload(): |
| 214 | + species_links_prefetch = Prefetch( |
| 215 | + "recordingannotationspecies_set", |
| 216 | + queryset=RecordingAnnotationSpecies.objects.select_related("species").order_by("order"), |
| 217 | + to_attr="ordered_species_links", |
| 218 | + ) |
| 219 | + annotations = ( |
| 220 | + RecordingAnnotation.objects.select_related("recording", "owner") |
| 221 | + .prefetch_related(species_links_prefetch) |
| 222 | + .order_by("recording_id", "id") |
| 223 | + ) |
| 224 | + |
| 225 | + recordings_by_id = {} |
| 226 | + for annotation in annotations: |
| 227 | + recording = annotation.recording |
| 228 | + recording_entry = recordings_by_id.get(recording.id) |
| 229 | + if recording_entry is None: |
| 230 | + recording_entry = { |
| 231 | + "recording_id": recording.id, |
| 232 | + "filename": recording.name, |
| 233 | + "grts_cell_id": recording.grts_cell_id, |
| 234 | + "sample_frame_id": recording.sample_frame_id, |
| 235 | + "submitted_annotations": 0, |
| 236 | + "unsubmitted_annotations": 0, |
| 237 | + "spectrogram_url": urljoin( |
| 238 | + settings.BATAI_WEB_URL, f"/recording/{recording.id}/spectrogram" |
| 239 | + ), |
| 240 | + "wav_download_url": ( |
| 241 | + default_storage.url(recording.audio_file.name) if recording.audio_file else None |
| 242 | + ), |
| 243 | + "annotations": [], |
| 244 | + } |
| 245 | + recordings_by_id[recording.id] = recording_entry |
| 246 | + |
| 247 | + if annotation.submitted: |
| 248 | + recording_entry["submitted_annotations"] += 1 |
| 249 | + else: |
| 250 | + recording_entry["unsubmitted_annotations"] += 1 |
| 251 | + |
| 252 | + species_codes = [ |
| 253 | + species_link.species.species_code for species_link in annotation.ordered_species_links |
| 254 | + ] |
| 255 | + recording_entry["annotations"].append( |
| 256 | + { |
| 257 | + "annotation_id": annotation.id, |
| 258 | + "user": annotation.owner.username, |
| 259 | + "species_codes": species_codes, |
| 260 | + "confidence": annotation.confidence, |
| 261 | + "additional_data": annotation.additional_data, |
| 262 | + "comment": annotation.comments, |
| 263 | + "submitted": annotation.submitted, |
| 264 | + } |
| 265 | + ) |
| 266 | + |
| 267 | + return list(recordings_by_id.values()) |
| 268 | + |
| 269 | + |
182 | 270 | def _collect_tag_summary_rows(): |
183 | 271 | tag_rows = [] |
184 | 272 | tag_user_rows = [] |
|
0 commit comments