915 lines
33 KiB
Python
915 lines
33 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()
|
|
|
|
# Global variable to track active training process
|
|
current_training_process = None
|
|
|
|
import re
|
|
|
|
# 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
|
|
bbox: dict = None # Optional client-side bounding box dictionary
|
|
|
|
class SaveArticlesRequest(BaseModel):
|
|
articles: dict
|
|
|
|
class ScanRequest(BaseModel):
|
|
class_name: str
|
|
|
|
class CalibrateRequest(BaseModel):
|
|
image: str
|
|
|
|
class DeleteCaptureRequest(BaseModel):
|
|
filename: 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)
|
|
|
|
def is_safe_filename(filename: str) -> bool:
|
|
return re.match(r"^user_capture_\d+\.jpg$", filename) is not None
|
|
|
|
def calculate_auto_crop_bbox(image_rgb: np.ndarray, background_path: str, threshold: int = 25) -> list:
|
|
"""
|
|
Compares the current image with the saved background.
|
|
Returns normalized YOLO bounding box [x_center, y_center, width, height] if found, else None.
|
|
"""
|
|
if not os.path.exists(background_path):
|
|
return None
|
|
try:
|
|
# Load background image
|
|
bg_img = cv2.imread(background_path)
|
|
if bg_img is None:
|
|
return None
|
|
# Convert to grayscale
|
|
bg_gray = cv2.cvtColor(bg_img, cv2.COLOR_BGR2GRAY)
|
|
|
|
# Ensure they have the same dimensions
|
|
h_img, w_img, _ = image_rgb.shape
|
|
if bg_gray.shape[0] != h_img or bg_gray.shape[1] != w_img:
|
|
bg_gray = cv2.resize(bg_gray, (w_img, h_img))
|
|
|
|
# Convert current image to grayscale (assuming input is RGB)
|
|
curr_gray = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2GRAY)
|
|
|
|
# Calculate absolute difference
|
|
diff = cv2.absdiff(bg_gray, curr_gray)
|
|
|
|
# Apply thresholding
|
|
_, thresh = cv2.threshold(diff, threshold, 255, cv2.THRESH_BINARY)
|
|
|
|
# Morphological operations to remove noise and fill holes
|
|
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
|
|
thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)
|
|
thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
|
|
|
|
# Find contours
|
|
contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
if not contours:
|
|
return None
|
|
|
|
# Filter contours by size and find the largest
|
|
large_contours = [c for c in contours if cv2.contourArea(c) > 150]
|
|
if not large_contours:
|
|
return None
|
|
|
|
largest_contour = max(large_contours, key=cv2.contourArea)
|
|
x, y, w, h = cv2.boundingRect(largest_contour)
|
|
|
|
# Add padding around the bounding box
|
|
padding = 15
|
|
x_min = max(0, x - padding)
|
|
y_min = max(0, y - padding)
|
|
x_max = min(w_img, x + w + padding)
|
|
y_max = min(h_img, y + h + padding)
|
|
|
|
box_w = x_max - x_min
|
|
box_h = y_max - y_min
|
|
|
|
# Calculate YOLO format values
|
|
x_center = (x_min + box_w / 2.0) / w_img
|
|
y_center = (y_min + box_h / 2.0) / h_img
|
|
norm_w = box_w / w_img
|
|
norm_h = box_h / h_img
|
|
|
|
return [x_center, y_center, norm_w, norm_h]
|
|
except Exception as e:
|
|
print(f"Error in calculate_auto_crop_bbox: {e}")
|
|
return None
|
|
|
|
|
|
@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)
|
|
|
|
# Check if empty background calibration exists and compare
|
|
bg_path = os.path.join(BASE_DIR, "dataset", "background.jpg")
|
|
if os.path.exists(bg_path):
|
|
# Use a slightly lower threshold (20) to ensure we capture any real item placed
|
|
bbox = calculate_auto_crop_bbox(img_np, bg_path, threshold=20)
|
|
if bbox is None:
|
|
# Scale is empty or no significant object placed
|
|
return {
|
|
"status": "success",
|
|
"predictions": []
|
|
}
|
|
|
|
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.
|
|
"""
|
|
global current_training_process
|
|
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 with reduced CPU priority
|
|
# so the FastAPI server stays responsive during CPU-heavy training
|
|
cmd = [
|
|
sys.executable,
|
|
os.path.join(BASE_DIR, "train.py"),
|
|
"--epochs", str(req.epochs),
|
|
"--batch-size", str(req.batch_size),
|
|
"--lr", str(req.lr)
|
|
]
|
|
|
|
env = os.environ.copy()
|
|
env["PYTHONUTF8"] = "1"
|
|
proc = subprocess.Popen(cmd, cwd=BASE_DIR, env=env)
|
|
current_training_process = proc
|
|
|
|
# On Windows: lower the training process priority so the web server
|
|
# remains responsive during CPU-intensive training
|
|
if sys.platform == "win32":
|
|
try:
|
|
import ctypes
|
|
BELOW_NORMAL_PRIORITY_CLASS = 0x00004000
|
|
ctypes.windll.kernel32.SetPriorityClass(
|
|
ctypes.windll.kernel32.OpenProcess(0x0200, False, proc.pid),
|
|
BELOW_NORMAL_PRIORITY_CLASS
|
|
)
|
|
print(f"Training process (PID {proc.pid}) set to BELOW_NORMAL priority.")
|
|
except Exception as e:
|
|
print(f"Could not set process priority: {e}")
|
|
|
|
return {"status": "success", "message": f"Training initiated for {req.epochs} epochs."}
|
|
|
|
@app.get("/api/train_status")
|
|
async def get_status():
|
|
"""
|
|
Returns current training status.
|
|
"""
|
|
global current_training_process
|
|
status = get_training_status()
|
|
|
|
# Auto-detect if process died unexpectedly
|
|
if status.get("status") == "training":
|
|
if current_training_process is None or current_training_process.poll() is not None:
|
|
status["status"] = "failed"
|
|
status["message"] = "Training-Prozess unerwartet beendet."
|
|
write_training_status(status)
|
|
current_training_process = None
|
|
|
|
return status
|
|
|
|
@app.post("/api/cancel_training")
|
|
async def cancel_training():
|
|
global current_training_process
|
|
if current_training_process is None or current_training_process.poll() is not None:
|
|
# Fallback: kill via PowerShell command on Windows
|
|
if sys.platform == "win32":
|
|
try:
|
|
import subprocess as sp
|
|
sp.run("powershell -Command \"Get-CimInstance Win32_Process | Where-Object { $_.CommandLine -like '*train.py*' } | ForEach-Object { Stop-Process -Id $_.ProcessId -Force }\"", shell=True, capture_output=True)
|
|
except Exception:
|
|
pass
|
|
|
|
status = get_training_status()
|
|
status["status"] = "failed"
|
|
status["message"] = "Training vom Benutzer abgebrochen."
|
|
write_training_status(status)
|
|
current_training_process = None
|
|
|
|
return {"status": "success", "message": "Training abgebrochen."}
|
|
|
|
try:
|
|
current_training_process.terminate()
|
|
current_training_process.wait(timeout=2.0)
|
|
except Exception:
|
|
try:
|
|
current_training_process.kill()
|
|
except Exception:
|
|
pass
|
|
|
|
current_training_process = None
|
|
|
|
# Read status, set to failed
|
|
status = get_training_status()
|
|
status["status"] = "failed"
|
|
status["message"] = "Training vom Benutzer abgebrochen."
|
|
write_training_status(status)
|
|
|
|
return {"status": "success", "message": "Training erfolgreich abgebrochen."}
|
|
|
|
@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)
|
|
env = os.environ.copy()
|
|
env["PYTHONUTF8"] = "1"
|
|
cmd = [sys.executable, os.path.join(BASE_DIR, "export.py")]
|
|
result = subprocess.run(cmd, cwd=BASE_DIR, capture_output=True, text=True, check=True, env=env)
|
|
|
|
# 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/calibrate_background")
|
|
async def calibrate_background(payload: CalibrateRequest):
|
|
"""
|
|
Saves a webcam snapshot of the empty scale to use as the background baseline.
|
|
"""
|
|
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")
|
|
|
|
bg_path = os.path.join(BASE_DIR, "dataset", "background.jpg")
|
|
os.makedirs(os.path.dirname(bg_path), exist_ok=True)
|
|
img.save(bg_path, "JPEG", quality=95)
|
|
|
|
return {
|
|
"status": "success",
|
|
"message": "Hintergrund erfolgreich kalibriert und als Null-Linie gespeichert!"
|
|
}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=f"Hintergrundkalibrierung fehlgeschlagen: {str(e)}")
|
|
|
|
@app.get("/api/has_background")
|
|
async def has_background():
|
|
"""
|
|
Checks if a baseline background image exists and returns metadata.
|
|
"""
|
|
bg_path = os.path.join(BASE_DIR, "dataset", "background.jpg")
|
|
exists = os.path.exists(bg_path)
|
|
timestamp = None
|
|
if exists:
|
|
timestamp = os.path.getmtime(bg_path)
|
|
return {
|
|
"has_background": exists,
|
|
"timestamp": timestamp
|
|
}
|
|
|
|
@app.post("/api/reset_background")
|
|
async def reset_background():
|
|
"""
|
|
Removes the baseline background image calibration.
|
|
"""
|
|
bg_path = os.path.join(BASE_DIR, "dataset", "background.jpg")
|
|
if os.path.exists(bg_path):
|
|
try:
|
|
os.remove(bg_path)
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=f"Fehler beim Zurücksetzen: {str(e)}")
|
|
return {"status": "success", "message": "Hintergrundkalibrierung zurückgesetzt."}
|
|
|
|
@app.post("/api/reset_dataset")
|
|
async def reset_dataset():
|
|
"""
|
|
Deletes all user-captured images, labels, and the baseline background image calibration.
|
|
"""
|
|
images_dir = os.path.join(BASE_DIR, "dataset", "images", "train")
|
|
labels_dir = os.path.join(BASE_DIR, "dataset", "labels", "train")
|
|
bg_path = os.path.join(BASE_DIR, "dataset", "background.jpg")
|
|
|
|
count_deleted = 0
|
|
try:
|
|
# Delete images starting with user_capture_
|
|
if os.path.exists(images_dir):
|
|
for f in os.listdir(images_dir):
|
|
if f.startswith("user_capture_"):
|
|
try:
|
|
os.remove(os.path.join(images_dir, f))
|
|
count_deleted += 1
|
|
except Exception:
|
|
pass
|
|
|
|
# Delete labels starting with user_capture_
|
|
if os.path.exists(labels_dir):
|
|
for f in os.listdir(labels_dir):
|
|
if f.startswith("user_capture_"):
|
|
try:
|
|
os.remove(os.path.join(labels_dir, f))
|
|
except Exception:
|
|
pass
|
|
|
|
# Delete background
|
|
if os.path.exists(bg_path):
|
|
try:
|
|
os.remove(bg_path)
|
|
except Exception:
|
|
pass
|
|
|
|
return {
|
|
"status": "success",
|
|
"message": f"Datenbank erfolgreich zurückgesetzt! {count_deleted} selbst erfasste Bilder gelöscht."
|
|
}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=f"Fehler beim Zurücksetzen: {str(e)}")
|
|
|
|
@app.post("/api/reset_model")
|
|
async def reset_model():
|
|
"""
|
|
Deletes the trained PyTorch model and exported ONNX model,
|
|
resets training status, and reloads the detector.
|
|
"""
|
|
pytorch_path = os.path.join(BASE_DIR, "models", "model.pt")
|
|
onnx_path = os.path.join(BASE_DIR, "models", "model.onnx")
|
|
|
|
try:
|
|
deleted_any = False
|
|
if os.path.exists(pytorch_path):
|
|
os.remove(pytorch_path)
|
|
deleted_any = True
|
|
if os.path.exists(onnx_path):
|
|
os.remove(onnx_path)
|
|
deleted_any = True
|
|
|
|
# Re-initialize training status to idle
|
|
write_training_status({
|
|
"status": "idle",
|
|
"epoch": 0,
|
|
"total_epochs": 0,
|
|
"train_loss": [],
|
|
"val_loss": [],
|
|
"progress": 0.0,
|
|
"message": "AI model has been reset to default/simulated state."
|
|
})
|
|
|
|
# Reload detector so it goes back to simulated mode
|
|
detector.load_model()
|
|
|
|
return {
|
|
"status": "success",
|
|
"message": "KI-Modell erfolgreich zurückgesetzt. Das System läuft nun wieder im Simulations-Modus."
|
|
}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=f"Fehler beim Zurücksetzen des Modells: {str(e)}")
|
|
|
|
@app.post("/api/capture_training_image")
|
|
async def capture_training_image(payload: CaptureRequest):
|
|
"""
|
|
Saves a webcam snapshot and generates a YOLO format bounding box label.
|
|
Uses OpenCV background subtraction if calibrated, falls back to client bbox or default center.
|
|
"""
|
|
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")
|
|
img_np = np.array(img)
|
|
|
|
# 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)
|
|
|
|
# Calculate bounding box using background subtraction
|
|
bg_path = os.path.join(BASE_DIR, "dataset", "background.jpg")
|
|
yolo_box = calculate_auto_crop_bbox(img_np, bg_path)
|
|
|
|
if yolo_box is None:
|
|
# Fallback to client-side bounding box if provided
|
|
if payload.bbox and all(k in payload.bbox for k in ["xmin", "ymin", "xmax", "ymax"]):
|
|
h_img, w_img, _ = img_np.shape
|
|
xmin = payload.bbox["xmin"]
|
|
ymin = payload.bbox["ymin"]
|
|
xmax = payload.bbox["xmax"]
|
|
ymax = payload.bbox["ymax"]
|
|
box_w = xmax - xmin
|
|
box_h = ymax - ymin
|
|
x_center = (xmin + box_w / 2.0) / w_img
|
|
y_center = (ymin + box_h / 2.0) / h_img
|
|
norm_w = box_w / w_img
|
|
norm_h = box_h / h_img
|
|
yolo_box = [x_center, y_center, norm_w, norm_h]
|
|
|
|
if yolo_box is None:
|
|
# Fallback to default centered box
|
|
yolo_box = [0.500000, 0.500000, 0.650000, 0.650000]
|
|
|
|
# Save YOLO label file
|
|
label_filename = f"user_capture_{timestamp}.txt"
|
|
with open(os.path.join(labels_dir, label_filename), "w") as f:
|
|
f.write(f"{class_idx} {yolo_box[0]:.6f} {yolo_box[1]:.6f} {yolo_box[2]:.6f} {yolo_box[3]:.6f}\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,
|
|
"filename": filename
|
|
}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=f"Fehler bei Bild-Erfassung: {str(e)}")
|
|
|
|
@app.post("/api/delete_captured_image")
|
|
async def delete_captured_image(payload: DeleteCaptureRequest):
|
|
"""
|
|
Deletes a user-captured image and its corresponding YOLO label file.
|
|
"""
|
|
filename = payload.filename
|
|
if not is_safe_filename(filename):
|
|
raise HTTPException(status_code=400, detail="Ungültiger Dateiname.")
|
|
|
|
images_dir = os.path.join(BASE_DIR, "dataset", "images", "train")
|
|
labels_dir = os.path.join(BASE_DIR, "dataset", "labels", "train")
|
|
|
|
img_path = os.path.join(images_dir, filename)
|
|
lbl_path = os.path.join(labels_dir, filename.replace(".jpg", ".txt"))
|
|
|
|
deleted = False
|
|
try:
|
|
if os.path.exists(img_path):
|
|
os.remove(img_path)
|
|
deleted = True
|
|
if os.path.exists(lbl_path):
|
|
os.remove(lbl_path)
|
|
deleted = True
|
|
|
|
if not deleted:
|
|
raise HTTPException(status_code=404, detail="Datei nicht gefunden.")
|
|
|
|
# Recount 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 {filename} und Label erfolgreich gelöscht.",
|
|
"count": count
|
|
}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=f"Fehler beim Löschen der Datei: {str(e)}")
|
|
|
|
@app.get("/api/captured_images/{filename}")
|
|
async def get_captured_image(filename: str):
|
|
"""
|
|
Serves a user-captured training image file securely.
|
|
"""
|
|
if not is_safe_filename(filename):
|
|
raise HTTPException(status_code=400, detail="Ungültiger Dateiname.")
|
|
|
|
img_path = os.path.join(BASE_DIR, "dataset", "images", "train", filename)
|
|
if not os.path.exists(img_path):
|
|
raise HTTPException(status_code=404, detail="Bild nicht gefunden.")
|
|
|
|
return FileResponse(img_path)
|
|
|
|
|
|
@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/background_image")
|
|
async def get_background_image():
|
|
"""
|
|
Serves the calibrated background image.
|
|
"""
|
|
bg_path = os.path.join(BASE_DIR, "dataset", "background.jpg")
|
|
if not os.path.exists(bg_path):
|
|
raise HTTPException(status_code=404, detail="Hintergrund nicht kalibriert.")
|
|
return FileResponse(bg_path)
|
|
|
|
@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"))
|