initial commit
This commit is contained in:
DanielS
523
app/main.py
Normal file
523
app/main.py
Normal file
@@ -0,0 +1,523 @@
|
||||
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"))
|
||||
Reference in New Issue
Block a user