fix: enable worker support for dimension reductionn

tattle-made · Sep 9, 2024 · 89df059 · 89df059
1 parent 9e5195a
commit 89df059
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 16 deletions.
diff --git a/src/worker/clustering_media/clustering_media_payload_writer.py b/src/worker/clustering_media/clustering_media_payload_writer.py
@@ -10,12 +10,14 @@
     dummy_input = {
         "path": "https://raw.githubusercontent.com/aatmanvaidya/audio-files/main/clustering-media/media_payload.json",
         "video": {
-            "n_clusters": 3
+            "n_clusters": 3,
+            "tsne": True
         },
         "audio": {
-            "n_clusters": 3
+            "n_clusters": 3,
+            "tsne": True
         }
     }
     feluda.queue.message(clustering_media_index_queue, dummy_input)
 except Exception as e:
-    print("Error Initializing Indexer", e)
+    print("Error Initializing Indexer", e)
diff --git a/src/worker/clustering_media/clustering_media_worker.py b/src/worker/clustering_media/clustering_media_worker.py
@@ -9,10 +9,10 @@
 
 log = Logger(__name__)
 
-def make_report_indexed(clustering_results_json, status):
+def make_report_indexed(clustering_results_json, dim_reduction_results_json, status):
     report = {}
     report["clustering_results"] = clustering_results_json
-    # report["dim_reduction_results"] = dim_reduction_results_json
+    report["dim_reduction_results"] = dim_reduction_results_json
     report["status"] = status
     report["status_code"] = 200
     return json.dumps(report)
@@ -128,13 +128,17 @@ def worker(ch, method, properties, body):
             "audio": clustering_results_audio,
             "video": clustering_results_video
         }
-        # log.info("Calculating t-SNE co-ordinates")
-        # dim_reduction_results_json = []
-        # if audio_config.get("tsne"):
-        #     dim_reduction_results_json.extend(dimension_reduction.perform_reduction(audio_embeddings))
-        # if video_config.get("tsne"):
-        #     dim_reduction_results_json.extend(dimension_reduction.perform_reduction(video_embeddings))
-        report = make_report_indexed(clustering_results_json, "indexed")
+        log.info("Calculating t-SNE co-ordinates")
+        if audio_config.get("tsne"):
+            dim_reduction_results_audio = dimension_reduction.run(audio_embeddings)
+        if video_config.get("tsne"):
+            dim_reduction_results_video = dimension_reduction.run(video_embeddings)
+
+        dim_reduction_results_json = {
+            "audio": dim_reduction_results_audio,
+            "video": dim_reduction_results_video
+        }
+        report = make_report_indexed(clustering_results_json, dim_reduction_results_json, "indexed")
         log.info("Report generated")
         feluda.queue.message(
             feluda.config.queue.parameters.queues[1]["name"], report
@@ -158,7 +162,7 @@ def worker(ch, method, properties, body):
     vid_vec_rep_clip.initialize(param={})
     classify_video_zero_shot.initialize(param={})
     cluster_embeddings.initialize(param={})
-    # dimension_reduction.setup_reduction(model_type='tsne', params={})
+    dimension_reduction.initialize(params={})
 
     # start listening to the queue
     feluda.queue.listen(clustering_media_index_queue, clustering_worker(feluda))

diff --git a/src/worker/clustering_media/config.yml b/src/worker/clustering_media/config.yml
@@ -22,7 +22,7 @@ operators :
     - name: "Clustering"
       type: "cluster_embeddings"
       parameters: {}
-    # - name: "tSNE"
-    #   type: "dimension_reduction"
-    #   parameters: {}
+    - name: "tSNE"
+      type: "dimension_reduction"
+      parameters: {}