From b622585fb9617fc3255fdd2535338a0599c53206 Mon Sep 17 00:00:00 2001
From: Joel Mathew Thomas <90510078+joelmathewthomas@users.noreply.github.com>
Date: Tue, 18 Mar 2025 20:37:20 +0530
Subject: [PATCH] add endpoint /api/spectrogram to calculate spectrograms

---
 api/api/tasks.py    | 14 ++++++++++++++
 api/api/views.py    | 30 ++++++++++++++++++++++++++++++
 api/backend/urls.py |  2 ++
 3 files changed, 46 insertions(+)

diff --git a/api/api/tasks.py b/api/api/tasks.py
index 0c023c0..0c5cc54 100644
--- a/api/api/tasks.py
+++ b/api/api/tasks.py
@@ -131,6 +131,20 @@ def noisereduce_task(file_path):
         return False
 
 @shared_task
+def generate_spectrogram_task(file_path):
+    """Celery task to generate spectrogram"""
+    try:
+        file_path = Path(file_path)
+    
+        # Generate spectrogram
+        spec_db, plot_data = generate_spectrogram(file_path)
+        spec_db = np.nan_to_num(spec_db, nan=-80.0, posinf=-80.0, neginf=-80.0)
+        spec_data_json = json.dumps(spec_db.tolist())
+
+        return True, spec_data_json, plot_data['sr'] 
+    except Exception as e:
+        return False
+@shared_task
 def cleanup_task(file_path):
     """Celery task to cleanup files"""
     file_path = Path(file_path)
diff --git a/api/api/views.py b/api/api/views.py
index cecfd43..6c0cf6f 100644
--- a/api/api/views.py
+++ b/api/api/views.py
@@ -12,6 +12,7 @@ from .tasks import trim_audio_task
 from .tasks import resample_audio_task
 from .tasks import music_separation_task
 from .tasks import noisereduce_task
+from .tasks import generate_spectrogram_task
 from .tasks import cleanup_task
 from freqsplit.input.format_checker import is_supported_format
 
@@ -147,6 +148,35 @@ def noisereduce(request):
     else:
         return Response({"error": "Failed to remove noise from audio"}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)
 
+# Endpoint to generate spectrograms
+@api_view(['POST'])
+def generate_spectrogram(request):
+    """Handle generation of spectrogram"""
+    file_uuid = request.data.get("file_uuid")
+    file_name = request.data.get("file_name")
+    file_path = os.path.join(UPLOAD_DIR, file_uuid, file_name)
+    
+    # Check if file exists
+    if os.path.exists(file_path):
+        # Call Celery task synchronously
+        task = generate_spectrogram_task.apply(args=(file_path,))
+
+        if task.ready() and task.successful():
+            result = task.result
+            if result[0]:
+                return Response(
+                    {
+                        "Status": "Spectrogram generated successfully",
+                        "spectrogram": result[1],
+                        "spec_sr": result[2]
+                    },
+                    status=status.HTTP_200_OK
+                )
+            else:
+                return Response({"error": "Failed to generate spectrogram"}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)
+    else:
+        return Response({"error": "File Not Found"}, status=status.HTTP_404_NOT_FOUND)
+    
 # Endpoint to download audio file or zipped directory
 @api_view(['GET'])
 def download_audio(request):
diff --git a/api/backend/urls.py b/api/backend/urls.py
index 727daa7..a675ee7 100644
--- a/api/backend/urls.py
+++ b/api/backend/urls.py
@@ -23,6 +23,7 @@ from api.views import resample_audio
 from api.views import separate_music
 from api.views import noisereduce
 from api.views import download_audio
+from api.views import generate_spectrogram
 from api.views import cleanup
 from api.views import cleanup_zip
 
@@ -35,6 +36,7 @@ urlpatterns = [
     path('api/separate', separate_music, name="separate_music"),
     path('api/noisereduce', noisereduce, name="noisreduce"),
     path('api/download', download_audio, name="download_audio"),
+    path('api/spectrogram', generate_spectrogram, name="generate_spectrogram"),
     path('api/cleanup', cleanup, name="cleanup"),
     path('api/cleanup_zip', cleanup_zip, name="cleanup_zip")
 ]