import os import sys import json import base64 import subprocess from io import BytesIO import numpy as np from PIL import Image, ImageDraw import cv2 import albumentations as A import ctypes import time import socket import threading from fastapi import FastAPI, HTTPException, BackgroundTasks from fastapi.staticfiles import StaticFiles from fastapi.responses import FileResponse, JSONResponse from pydantic import BaseModel from inference import FruitVegetableDetector, CLASSES app = FastAPI(title="POS Object Detection API", description="Fruit & Vegetable checkout detection microservice") # Initialize detector detector = FruitVegetableDetector() # Paths BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) STATUS_FILE = os.path.join(BASE_DIR, "training_status.json") # Data Models class DetectionRequest(BaseModel): image: str # Base64 encoded image or Data URL class TrainRequest(BaseModel): epochs: int = 15 batch_size: int = 8 lr: float = 0.001 class CaptureRequest(BaseModel): image: str class_name: str class SaveArticlesRequest(BaseModel): articles: dict class ScanRequest(BaseModel): class_name: str # Serve frontend static files static_path = os.path.join(os.path.dirname(__file__), "static") os.makedirs(static_path, exist_ok=True) # Helper function to read training status def get_training_status(): if not os.path.exists(STATUS_FILE): return { "status": "idle", "epoch": 0, "total_epochs": 0, "train_loss": [], "val_loss": [], "progress": 0.0, "message": "No training history found." } try: with open(STATUS_FILE, "r") as f: return json.load(f) except Exception: return {"status": "error", "message": "Failed to read status file."} # Helper to write status def write_training_status(status_dict): with open(STATUS_FILE, "w") as f: json.dump(status_dict, f, indent=2) @app.post("/api/detect") async def detect_objects(payload: DetectionRequest): """ Accepts base64 encoded frame, runs ONNX inference, and returns detections. """ try: data_url = payload.image if "," in data_url: data_url = data_url.split(",")[1] img_bytes = base64.b64decode(data_url) img = Image.open(BytesIO(img_bytes)).convert("RGB") img_np = np.array(img) predictions = detector.detect(img_np) return { "status": "success", "predictions": predictions } except Exception as e: return JSONResponse( status_code=400, content={"status": "error", "message": f"Inference failed: {str(e)}"} ) def bg_generate_dataset(): try: from data_generator import build_dataset build_dataset(base_dir=os.path.join(BASE_DIR, "dataset"), train_count=200, val_count=50) write_training_status({ "status": "idle", "epoch": 0, "total_epochs": 0, "train_loss": [], "val_loss": [], "progress": 100.0, "message": "Dataset generation completed successfully!" }) except Exception as e: write_training_status({ "status": "failed", "epoch": 0, "total_epochs": 0, "train_loss": [], "val_loss": [], "progress": 0.0, "message": f"Dataset generation failed: {str(e)}" }) @app.post("/api/generate_dataset") async def trigger_dataset_generation(background_tasks: BackgroundTasks): """ Triggers dataset generation asynchronously. """ status = get_training_status() if status.get("status") in ["training", "generating_data"]: raise HTTPException(status_code=400, detail="An operation is already in progress.") write_training_status({ "status": "generating_data", "epoch": 0, "total_epochs": 0, "train_loss": [], "val_loss": [], "progress": 10.0, "message": "Starting synthetic data generation process..." }) background_tasks.add_task(bg_generate_dataset) return {"status": "success", "message": "Dataset generation started in the background."} @app.post("/api/train") async def trigger_training(req: TrainRequest): """ Triggers model training. """ status = get_training_status() if status.get("status") in ["training", "generating_data"]: raise HTTPException(status_code=400, detail="Training or data generation is already in progress.") write_training_status({ "status": "training", "epoch": 0, "total_epochs": req.epochs, "train_loss": [], "val_loss": [], "progress": 5.0, "message": "Initializing training background process..." }) # Launch training process in background cmd = [ sys.executable, os.path.join(BASE_DIR, "train.py"), "--epochs", str(req.epochs), "--batch-size", str(req.batch_size), "--lr", str(req.lr) ] subprocess.Popen(cmd, cwd=BASE_DIR) return {"status": "success", "message": f"Training initiated for {req.epochs} epochs."} @app.get("/api/train_status") async def get_status(): """ Returns current training status. """ # Check if background processes are actually running # If status says training/generating but no process matches, we can update status to failed/idle. status = get_training_status() return status @app.post("/api/export") async def trigger_export(): """ Exports trained PyTorch weights to ONNX format and reloads the detector. """ pytorch_path = os.path.join(BASE_DIR, "models", "model.pt") if not os.path.exists(pytorch_path): raise HTTPException(status_code=404, detail="No trained PyTorch weights found at models/model.pt. Train a model first.") try: # Run export script synchronously since it is fast (few seconds) cmd = [sys.executable, os.path.join(BASE_DIR, "export.py")] result = subprocess.run(cmd, cwd=BASE_DIR, capture_output=True, text=True, check=True) # Reload the detector model detector.load_model() return { "status": "success", "message": "Model successfully exported to ONNX and loaded into memory.", "log": result.stdout } except Exception as e: return JSONResponse( status_code=500, content={"status": "error", "message": f"ONNX Export failed: {str(e)}"} ) @app.get("/api/augmented_preview") async def get_augmented_preview(): """ Generates a sample preview image from the generator, applies augmentations, draws bounding boxes on both, and returns them side-by-side. """ try: from data_generator import generate_sample # Generate a raw sample image and its YOLO annotations img, annotations = generate_sample(width=320, height=320) img_np = np.array(img) # Setup the same augmentation pipeline used in training transform = A.Compose([ A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.15, rotate_limit=30, p=1.0), A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=1.0), A.GaussNoise(p=1.0) ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['category_ids'])) # Parse annotations to box absolute coords boxes = [] classes = [] for ann in annotations: parts = ann.strip().split() if len(parts) == 5: class_id = int(parts[0]) xc, yc, w, h = map(float, parts[1:]) xmin = (xc - w/2) * 320 ymin = (yc - h/2) * 320 xmax = (xc + w/2) * 320 ymax = (yc + h/2) * 320 boxes.append([xmin, ymin, xmax, ymax]) classes.append(class_id) boxes = np.array(boxes, dtype=np.float32).reshape(-1, 4) # Create augmented image augmented_np = img_np.copy() aug_boxes = boxes.copy() if len(boxes) > 0: try: augmented = transform(image=img_np, bboxes=boxes, category_ids=classes) augmented_np = augmented["image"] aug_boxes = np.array(augmented["bboxes"], dtype=np.float32).reshape(-1, 4) except Exception: pass # Draw bounding boxes on original image orig_draw = img.copy() draw_tool = ImageDraw.Draw(orig_draw) for idx, box in enumerate(boxes): draw_tool.rectangle(box.tolist(), outline=(255, 0, 0), width=3) draw_tool.text((box[0]+2, box[1]+2), CLASSES[classes[idx]], fill=(255, 0, 0)) # Draw bounding boxes on augmented image aug_img = Image.fromarray(augmented_np) draw_tool_aug = ImageDraw.Draw(aug_img) for idx, box in enumerate(aug_boxes): draw_tool_aug.rectangle(box.tolist(), outline=(0, 255, 0), width=3) draw_tool_aug.text((box[0]+2, box[1]+2), CLASSES[classes[idx]], fill=(0, 255, 0)) # Create side-by-side preview canvas preview = Image.new("RGB", (640, 320)) preview.paste(orig_draw, (0, 0)) preview.paste(aug_img, (320, 0)) # Buffer and return as base64 buffered = BytesIO() preview.save(buffered, format="JPEG") preview_base64 = base64.b64encode(buffered.getvalue()).decode() return { "status": "success", "image": f"data:image/jpeg;base64,{preview_base64}" } except Exception as e: raise HTTPException(status_code=500, detail=f"Failed to generate preview: {str(e)}") @app.get("/api/simulated_frame") async def get_simulated_frame(): """ Generates a raw synthetic scale frame (without annotations) for client-side POS testing. """ try: from data_generator import generate_sample # Generate raw sample image (640x480) img, _ = generate_sample(width=640, height=480) buffered = BytesIO() img.save(buffered, format="JPEG") img_base64 = base64.b64encode(buffered.getvalue()).decode() return { "status": "success", "image": f"data:image/jpeg;base64,{img_base64}" } except Exception as e: raise HTTPException(status_code=500, detail=f"Failed to generate simulated frame: {str(e)}") @app.post("/api/capture_training_image") async def capture_training_image(payload: CaptureRequest): """ Saves a webcam snapshot and generates a centered YOLO format bounding box label. """ try: class_name = payload.class_name if class_name not in CLASSES: raise HTTPException(status_code=400, detail=f"Ungültige Produktklasse: {class_name}") class_idx = CLASSES.index(class_name) # Decode base64 data_url = payload.image if "," in data_url: data_url = data_url.split(",")[1] img_bytes = base64.b64decode(data_url) img = Image.open(BytesIO(img_bytes)).convert("RGB") # Save image with unique filename import time timestamp = int(time.time() * 1000) filename = f"user_capture_{timestamp}.jpg" images_dir = os.path.join(BASE_DIR, "dataset", "images", "train") labels_dir = os.path.join(BASE_DIR, "dataset", "labels", "train") os.makedirs(images_dir, exist_ok=True) os.makedirs(labels_dir, exist_ok=True) img.save(os.path.join(images_dir, filename), quality=95) # Write YOLO label centered at scale # Bounding box: class_idx, x_center=0.5, y_center=0.5, width=0.6, height=0.6 label_filename = f"user_capture_{timestamp}.txt" with open(os.path.join(labels_dir, label_filename), "w") as f: f.write(f"{class_idx} 0.500000 0.500000 0.650000 0.650000\n") # Count total user captured images user_images = [f for f in os.listdir(images_dir) if f.startswith("user_capture_")] count = len(user_images) return { "status": "success", "message": f"Bild erfolgreich als Training für {class_name} erfasst!", "count": count } except Exception as e: raise HTTPException(status_code=500, detail=f"Fehler bei Bild-Erfassung: {str(e)}") @app.get("/api/capture_count") async def get_capture_count(): """ Returns the number of user captured training images. """ images_dir = os.path.join(BASE_DIR, "dataset", "images", "train") if not os.path.exists(images_dir): return {"count": 0} user_images = [f for f in os.listdir(images_dir) if f.startswith("user_capture_")] return {"count": len(user_images)} # --- BARCODE SCANNER EMULATION (Windows Keyboard Hook) --- KEYEVENTF_KEYUP = 0x0002 VK_RETURN = 0x0D ARTICLES_FILE = os.path.join(BASE_DIR, "articles.json") def simulate_keyboard_type(text: str): """ Types the EAN code and presses Enter globally on Windows. Uses hardware scan codes mapped from virtual key codes for highest compatibility. """ try: # Give a small 200ms delay to allow cashier window to gain focus if clicked manually time.sleep(0.2) KEYEVENTF_KEYUP = 0x0002 MAPVK_VK_TO_VSC = 0 for char in text: vk = None if '0' <= char <= '9': vk = 0x30 + int(char) elif 'A' <= char <= 'Z': vk = ord(char) elif 'a' <= char <= 'z': vk = ord(char.upper()) if vk is not None: # Map to hardware scan code scan_code = ctypes.windll.user32.MapVirtualKeyW(vk, MAPVK_VK_TO_VSC) # Key Down ctypes.windll.user32.keybd_event(vk, scan_code, 0, 0) time.sleep(0.001) # Key Up ctypes.windll.user32.keybd_event(vk, scan_code, KEYEVENTF_KEYUP, 0) time.sleep(0.001) # Send VK_RETURN (Enter) with scan code vk_ret = 0x0D scan_ret = ctypes.windll.user32.MapVirtualKeyW(vk_ret, MAPVK_VK_TO_VSC) ctypes.windll.user32.keybd_event(vk_ret, scan_ret, 0, 0) time.sleep(0.001) ctypes.windll.user32.keybd_event(vk_ret, scan_ret, KEYEVENTF_KEYUP, 0) return True except Exception as e: print(f"Failed to type keyboard input: {e}", file=sys.stderr) return False @app.get("/api/articles") async def get_articles(): if not os.path.exists(ARTICLES_FILE): return {} try: with open(ARTICLES_FILE, "r") as f: return json.load(f) except Exception as e: raise HTTPException(status_code=500, detail=f"Fehler beim Laden der Artikeldatenbank: {str(e)}") @app.post("/api/articles") async def save_articles(payload: SaveArticlesRequest): try: with open(ARTICLES_FILE, "w") as f: json.dump(payload.articles, f, indent=2) from data_generator import reload_classes reload_classes() return {"status": "success", "message": "Artikeldatenbank erfolgreich gespeichert!"} except Exception as e: raise HTTPException(status_code=500, detail=f"Fehler beim Speichern der Artikeldatenbank: {str(e)}") # --- TCP SOCKET CLIENT SENDER --- TCP_PORT = 9000 def send_ean_via_tcp(ean_code: str): """ Sends the EAN code to the cash register over TCP on port 9000. Since the cash register is the server (listening on port 9000), we connect as a client, transmit the barcode, and close the connection. """ message = (ean_code + "\r\n").encode('utf-8') print(f"TCP Client Mode: Connecting to cash register at 127.0.0.1:{TCP_PORT}...", flush=True) try: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.settimeout(0.5) # Fast timeout s.connect(('127.0.0.1', TCP_PORT)) s.sendall(message) print(f"TCP Client: Successfully sent EAN {ean_code} to cash register.", flush=True) return True except Exception as e: print(f"TCP Client: Connection to cash register failed (is register listening on port {TCP_PORT}?): {e}", flush=True) return False @app.post("/api/scan_ean") async def scan_ean_endpoint(payload: ScanRequest): try: class_name = payload.class_name if not os.path.exists(ARTICLES_FILE): raise HTTPException(status_code=404, detail="Artikeldatenbank existiert nicht.") with open(ARTICLES_FILE, "r") as f: articles = json.load(f) if class_name not in articles: raise HTTPException(status_code=404, detail=f"Klasse {class_name} nicht in Artikeldatenbank.") ean_code = articles[class_name].get("ean", "") if not ean_code: raise HTTPException(status_code=400, detail=f"Kein EAN-Code für {class_name} hinterlegt.") # 1. Send via TCP (Direct socket integration) tcp_sent = send_ean_via_tcp(ean_code) # 2. Parallel keyboard emulation fallback on Windows keyboard_sent = simulate_keyboard_type(ean_code) if tcp_sent or keyboard_sent: methods = [] if tcp_sent: methods.append("TCP") if keyboard_sent: methods.append("Tastatur") method_desc = " + ".join(methods) return { "status": "success", "class": class_name, "name": articles[class_name].get("name", class_name), "ean": ean_code, "message": f"Barcode {ean_code} gesendet via {method_desc}!" } else: raise HTTPException(status_code=500, detail="Übertragung fehlgeschlagen (TCP und Tastaturemulation gescheitert).") except Exception as e: raise HTTPException(status_code=500, detail=str(e)) # Mount static files (must be after api definitions) app.mount("/", StaticFiles(directory=static_path, html=True), name="static") @app.get("/") async def root(): return FileResponse(os.path.join(static_path, "index.html"))