add snapshots & counter

alpr-service/main.py

@@ -5,16 +5,18 @@ import os
 import time
 import threading
 import re
+from datetime import datetime
 from queue import Queue
-from flask import Flask, Response
+from flask import Flask, Response, jsonify
 from flask_cors import CORS
 from ultralytics import YOLO
 
 # Configuration
 BACKEND_URL = os.environ.get('BACKEND_URL', 'http://localhost:3000')
 CAMERA_ID = 0
-PROCESS_INTERVAL = 1.5  # Más reactivo
+PROCESS_INTERVAL = 1.5
 MODEL_PATH = 'best.pt'
+DATASET_DIR = '/app/dataset'  # Carpeta para guardar capturas
 
 app = Flask(__name__)
 CORS(app)
@@ -25,9 +27,32 @@ frame_lock = threading.Lock()
 latest_detections = []
 detection_lock = threading.Lock()
 
-# Cola para procesamiento OCR asíncrono
+# Cola para procesamiento OCR asíncrono (ahora incluye frame completo)
 ocr_queue = Queue(maxsize=5)
 
+# Crear carpeta de dataset si no existe
+os.makedirs(DATASET_DIR, exist_ok=True)
+print(f"📁 Dataset directory: {DATASET_DIR}")
+
+def save_plate_capture(plate_number, plate_img, full_frame):
+    """Guarda la captura de la patente para el dataset"""
+    try:
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        
+        # Guardar imagen recortada de la patente
+        plate_filename = f"{DATASET_DIR}/{plate_number}_{timestamp}_plate.jpg"
+        cv2.imwrite(plate_filename, plate_img, [cv2.IMWRITE_JPEG_QUALITY, 95])
+        
+        # Guardar frame completo con contexto
+        frame_filename = f"{DATASET_DIR}/{plate_number}_{timestamp}_full.jpg"
+        cv2.imwrite(frame_filename, full_frame, [cv2.IMWRITE_JPEG_QUALITY, 90])
+        
+        print(f"📸 Saved to dataset: {plate_number}")
+        return True
+    except Exception as e:
+        print(f"❌ Error saving capture: {e}")
+        return False
+
 def send_plate(plate_number):
     """Envía la patente detectada al backend"""
     try:
@@ -37,22 +62,21 @@ def send_plate(plate_number):
     except Exception as e:
         print(f"✗ Error sending plate: {e}")
 
-def validate_and_send(text):
-    """Valida formato chileno y envía"""
+def validate_plate(text):
+    """Valida formato chileno"""
     # Formato nuevo: XXXX-00 | Formato antiguo: XX-0000
-    if re.match(r'^[A-Z]{4}\d{2}$', text) or re.match(r'^[A-Z]{2}\d{4}$', text):
-        send_plate(text)
-        return True
-    return False
+    return bool(re.match(r'^[A-Z]{4}\d{2}$', text) or re.match(r'^[A-Z]{2}\d{4}$', text))
 
 def ocr_worker(reader):
     """Hilo dedicado para OCR - no bloquea el stream"""
     while True:
         try:
-            plate_img = ocr_queue.get(timeout=1)
-            if plate_img is None:
+            data = ocr_queue.get(timeout=1)
+            if data is None:
                 continue
             
+            plate_img, full_frame = data
+            
             # Preprocesamiento para mejor OCR
             gray = cv2.cvtColor(plate_img, cv2.COLOR_BGR2GRAY)
             
@@ -60,8 +84,11 @@ def ocr_worker(reader):
                                           allowlist='ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789')
             for text in ocr_results:
                 clean_text = ''.join(e for e in text if e.isalnum()).upper()
-                if len(clean_text) >= 6:
-                    validate_and_send(clean_text)
+                if len(clean_text) >= 6 and validate_plate(clean_text):
+                    # Enviar al backend
+                    send_plate(clean_text)
+                    # Guardar captura para dataset
+                    save_plate_capture(clean_text, plate_img, full_frame)
         except:
             pass
 
@@ -96,10 +123,9 @@ def camera_loop():
     print("✅ System ready!")
     
     last_process_time = 0
-    frame_count = 0
     
     while True:
-        # Captura eficiente - solo 2 grabs
+        # Captura eficiente
         cap.grab()
         cap.grab()
         ret, frame = cap.retrieve()
@@ -108,14 +134,13 @@ def camera_loop():
             time.sleep(0.01)
             continue
         
-        frame_count += 1
         current_time = time.time()
         
         # Procesar ALPR cada PROCESS_INTERVAL segundos
         if current_time - last_process_time > PROCESS_INTERVAL:
             last_process_time = current_time
             
-            # YOLO detection - usar imgsz pequeño para velocidad
+            # YOLO detection
             results = model(frame, verbose=False, imgsz=320, conf=0.5)
             
             new_detections = []
@@ -125,15 +150,16 @@ def camera_loop():
                     conf = float(box.conf[0])
                     new_detections.append((x1, y1, x2, y2, conf))
                     
-                    # Extraer imagen de placa y enviar a cola OCR
+                    # Extraer imagen de placa
                     plate_img = frame[y1:y2, x1:x2].copy()
                     if plate_img.size > 0 and not ocr_queue.full():
-                        ocr_queue.put(plate_img)
+                        # Enviar placa Y frame completo para dataset
+                        ocr_queue.put((plate_img, frame.copy()))
             
             with detection_lock:
                 latest_detections = new_detections
         
-        # Actualizar frame para streaming (sin bloquear)
+        # Actualizar frame para streaming
         display_frame = frame
         with detection_lock:
             for (x1, y1, x2, y2, conf) in latest_detections:
@@ -148,7 +174,7 @@ def generate():
     """Generador para streaming MJPEG"""
     global outputFrame
     while True:
-        time.sleep(0.033)  # ~30 FPS para el stream
+        time.sleep(0.033)
         with frame_lock:
             if outputFrame is None:
                 continue
@@ -163,6 +189,16 @@ def video_feed():
 def health():
     return {"status": "ok", "service": "alpr"}
 
+@app.route("/dataset/count")
+def dataset_count():
+    """Endpoint para ver cuántas capturas hay en el dataset"""
+    try:
+        files = os.listdir(DATASET_DIR)
+        plates = len([f for f in files if f.endswith('_plate.jpg')])
+        return {"plates_captured": plates, "total_files": len(files)}
+    except:
+        return {"plates_captured": 0, "total_files": 0}
+
 if __name__ == "__main__":
     t = threading.Thread(target=camera_loop, daemon=True)
     t.start()