Update gain_viz/app.py

This commit is contained in:
G gael 2025-12-17 22:42:21 -05:00
parent dbaf222ebc
commit ee27d5cc7c

View File

@ -9,9 +9,17 @@ import time
import serial
import json
import subprocess
import io
import base64 # ADD THIS IMPORT
from PIL import Image
from flask_socketio import SocketIO, emit
# Define PLOT_PATH at the top level
PLOT_PATH = os.path.join(os.getcwd(), "plot.png")
app = Flask(__name__)
PLOT_PATH = os.path.join(os.getcwd(), "plot.png")
# Initialize SocketIO with proper configuration
socketio = SocketIO(app, cors_allowed_origins="*", async_mode='threading')
# ----------------- Shared Config -----------------
config = {
@ -24,7 +32,7 @@ config = {
"center_freq": 3.415e9,
"NFFT": 1024,
"tcp_port": 5556,
"streaming": False, # Added streaming state
"streaming": False,
}
config_lock = threading.Lock()
@ -45,6 +53,13 @@ tmux_lock = threading.Lock()
tmux_thread = None
tmux_stop_event = threading.Event()
# In-memory plot storage
plot_lock = threading.Lock()
plot_buffer = io.BytesIO()
# WebSocket sender thread
websocket_thread = None
# ----------------- Serial / SCM -----------------
def connect_serial(port, baudrate=115200, timeout=1):
"""Connect to a serial port with even parity."""
@ -247,19 +262,22 @@ def generate_spectrum_plot():
else:
iq_sample = np.pad(iq_all, (window_samples - len(iq_all), 0))
# Create plot
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 6))
fig.subplots_adjust(hspace=0.4)
# Create plot with optimized settings
plt.rcParams['savefig.dpi'] = 80
plt.rcParams['figure.dpi'] = 80
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 5))
fig.subplots_adjust(hspace=0.3)
# Time-domain plot
times_ms = np.arange(len(iq_sample)) * 1000 / sample_rate
ax1.plot(times_ms, np.real(iq_sample), label="Real", color='b')
ax1.plot(times_ms, np.imag(iq_sample), label="Imag", color='r')
ax1.plot(times_ms, np.real(iq_sample), label="Real", color='b', linewidth=0.8)
ax1.plot(times_ms, np.imag(iq_sample), label="Imag", color='r', linewidth=0.8)
ax1.set_xlim(0, window_ms)
ax1.set_xlabel("Time (ms)")
ax1.set_ylabel("IQ Amplitude")
ax1.grid(True, which='both', linestyle='--', linewidth=0.5)
ax1.legend()
ax1.legend(fontsize=8)
# Spectrogram
cmap = plt.get_cmap('twilight')
@ -281,25 +299,60 @@ def generate_spectrum_plot():
)
ax2.xaxis.set_minor_locator(ticker.AutoMinorLocator())
plt.savefig(PLOT_PATH, bbox_inches='tight')
# Save to buffer
buf = io.BytesIO()
plt.savefig(buf, format='png', bbox_inches='tight', dpi=80)
buf.seek(0)
# Store in global buffer and send via WebSocket
with plot_lock:
plot_buffer.seek(0)
plot_buffer.truncate()
plot_buffer.write(buf.getvalue())
plot_buffer.seek(0)
# Send via SocketIO
img_data = base64.b64encode(buf.getvalue()).decode('utf-8')
socketio.emit('plot_update', {'image': img_data})
plt.close(fig)
buf.close()
except zmq.Again:
# No new data
fig, ax = plt.subplots(figsize=(12, 6))
fig, ax = plt.subplots(figsize=(10, 5))
ax.text(0.5, 0.5, "Waiting for data...",
ha='center', va='center', transform=ax.transAxes, fontsize=16)
ha='center', va='center', transform=ax.transAxes, fontsize=14)
ax.set_title("Spectrum Analyzer - No Data (Streaming Active)")
plt.savefig(PLOT_PATH, bbox_inches='tight')
buf = io.BytesIO()
plt.savefig(buf, format='png', bbox_inches='tight', dpi=80)
buf.seek(0)
with plot_lock:
plot_buffer.seek(0)
plot_buffer.truncate()
plot_buffer.write(buf.getvalue())
plt.close(fig)
buf.close()
except Exception as e:
print(f"Plot generation error: {e}")
fig, ax = plt.subplots(figsize=(12, 6))
fig, ax = plt.subplots(figsize=(10, 5))
ax.text(0.5, 0.5, f"Error: {str(e)}",
ha='center', va='center', transform=ax.transAxes, fontsize=12)
ax.set_title("Spectrum Analyzer - Error")
plt.savefig(PLOT_PATH, bbox_inches='tight')
buf = io.BytesIO()
plt.savefig(buf, format='png', bbox_inches='tight', dpi=80)
buf.seek(0)
with plot_lock:
plot_buffer.seek(0)
plot_buffer.truncate()
plot_buffer.write(buf.getvalue())
plt.close(fig)
buf.close()
time.sleep(0.1)
@ -338,12 +391,22 @@ def stop_plotting():
plot_thread.join(timeout=2.0)
# Create stopped message plot
fig, ax = plt.subplots(figsize=(12, 6))
fig, ax = plt.subplots(figsize=(10, 5))
ax.text(0.5, 0.5, "Streaming Stopped\nClick Start to begin",
ha='center', va='center', transform=ax.transAxes, fontsize=16)
ha='center', va='center', transform=ax.transAxes, fontsize=14)
ax.set_title("Spectrum Analyzer - Stopped")
plt.savefig(PLOT_PATH, bbox_inches='tight')
buf = io.BytesIO()
plt.savefig(buf, format='png', bbox_inches='tight', dpi=80)
buf.seek(0)
with plot_lock:
plot_buffer.seek(0)
plot_buffer.truncate()
plot_buffer.write(buf.getvalue())
plt.close(fig)
buf.close()
print("Plotting thread stopped")
return True
@ -355,7 +418,7 @@ def pause_plotting():
if pause_event.is_set():
pause_event.clear()
print("Plotting resumed")
return "resumed"
return "running"
else:
pause_event.set()
print("Plotting paused")
@ -364,7 +427,11 @@ def pause_plotting():
# ----------------- Flask Routes -----------------
@app.route('/')
def index():
return render_template('index.html')
# Read the template file and inject Socket.IO script
template_path = 'templates/index.html'
if os.path.exists(template_path):
return render_template('index.html')
return "Template not found", 404
@app.route('/update_gains', methods=['POST'])
def update_gains():
@ -397,9 +464,32 @@ def update_gains():
@app.route('/plot')
def plot():
try:
return send_file(PLOT_PATH, mimetype='image/png')
with plot_lock:
plot_buffer.seek(0)
img_data = plot_buffer.read()
if not img_data:
# Return placeholder if buffer is empty
return send_file(PLOT_PATH, mimetype='image/png')
# Create a new BytesIO object for this request
img_io = io.BytesIO(img_data)
img_io.seek(0)
response = send_file(
img_io,
mimetype='image/png',
cache_timeout=0
)
response.headers['Cache-Control'] = 'no-cache, no-store, must-revalidate'
response.headers['Pragma'] = 'no-cache'
response.headers['Expires'] = '0'
return response
except Exception as e:
return send_file(PLOT_PATH, mimetype='image/png')
print(f"Error serving plot: {e}")
try:
return send_file(PLOT_PATH, mimetype='image/png')
except:
return "Error serving plot", 500
@app.route('/get_gains')
def get_gains():
@ -452,11 +542,12 @@ def start_stream():
try:
success = start_plotting()
if success:
start_tmux_capture() # Start capturing tmux output
start_tmux_capture()
return jsonify({"status": "success", "message": "Streaming started"})
else:
return jsonify({"status": "error", "message": "Failed to start streaming"}), 500
except Exception as e:
print(f"Error starting stream: {e}")
return jsonify({"status": "error", "message": f"Error: {str(e)}"}), 500
@app.route('/stop_stream', methods=['POST'])
@ -464,11 +555,12 @@ def stop_stream():
try:
success = stop_plotting()
if success:
stop_tmux_capture() # Stop capturing tmux output
stop_tmux_capture()
return jsonify({"status": "success", "message": "Streaming stopped"})
else:
return jsonify({"status": "error", "message": "Failed to stop streaming"}), 500
except Exception as e:
print(f"Error stopping stream: {e}")
return jsonify({"status": "error", "message": f"Error: {str(e)}"}), 500
@app.route('/pause_stream', methods=['POST'])
@ -477,6 +569,7 @@ def pause_stream():
result = pause_plotting()
return jsonify({"status": "success", "message": f"Streaming {result}", "state": result})
except Exception as e:
print(f"Error pausing stream: {e}")
return jsonify({"status": "error", "message": f"Error: {str(e)}"}), 500
@app.route('/get_stream_state', methods=['GET'])
@ -499,6 +592,15 @@ def get_tmux_output():
with tmux_lock:
return jsonify({"output": tmux_output})
# WebSocket event handlers
@socketio.on('connect')
def handle_connect():
print('Client connected via WebSocket')
@socketio.on('disconnect')
def handle_disconnect():
print('Client disconnected from WebSocket')
def save_config():
with config_lock:
cfg = dict(config)
@ -513,14 +615,16 @@ def save_config():
def main():
# Ensure placeholder image exists
if not os.path.exists(PLOT_PATH):
fig, ax = plt.subplots(figsize=(12, 6))
ax.text(0.5, 0.5, "Click Start to begin streaming", ha='center', va='center', fontsize=16)
fig, ax = plt.subplots(figsize=(10, 5))
ax.text(0.5, 0.5, "Click Start to begin streaming", ha='center', va='center', fontsize=14)
ax.set_title("Gain-Viz Spectrum Analyzer - Ready")
plt.savefig(PLOT_PATH)
plt.savefig(PLOT_PATH, bbox_inches='tight', dpi=80)
plt.close(fig)
print("Gain-Viz server started. Use the web interface to control streaming.")
app.run(host="0.0.0.0", port=5000, debug=True, use_reloader=False)
print("Gain-Viz server starting on http://0.0.0.0:5000")
print("WebSocket support enabled")
# Run the SocketIO server
socketio.run(app, host="0.0.0.0", port=5000, debug=False, use_reloader=False, allow_unsafe_werkzeug=True)
if __name__ == '__main__':
main()