update

pipeline.py

@@ -15,9 +15,6 @@ class PreTrainedPipeline():
         # Preload all the elements you are going to need at inference.
         # For instance your model, processors, tokenizer that might be needed.
         # This function is only called once, so do all the heavy processing I/O here"""
-        # raise NotImplementedError(
-        #     "Please implement PreTrainedPipeline __init__ function"
-        # )
 
         ## load model
         tf.compat.v1.reset_default_graph()
@@ -32,7 +29,6 @@ class PreTrainedPipeline():
         print(f"restoring model {latest_check_point}")
         saver.restore(sess, latest_check_point)
 
-        ## 
         self.sess = sess
         self.model = model
 
@@ -51,21 +47,45 @@ class PreTrainedPipeline():
         #     "Please implement PreTrainedPipeline __call__ function"
         # )
 
-        vec = np.asarray(json.loads(inputs))[np.newaxis, :, np.newaxis, :]
+        vec = np.asarray(json.loads(inputs))
+        vec = self.reshape_input(vec) # (nb, nt, nsta, nch)
+        vec = self.normalize(vec)
 
         feed = {self.model.X: vec, self.model.drop_rate: 0, self.model.is_training: False}
         preds = self.sess.run(self.model.preds, feed_dict=feed)
 
         picks = extract_picks(preds)#, station_ids=data.id, begin_times=data.timestamp, waveforms=vec_raw)
-
+        picks = [{'phase_index': x['phase_index'], 'phase_score': x['phase_score'], 'phase_type': x['phase_type']} for x in picks]
         # picks = [{k: v for k, v in pick.items() if k in ["station_id", "phase_time", "phase_score", "phase_type", "dt"]} for pick in picks]
 
         # return picks
-        return [[picks, {"label": "debug", "score": 0.1}]]
+        # return [[picks, {"label": "debug", "score": 0.1}]]
+        return [picks]
+
+    def normalize(self, vec):
+        mu = np.mean(vec, axis=1, keepdims=True)
+        std = np.std(vec, axis=1, keepdims=True)
+        std[std == 0] = 1.0
+        vec = (vec - mu) / std
+        return vec
+
+    def reshape_input(self, vec):
+        if len(vec.shape) == 2:
+            vec = vec[np.newaxis, :, np.newaxis, :]
+        elif len(vec.shape) == 3:
+            vec = vec[np.newaxis, :, :, :]
+        else:
+            pass
+        return vec
 
 
 if __name__ == "__main__":
+    import obspy
+    waveform = obspy.read()
+    array = np.array([x.data for x in waveform]).T
+    
     pipeline = PreTrainedPipeline()
-    inputs = np.random.rand(1000, 3).tolist()
+    inputs = array.tolist()
     inputs = json.dumps(inputs)
     picks = pipeline(inputs)
+    print(picks)

tests/test_api.ipynb

@@ -2,14 +2,60 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 30,
+   "execution_count": 20,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import obspy\n",
+    "import numpy as np"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "(3000, 3)\n"
+     ]
+    }
+   ],
+   "source": [
+    "waveform = obspy.read()\n",
+    "array = np.array([x.data for x in waveform]).T\n",
+    "print(array.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "(3000, 3)\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(array.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "[[{'label': 'debug', 'score': 0.1}]]\n"
+      "[[[{'file_name': '0000', 'station_id': '0000', 'begin_time': '1970-01-01T00:00:00.000+00:00', 'phase_index': 573, 'phase_time': '1970-01-01T00:00:05.730+00:00', 'phase_score': 0.999, 'phase_type': 'S', 'dt': 0.01}], {'label': 'debug', 'score': 0.1}]]\n"
      ]
     }
    ],
@@ -27,8 +73,8 @@
     "    return response.json()\n",
     "    # return json.loads(response.content.decode(\"utf-8\"))\n",
     "\n",
-    "array = np.random.rand(10, 3).tolist()\n",
-    "inputs = json.dumps(array)\n",
+    "# array = np.random.rand(10, 3).tolist()\n",
+    "inputs = json.dumps(array.tolist())\n",
     "data = {\n",
     "\t# \"inputs\": \"I like you. I love you\",\n",
     "    \"inputs\": inputs,\n",
@@ -38,6 +84,27 @@
     "output = query(data)\n",
     "print(output)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {