Selaa lähdekoodia

feat: add JSONEncoder class

Handles serialisation of custom classes via to_json() method if present
main
Rob Hallam 2 kuukautta sitten
vanhempi
commit
8d188cab19
1 muutettua tiedostoa jossa 11 lisäystä ja 0 poistoa
  1. +11
    -0
      pipeline/video_producers.py

+ 11
- 0
pipeline/video_producers.py Näytä tiedosto

@@ -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):


Ladataan…
Peruuta
Tallenna