351 lines
12 KiB
Python
351 lines
12 KiB
Python
import logging
|
|
import sys
|
|
import os
|
|
import threading
|
|
import json
|
|
import time
|
|
import traceback
|
|
import base64
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import paho.mqtt.client as mqtt
|
|
from flask import Flask, render_template, jsonify, request, Response
|
|
|
|
# Import Config, Manager, and NEW Inference Worker
|
|
from config import Config
|
|
from manager import CameraManager
|
|
from inference import InferenceWorker
|
|
|
|
# --- Logging Setup ---
|
|
logging.basicConfig(
|
|
level=Config.LOG_LEVEL,
|
|
format='%(asctime)s [%(levelname)s] %(message)s',
|
|
handlers=[logging.StreamHandler(sys.stdout)]
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
app = Flask(__name__)
|
|
|
|
# --- Initialize Components ---
|
|
camera_manager = CameraManager()
|
|
inference_worker = InferenceWorker() # <--- NEW
|
|
inference_worker.start() # <--- Start the background thread
|
|
|
|
# --- MQTT Setup ---
|
|
mqtt_client = mqtt.Client()
|
|
if Config.MQTT_USERNAME and Config.MQTT_PASSWORD:
|
|
mqtt_client.username_pw_set(Config.MQTT_USERNAME, Config.MQTT_PASSWORD)
|
|
|
|
try:
|
|
mqtt_client.connect(Config.MQTT_BROKER, Config.MQTT_PORT, 60)
|
|
mqtt_client.loop_start() # START THE LOOP HERE
|
|
logger.info(f"Connected to MQTT Broker at {Config.MQTT_BROKER}:{Config.MQTT_PORT}")
|
|
except Exception as e:
|
|
logger.error(f"Failed to connect to MQTT Broker: {e}")
|
|
|
|
# --- Helper Functions (UI Only) ---
|
|
def crop_image_for_ui(image, roi_list, scaleX, scaleY):
|
|
"""Helper for the /crop endpoint (UI preview only)."""
|
|
cropped_images = []
|
|
for roi in roi_list:
|
|
try:
|
|
x = int(roi['x'] * scaleX)
|
|
y = int(roi['y'] * scaleY)
|
|
width = int(roi['width'] * scaleX)
|
|
height = int(roi['height'] * scaleY)
|
|
cropped = image[y:y + height, x:x + width]
|
|
if cropped.size > 0:
|
|
cropped_images.append(cropped)
|
|
except Exception:
|
|
pass
|
|
return cropped_images
|
|
|
|
def publish_detected_number(camera_id, detected_number):
|
|
"""Publish result to MQTT."""
|
|
topic = f"{Config.MQTT_TOPIC}/{camera_id}"
|
|
payload = json.dumps({"value": detected_number})
|
|
try:
|
|
mqtt_client.publish(topic, payload)
|
|
logger.info(f"Published to {topic}: {detected_number}")
|
|
except Exception as e:
|
|
logger.error(f"MQTT Publish failed: {e}")
|
|
|
|
# --- Main Processing Loop (Refactored) ---
|
|
# Add this global dictionary at the top of app.py (near other globals)
|
|
last_processed_time = {}
|
|
|
|
# Update process_all_cameras function
|
|
def process_all_cameras():
|
|
"""
|
|
Revised Loop with Rate Limiting
|
|
"""
|
|
# Configurable interval (seconds)
|
|
DETECTION_INTERVAL = 10
|
|
|
|
while True:
|
|
try:
|
|
# --- Part 1: Process Results ---
|
|
while True:
|
|
result = inference_worker.get_result()
|
|
if not result:
|
|
break
|
|
|
|
cam_id = result['camera_id']
|
|
val = result['value']
|
|
|
|
camera_manager.results[cam_id] = val
|
|
publish_detected_number(cam_id, val)
|
|
|
|
# --- Part 2: Feed Frames ---
|
|
camera_manager.load_roi_config()
|
|
|
|
for camera_id, camera_data in camera_manager.cameras.items():
|
|
if not camera_data.get("active", True):
|
|
continue
|
|
|
|
# RATE LIMIT CHECK
|
|
current_time = time.time()
|
|
last_time = last_processed_time.get(camera_id, 0)
|
|
|
|
if current_time - last_time < DETECTION_INTERVAL:
|
|
continue # Skip this camera, it's too soon
|
|
|
|
stream = camera_data.get("stream")
|
|
if not stream: continue
|
|
|
|
# Warmup Check
|
|
if (current_time - stream.start_time) < 5:
|
|
continue
|
|
|
|
frame = stream.read()
|
|
|
|
if frame is None:
|
|
continue
|
|
|
|
if np.std(frame) < 10:
|
|
continue
|
|
|
|
roi_list = camera_manager.rois.get(camera_id, [])
|
|
if not roi_list:
|
|
continue
|
|
|
|
# SEND TO WORKER
|
|
inference_worker.add_task(camera_id, roi_list, frame)
|
|
|
|
# Update last processed time
|
|
last_processed_time[camera_id] = current_time
|
|
|
|
# Sleep briefly to prevent CPU spinning, but keep it responsive for results
|
|
time.sleep(0.1)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Global process loop error: {e}")
|
|
traceback.print_exc()
|
|
time.sleep(5)
|
|
|
|
|
|
# --- Flask Routes (Unchanged logic, just imports) ---
|
|
|
|
@app.route('/')
|
|
def index():
|
|
return render_template('index.html')
|
|
|
|
@app.route('/cameras', methods=['GET'])
|
|
def get_cameras():
|
|
return jsonify(camera_manager.get_camera_list())
|
|
|
|
@app.route('/video/<camera_id>')
|
|
def video_feed(camera_id):
|
|
def generate():
|
|
while True:
|
|
frame = camera_manager.get_frame(camera_id)
|
|
if frame is not None:
|
|
ret, jpeg = cv2.imencode('.jpg', frame)
|
|
if ret:
|
|
yield (b'--frame\r\n' b'Content-Type: image/jpeg\r\n\r\n' + jpeg.tobytes() + b'\r\n\r\n')
|
|
else:
|
|
time.sleep(0.1)
|
|
return Response(generate(), mimetype='multipart/x-mixed-replace; boundary=frame')
|
|
|
|
@app.route('/snapshot/<camera_id>')
|
|
def snapshot(camera_id):
|
|
frame = camera_manager.get_frame(camera_id)
|
|
if frame is not None:
|
|
ret, jpeg = cv2.imencode('.jpg', frame)
|
|
if ret:
|
|
return Response(jpeg.tobytes(), mimetype='image/jpeg')
|
|
return 'No frame available', 404
|
|
|
|
@app.route('/rois/<camera_id>', methods=['GET'])
|
|
def get_rois(camera_id):
|
|
# ... (Same logic as Step 3, just ensure it uses camera_manager) ...
|
|
try:
|
|
camera_manager.load_roi_config()
|
|
all_rois = camera_manager.rois
|
|
img_width = request.args.get("img_width", type=float)
|
|
img_height = request.args.get("img_height", type=float)
|
|
|
|
if not img_width or not img_height:
|
|
return jsonify(all_rois.get(camera_id, []))
|
|
|
|
cam = camera_manager.cameras.get(camera_id)
|
|
if cam and cam.get("stream"):
|
|
real_w = cam["stream"].width or cam["width"]
|
|
real_h = cam["stream"].height or cam["height"]
|
|
else:
|
|
return jsonify({"error": "Camera not ready"}), 500
|
|
|
|
scaleX = img_width / real_w
|
|
scaleY = img_height / real_h
|
|
scaled_rois = []
|
|
for roi in all_rois.get(camera_id, []):
|
|
scaled_rois.append({
|
|
"id": roi["id"],
|
|
"x": int(round(roi["x"] * scaleX)),
|
|
"y": int(round(roi["y"] * scaleY)),
|
|
"width": int(round(roi["width"] * scaleX)),
|
|
"height": int(round(roi["height"] * scaleY)),
|
|
"angle": roi["angle"]
|
|
})
|
|
return jsonify(scaled_rois)
|
|
except Exception as e:
|
|
return jsonify({"error": str(e)}), 500
|
|
|
|
@app.route("/save_rois", methods=["POST"])
|
|
def save_rois_api():
|
|
# ... (Same logic as Step 3) ...
|
|
data = request.json
|
|
camera_id = data.get("camera_id")
|
|
new_rois = data.get("rois")
|
|
img_width = data.get("img_width")
|
|
img_height = data.get("img_height")
|
|
|
|
if not camera_id or new_rois is None: return jsonify({"success": False})
|
|
|
|
cam = camera_manager.cameras.get(camera_id)
|
|
if not cam: return jsonify({"success": False})
|
|
|
|
stream = cam.get("stream")
|
|
real_w = stream.width if stream and stream.width else cam["width"]
|
|
real_h = stream.height if stream and stream.height else cam["height"]
|
|
|
|
scaleX = real_w / img_width if img_width else 1
|
|
scaleY = real_h / img_height if img_height else 1
|
|
|
|
scaled_rois = []
|
|
for roi in new_rois:
|
|
scaled_rois.append({
|
|
"id": roi["id"],
|
|
"x": int(round(roi["x"] * scaleX)),
|
|
"y": int(round(roi["y"] * scaleY)),
|
|
"width": int(round(roi["width"] * scaleX)),
|
|
"height": int(round(roi["height"] * scaleY)),
|
|
"angle": roi["angle"]
|
|
})
|
|
camera_manager.rois[camera_id] = scaled_rois
|
|
return jsonify(camera_manager.save_roi_config())
|
|
|
|
@app.route('/crop', methods=['POST'])
|
|
def crop():
|
|
# Helper for UI
|
|
data = request.json
|
|
camera_id = data.get('camera_id')
|
|
scaleX = data.get('scaleX', 1)
|
|
scaleY = data.get('scaleY', 1)
|
|
|
|
frame = camera_manager.get_frame(camera_id)
|
|
if frame is None: return jsonify({'error': 'No frame'}), 500
|
|
|
|
roi_list = camera_manager.rois.get(camera_id, [])
|
|
# Use the local UI helper function
|
|
cropped_images = crop_image_for_ui(frame, roi_list, scaleX, scaleY)
|
|
|
|
cropped_base64_list = []
|
|
for cropped_img in cropped_images:
|
|
ret, buffer = cv2.imencode('.jpg', cropped_img)
|
|
if ret:
|
|
cropped_base64_list.append(base64.b64encode(buffer).decode('utf-8'))
|
|
return jsonify({'cropped_images': cropped_base64_list})
|
|
|
|
@app.route('/detect_digits', methods=['POST'])
|
|
def detect_digits():
|
|
"""Manual trigger: Runs inference immediately and returns result."""
|
|
data = request.json
|
|
camera_id = data.get('camera_id')
|
|
|
|
if not camera_id:
|
|
return jsonify({'error': 'Invalid camera ID'}), 400
|
|
|
|
# 1. Get Frame
|
|
frame = camera_manager.get_frame(camera_id)
|
|
if frame is None:
|
|
return jsonify({'error': 'Failed to capture image'}), 500
|
|
|
|
# 2. Get ROIs
|
|
roi_list = camera_manager.rois.get(camera_id, [])
|
|
if not roi_list:
|
|
return jsonify({'error': 'No ROIs defined'}), 400
|
|
|
|
# 3. Crop (Using the UI helper is fine here)
|
|
cropped_images = crop_image_for_ui(frame, roi_list, scaleX=1, scaleY=1)
|
|
if not cropped_images:
|
|
return jsonify({'error': 'Failed to crop ROIs'}), 500
|
|
|
|
# 4. Run Inference Synchronously
|
|
# Note: We access the worker directly.
|
|
# Thread safety: 'predict_batch' uses 'self.interpreter'.
|
|
# If the background thread is also using it, TFLite might complain or crash.
|
|
# PROPER FIX: Pause the worker or use a Lock.
|
|
|
|
# Since adding a Lock is complex now, a simple hack is to just add it to the queue
|
|
# and WAIT for the result? No, that's hard to correlate.
|
|
|
|
# SAFE APPROACH: Use a Lock in InferenceWorker.
|
|
# For now, let's assume TFLite is robust enough or race conditions are rare for manual clicks.
|
|
# CALL THE PUBLIC METHOD:
|
|
try:
|
|
detected_digits = inference_worker.predict_batch(cropped_images)
|
|
|
|
valid_digits = [d for d in detected_digits if d.isdigit()]
|
|
|
|
if not valid_digits:
|
|
return jsonify({'error': 'No valid digits detected', 'raw': detected_digits}), 500
|
|
|
|
final_number = int("".join(valid_digits))
|
|
|
|
# Publish and Update State
|
|
publish_detected_number(camera_id, final_number)
|
|
camera_manager.results[camera_id] = final_number
|
|
|
|
logger.info(f"Manual detection for {camera_id}: {final_number}")
|
|
|
|
return jsonify({
|
|
'detected_digits': valid_digits,
|
|
'final_number': final_number
|
|
})
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error during manual detection: {e}")
|
|
return jsonify({'error': str(e)}), 500
|
|
|
|
|
|
@app.route('/update_camera_config', methods=['POST'])
|
|
def update_camera_config():
|
|
data = request.json
|
|
success = camera_manager.update_camera_flip(data.get("camera_id"), data.get("flip_type"))
|
|
return jsonify({"success": success})
|
|
|
|
# --- Main ---
|
|
if __name__ == '__main__':
|
|
# Threading:
|
|
# 1. Video Threads (in Manager)
|
|
# 2. Inference Thread (in Worker)
|
|
# 3. Main Loop (process_all_cameras - handles feeding)
|
|
|
|
t = threading.Thread(target=process_all_cameras, daemon=True)
|
|
t.start()
|
|
|
|
logger.info("Starting Flask Server...")
|
|
app.run(host='0.0.0.0', port=5000, threaded=True)
|