524 lines
18 KiB
Python
524 lines
18 KiB
Python
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"))
|