From 785584aff585bd3362e99fc14681a59ae501cb08 Mon Sep 17 00:00:00 2001
From: Junzi Sun <junzisun@gmail.com>
Date: Thu, 22 Aug 2019 16:59:12 +0200
Subject: [PATCH] improve singal process efficiency

---
 pyModeS/extra/rtlreader.py | 91 ++++++++++++++++++++++----------------
 1 file changed, 53 insertions(+), 38 deletions(-)

diff --git a/pyModeS/extra/rtlreader.py b/pyModeS/extra/rtlreader.py
index df9e861..ec3747d 100644
--- a/pyModeS/extra/rtlreader.py
+++ b/pyModeS/extra/rtlreader.py
@@ -1,24 +1,25 @@
-import sys
 import numpy as np
 import pyModeS as pms
 from rtlsdr import RtlSdr
 from threading import Thread
 import time
 
-amplitude_threshold = 0.2
 modes_sample_rate = 2e6
 modes_frequency = 1090e6
 buffer_size = 1024 * 100
-read_size = 1024 * 8
+read_size = 1024 * 20
+
 pbits = 8
 fbits = 112
-preamble = "1010000101000000"
+preamble = [1, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0]
+th_amp = 0.2  # signal amplitude threshold for 0 and 1 bit
+th_amp_diff = 0.8  # signal amplitude threshold difference between 0 and 1 bit
 
 
 class RtlReader(Thread):
     def __init__(self, debug=False):
         super(RtlReader, self).__init__()
-        self.signal_buffer = np.array([])
+        self.signal_buffer = []
         self.debug = debug
         self.sdr = RtlSdr()
         self.sdr.sample_rate = modes_sample_rate
@@ -29,37 +30,36 @@ class RtlReader(Thread):
     def _process_buffer(self):
         messages = []
 
-        pulses_array = np.where(self.signal_buffer < amplitude_threshold, 0, 1)
-        pulses = "".join(str(x) for x in pulses_array)
+        # signal_array = np.array(self.signal_buffer)
+        # pulses_array = np.where(np.array(self.signal_buffer) < th_amp, 0, 1)
+        # pulses = "".join(str(x) for x in pulses_array)
+        buffer_length = len(self.signal_buffer)
 
         i = 0
-        while i < len(pulses):
-            if pulses[i] == 0:
+        while i < buffer_length:
+            if self.signal_buffer[i] < th_amp:
                 i += 1
                 continue
 
-            if pulses[i : i + pbits * 2] == preamble:
+            # if pulses[i : i + pbits * 2] == preamble:
+            if self._check_preamble(self.signal_buffer[i : i + pbits * 2]):
                 frame_start = i + pbits * 2
                 frame_end = i + pbits * 2 + (fbits + 1) * 2
-                frame_pulses = pulses[frame_start:frame_end]
+                frame_length = (fbits + 1) * 2
+                frame_pulses = self.signal_buffer[frame_start:frame_end]
 
                 msgbin = ""
-                for j in range(0, len(frame_pulses), 2):
+                for j in range(0, frame_length, 2):
                     p2 = frame_pulses[j : j + 2]
-                    if p2 == "10":
-                        c = "1"
-                    elif p2 == "01":
-                        c = "0"
-                    elif p2 == "11":
-                        a2 = self.signal_buffer[
-                            frame_start + j : frame_start + j + 2
-                        ]
-                        if a2[0] > a2[1]:
-                            c = "1"
-                        else:
-                            c = "0"
-                    elif p2 == "00":
+                    if len(p2) < 2:
                         break
+
+                    if p2[0] < th_amp and p2[1] < th_amp:
+                        break
+                    elif p2[0] >= p2[1]:
+                        c = "1"
+                    elif p2[0] < p2[1]:
+                        c = "0"
                     else:
                         msgbin = ""
                         break
@@ -75,7 +75,8 @@ class RtlReader(Thread):
                     if self.debug:
                         self._debug_msg(msghex)
 
-            elif i > len(self.signal_buffer) - pbits * 2 - fbits * 2:
+            elif i > buffer_length - 500:
+                # save some for next process
                 break
             else:
                 i += 1
@@ -84,6 +85,16 @@ class RtlReader(Thread):
         self.signal_buffer = self.signal_buffer[i:]
         return messages
 
+    def _check_preamble(self, pulses):
+        if len(pulses) != 16:
+            return False
+
+        for i in range(16):
+            if abs(pulses[i] - preamble[i]) > th_amp_diff:
+                return False
+
+        return True
+
     def _check_msg(self, msg):
         df = pms.df(msg)
         msglen = len(msg)
@@ -105,30 +116,34 @@ class RtlReader(Thread):
         elif df in [4, 5, 11] and msglen == 14:
             print(msg, pms.icao(msg))
         else:
-            print("[*]", msg)
+            # print("[*]", msg)
             pass
 
     def _read_callback(self, data, rtlsdr_obj):
-        self.signal_buffer = np.concatenate(
-            (self.signal_buffer, np.absolute(data))
-        )
+        # scaling signal (imporatant)
+        amp = np.absolute(data)
+        amp_norm = np.interp(amp, (amp.min(), amp.max()), (0, 1))
+        self.signal_buffer.extend(amp_norm.tolist())
 
         if len(self.signal_buffer) >= buffer_size:
-            try:
-                messages = self._process_buffer()
-                self.handle_messages(messages)
-            except KeyboardInterrupt:
-                sys.exit(1)
+            messages = self._process_buffer()
+            self.handle_messages(messages)
 
     def handle_messages(self, messages):
         """re-implement this method to handle the messages"""
         for msg, t in messages:
-            pass
             # print("%15.9f %s" % (t, msg))
+            pass
+
+    def stop(self):
+        self.sdr.cancel_read_async()
 
     def run(self):
         self.sdr.read_samples_async(self._read_callback, read_size)
-        # while True:
+
+        # count = 1
+        # while count < 1000:
+        #     count += 1
         #     data = self.sdr.read_samples(read_size)
         #     self._read_callback(data, None)
 
@@ -136,5 +151,5 @@ class RtlReader(Thread):
 if __name__ == "__main__":
     rtl = RtlReader()
     rtl.debug = True
-    # rtl.daemon = True
     rtl.start()
+    rtl.join()