From 8d188cab190527616e4bc4217a120b4504d84d36 Mon Sep 17 00:00:00 2001
From: Rob Hallam <0504004h@student.gla.ac.uk>
Date: Tue, 13 Aug 2024 10:54:48 +0100
Subject: [PATCH] feat: add JSONEncoder class

Handles serialisation of custom classes via to_json() method if present
---
 pipeline/video_producers.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/pipeline/video_producers.py b/pipeline/video_producers.py
index 45c09e8..298a3f7 100644
--- a/pipeline/video_producers.py
+++ b/pipeline/video_producers.py
@@ -10,6 +10,10 @@ import tempfile
 # for visualisations:
 import matplotlib.pyplot as plt
 
+# for encoding as JSON
+from pipeline.utils import Feature
+
+
 class Producer(ABC):
     """Generic producer interface."""
     def __init__(self, features):
@@ -141,6 +145,13 @@ class VisualisationProducer(Producer):
         plt.savefig("/tmp/visualisation.png")
         plt.close()
 
+class PipelineJSONEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if hasattr(obj, 'to_json'):
+            return obj.to_json()
+        else:
+            return json.JSONEncoder.default(self, obj)
+
 class JSONProducer(Producer):
     """Produce JSON output"""
     def __init__(self, features):