diff --git a/.gitignore b/.gitignore index 6451481..00424e9 100755 --- a/.gitignore +++ b/.gitignore @@ -29,6 +29,7 @@ server/venv server/output/categories/* server/__pycache__/ server/db/__pycache__/ +server/handlers/__pycache__/ secret.* secret.story.* *.secret.* diff --git a/client/src/ImageUpload/index.jsx b/client/src/ImageUpload/index.jsx index 7acbb1c..45b64fe 100644 --- a/client/src/ImageUpload/index.jsx +++ b/client/src/ImageUpload/index.jsx @@ -1,174 +1,314 @@ -import React, { useCallback, useState, useEffect } from "react" -import { useDropzone } from "react-dropzone" -import { createTheme } from "@mui/material/styles" -import { CssBaseline, Box, Typography, IconButton } from "@mui/material" -import DeleteIcon from "@mui/icons-material/Delete" -import axios from "axios" -import { useSnackbar } from "../SnackbarContext" -import { useTranslation, Trans } from "react-i18next" -import useMediaQuery from "@mui/material/useMediaQuery" -import config from "../config.js" - -const theme = createTheme() +import React, { useCallback, useState, useEffect, useRef } from "react"; +import { useDropzone } from "react-dropzone"; +import { createTheme } from "@mui/material/styles"; +import { CssBaseline, Box, Typography, IconButton, Button, CircularProgress } from "@mui/material"; +import DeleteIcon from "@mui/icons-material/Delete"; +import PlayArrowIcon from "@mui/icons-material/PlayArrow"; +import axios from "axios"; +import { useSnackbar } from "../SnackbarContext"; +import { useTranslation, Trans } from "react-i18next"; +import useMediaQuery from "@mui/material/useMediaQuery"; +import config from "../config.js"; + +const theme = createTheme(); + +// Helper: generate a consistent color from a class name (HSL) +const getColorForClass = (className) => { + let hash = 0; + for (let i = 0; i < className.length; i++) { + hash = className.charCodeAt(i) + ((hash << 5) - hash); + } + const hue = Math.abs(hash % 360); + return `hsl(${hue}, 70%, 50%)`; +}; + +// Draw image with bounding boxes, scaled to fit within maxWidth/maxHeight +const drawImageWithBoxes = (canvas, imageUrl, detections, maxWidth = 800, maxHeight = 600) => { + const ctx = canvas.getContext("2d"); + const img = new Image(); + img.crossOrigin = "Anonymous"; + img.onload = () => { + // Calculate display size preserving aspect ratio + let width = img.width; + let height = img.height; + if (width > maxWidth) { + height = (height * maxWidth) / width; + width = maxWidth; + } + if (height > maxHeight) { + width = (width * maxHeight) / height; + height = maxHeight; + } + canvas.width = width; + canvas.height = height; + + // Draw image scaled to display size + ctx.drawImage(img, 0, 0, width, height); + + if (detections && detections.length) { + const scaleX = width / img.width; + const scaleY = height / img.height; + + detections.forEach((det) => { + const [x1, y1, x2, y2] = det.bbox; + // Scale coordinates to display size + const sx1 = x1 * scaleX; + const sy1 = y1 * scaleY; + const sx2 = x2 * scaleX; + const sy2 = y2 * scaleY; + const boxWidth = sx2 - sx1; + const boxHeight = sy2 - sy1; + + const classColor = getColorForClass(det.class); + ctx.strokeStyle = classColor; + ctx.lineWidth = 3; + ctx.strokeRect(sx1, sy1, boxWidth, boxHeight); + + // Draw label background + const label = `[${det.class}] ${(det.confidence * 100).toFixed(1)}%`; + ctx.font = "bold 16px Arial"; + const textMetrics = ctx.measureText(label); + const textWidth = textMetrics.width; + const textHeight = 20; + ctx.fillStyle = "rgba(0,0,0,0.6)"; + ctx.fillRect(sx1, sy1 - textHeight - 4, textWidth + 8, textHeight + 4); + + // Draw label text + ctx.fillStyle = classColor; + ctx.fillText(label, sx1 + 4, sy1 - 6); + }); + } + }; + img.src = imageUrl; +}; const ImageUpload = ({ onImageUpload, settingsImages }) => { - const { t } = useTranslation() - const { showSnackbar } = useSnackbar() - const [images, setImages] = useState(settingsImages) - const [loading, setLoading] = useState(false) - const [progress, setProgress] = useState(0) - const isSmallDevice = useMediaQuery(theme.breakpoints.down("sm")) - - // Sync with parent prop + const { t } = useTranslation(); + const { showSnackbar } = useSnackbar(); + const [images, setImages] = useState(settingsImages); + const [loading, setLoading] = useState(false); + const [progress, setProgress] = useState(0); + const [inferenceLoading, setInferenceLoading] = useState(false); + const [selectedIndex, setSelectedIndex] = useState(null); + const [inferredImageUrl, setInferredImageUrl] = useState(null); + const [inferredDetections, setInferredDetections] = useState([]); + const largeCanvasRef = useRef(null); + const isSmallDevice = useMediaQuery(theme.breakpoints.down("sm")); + useEffect(() => { - setImages(settingsImages) - }, [settingsImages]) + setImages(settingsImages); + // Reset inference preview when external images change + setInferredImageUrl(null); + setInferredDetections([]); + setSelectedIndex(null); + }, [settingsImages]); const onDrop = useCallback( (acceptedFiles, fileRejections) => { if (fileRejections.length) { - const { errors } = fileRejections[0] + const { errors } = fileRejections[0]; if (errors.length) { - showSnackbar(errors[0].message, "error") - return + showSnackbar(errors[0].message, "error"); + return; } } - // Prevent new uploads while a previous one is in progress if (loading) { - showSnackbar(t("error.upload_in_progress"), "info") - return + showSnackbar(t("error.upload_in_progress"), "info"); + return; } - const totalImages = images.length + acceptedFiles.length + const totalImages = images.length + acceptedFiles.length; if (totalImages > config.UPLOAD_LIMIT) { - showSnackbar(t("error.configuration.image_upload.max"), "error") - return + showSnackbar(t("error.configuration.image_upload.max"), "error"); + return; } - const newImages = acceptedFiles.map((file) => { - return Object.assign(file, { - preview: URL.createObjectURL(file), - imageName: file.name, - }) - }) + const newImages = acceptedFiles.map((file) => ({ + file: file, + preview: URL.createObjectURL(file), + imageName: file.name, + })); - uploadImages(newImages) + uploadImages(newImages); }, - [images, onImageUpload, showSnackbar, loading], - ) + [images, onImageUpload, showSnackbar, loading] + ); const uploadImages = async (newImages) => { - const formData = new FormData() - + const formData = new FormData(); newImages.forEach((image) => { - formData.append("file", image) - }) + formData.append("file", image.file); + }); try { - setLoading(true) - const response = await axios.post( - `${config.SERVER_URL}/upload`, - formData, - { - headers: { - "Content-Type": "multipart/form-data", - }, - onUploadProgress: (progressEvent) => { - const { loaded, total } = progressEvent - let percentCompleted = Math.floor((loaded * 100) / total) - setProgress(percentCompleted) - }, + setLoading(true); + const response = await axios.post(`${config.SERVER_URL}/upload`, formData, { + headers: { "Content-Type": "multipart/form-data" }, + onUploadProgress: (progressEvent) => { + const { loaded, total } = progressEvent; + let percentCompleted = Math.floor((loaded * 100) / total); + setProgress(percentCompleted); }, - ) - showSnackbar(response.data.message, "success") + }); + showSnackbar(response.data.message, "success"); - const uploadedFiles = response.data.files - const uploadedImages = uploadedFiles.map((file) => ({ + const uploadedFiles = response.data.files; + const uploadedImages = uploadedFiles.map((file, idx) => ({ preview: file.url, filename: file.filename, - })) + imageName: newImages[idx]?.imageName || file.filename, + })); - // Merge existing images with new ones - const updatedImages = [...images, ...uploadedImages] - setImages(updatedImages) - onImageUpload(updatedImages) + const updatedImages = [...images, ...uploadedImages]; + setImages(updatedImages); + onImageUpload(updatedImages); + // Clear inference preview after new upload + setInferredImageUrl(null); + setInferredDetections([]); + setSelectedIndex(null); } catch (error) { - const errorResponse = error?.response?.data - if (errorResponse) { - showSnackbar(errorResponse?.message, "error") - } else { - showSnackbar(t("error.server_connection"), "error") - } - console.error("Error uploading images:", error) + const errorResponse = error?.response?.data; + showSnackbar(errorResponse?.message || t("error.server_connection"), "error"); + console.error("Error uploading images:", error); } finally { - setLoading(false) - setProgress(0) + setLoading(false); + setProgress(0); } - } + }; const deleteImage = async (filename, isNotFound = false) => { try { if (isNotFound) { - const updatedImages = images.filter( - (image) => image.filename !== filename, - ) - setImages(updatedImages) - onImageUpload(updatedImages) + const updatedImages = images.filter((image) => image.filename !== filename); + setImages(updatedImages); + onImageUpload(updatedImages); } else { - const response = await axios.delete( - `${config.SERVER_URL}/uploads/${filename}`, - ) - showSnackbar(response.data.message, "success") - - const updatedImages = images.filter( - (image) => image.filename !== filename, - ) - setImages(updatedImages) - onImageUpload(updatedImages) + const response = await axios.delete(`${config.SERVER_URL}/uploads/${filename}`); + showSnackbar(response.data.message, "success"); + const updatedImages = images.filter((image) => image.filename !== filename); + setImages(updatedImages); + onImageUpload(updatedImages); } - } catch (error) { - if (error?.response?.data) { - showSnackbar(error.response.data.message, "error") - } else { - showSnackbar(t("error.server_connection"), "error") + // If the deleted image was the one showing inference, clear preview + if (selectedIndex !== null && images[selectedIndex]?.filename === filename) { + setSelectedIndex(null); + setInferredImageUrl(null); + setInferredDetections([]); } - console.error("Error deleting image:", error) + } catch (error) { + showSnackbar(error?.response?.data?.message || t("error.server_connection"), "error"); + console.error("Error deleting image:", error); } - } + }; const handleImageError = (index) => { - const updatedImages = [...images] - updatedImages[index].isNotFound = true - setImages(updatedImages) - showSnackbar(t("error.image_not_found"), "error") - } + const updatedImages = [...images]; + updatedImages[index].isNotFound = true; + setImages(updatedImages); + showSnackbar(t("error.image_not_found"), "error"); + }; const handleRemoveImage = (index) => { - const imageToRemove = images[index] + const imageToRemove = images[index]; if (imageToRemove && !imageToRemove.filename) { - if (imageToRemove.src) { - let parts = imageToRemove.src.split("/") - let filename = parts[parts.length - 1] - imageToRemove.filename = filename - } + let parts = imageToRemove.src?.split("/") || []; + let filename = parts[parts.length - 1]; + imageToRemove.filename = filename; } if (imageToRemove && imageToRemove.filename) { - deleteImage(imageToRemove.filename, imageToRemove.isNotFound) + deleteImage(imageToRemove.filename, imageToRemove.isNotFound); } else { - console.error( - "Error deleting image: imageToRemove or imageToRemove.filename is undefined", - ) + console.error("Error deleting image: filename missing"); + } + }; + + // Run inference on selected image and display result in large preview area + const runInferenceOnSelected = async () => { + if (selectedIndex === null) { + showSnackbar(t("menu.no_image_selected"), "warning"); + return; + } + + const image = images[selectedIndex]; + if (!image) { + showSnackbar(t("error.inference_failed"), "error"); + return; } - } + + // Extract filename from src + let filename = null; + if (image.src) { + const parts = image.src.split('/'); + filename = parts[parts.length - 1]; + } + if (!filename && image.filename) filename = image.filename; + if (!filename && image.name) filename = image.name; + + if (!filename) { + console.error("Could not determine filename for image", image); + showSnackbar(t("error.inference_failed"), "error"); + return; + } + + // Get labels from localStorage + let labels = []; + try { + const settings = localStorage.getItem("settings_settings"); + if (settings) { + const parsed = JSON.parse(settings); + if (parsed.configuration && parsed.configuration.labels) { + labels = parsed.configuration.labels.map(labelObj => labelObj.id); + } + } + } catch (err) { + console.error("Failed to read labels from localStorage", err); + } + + setInferenceLoading(true); + try { + const response = await axios.post(`${config.SERVER_URL}/infer`, { + filenames: [filename], + labels: labels, + }); + const detections = response.data.results?.[0] || []; + + setInferredImageUrl(image.preview || image.src); + setInferredDetections(detections); + + // Draw on large canvas after a short delay (ensures ref is ready) + setTimeout(() => { + if (largeCanvasRef.current) { + drawImageWithBoxes(largeCanvasRef.current, image.preview || image.src, detections, 800, 600); + } + }, 50); + + showSnackbar(t("menu.inference_success"), "success"); + } catch (error) { + console.error("Inference error:", error); + const errorMsg = error?.response?.data?.message || error.message || t("error.inference_failed"); + showSnackbar(errorMsg, "error"); + } finally { + setInferenceLoading(false); + } + }; + + const clearInference = () => { + setInferredImageUrl(null); + setInferredDetections([]); + if (largeCanvasRef.current) { + const ctx = largeCanvasRef.current.getContext("2d"); + ctx.clearRect(0, 0, largeCanvasRef.current.width, largeCanvasRef.current.height); + } + }; const { getRootProps, getInputProps, isDragActive } = useDropzone({ onDrop, accept: "image/*", multiple: true, maxFiles: config.UPLOAD_LIMIT, - disabled: loading, // Disable dropzone during upload - }) + disabled: loading, + }); return ( <> @@ -210,9 +350,7 @@ const ImageUpload = ({ onImageUpload, settingsImages }) => { )} ) : ( - + { )} - - {images.map((image, index) => ( + + {/* Inference Controls */} + + + {inferredDetections.length > 0 && ( + + )} + + + {/* Large Inference Preview Area */} + {inferenceLoading ? ( + + + + ) : ( + inferredImageUrl && ( - preview handleImageError(index)} + + {t("menu.inference_result")} + + - handleRemoveImage(index)} - sx={{ position: "absolute", top: "0", right: "0" }} - > - - - ))} + ) + )} + + {/* Image Gallery (Thumbnails) */} + + {images.map((image, index) => { + const isSelected = selectedIndex === index; + return ( + setSelectedIndex(index)} + sx={{ + cursor: "pointer", + border: isSelected ? "3px solid #1976d2" : "3px solid transparent", + borderRadius: "8px", + padding: "2px", + transition: "all 0.2s", + "&:hover": { opacity: 0.9 }, + }} + > + preview handleImageError(index)} + style={{ + width: isSmallDevice ? "65px" : "82px", + height: isSmallDevice ? "65px" : "82px", + objectFit: "cover", + borderRadius: "4px", + marginBottom: "0.5rem", + }} + /> + { + e.stopPropagation(); + handleRemoveImage(index); + }} + sx={{ position: "absolute", top: "0", right: "0" }} + > + + + + ); + })} - ) -} + ); +}; -export default ImageUpload \ No newline at end of file +export default ImageUpload; \ No newline at end of file diff --git a/client/src/Localization/translation-de-DE.js b/client/src/Localization/translation-de-DE.js index 3bf51eb..0574d4f 100644 --- a/client/src/Localization/translation-de-DE.js +++ b/client/src/Localization/translation-de-DE.js @@ -6,6 +6,15 @@ const translationDeDE = { "menu.class_distribution_count": "Anzahl", "menu.keyframes": "Schlüsselbilder", "menu.images": "Bilder", + "menu.inference_button": "Inferenz starten", + "menu.clear_results": "Ergebnisse löschen", + "menu.no_images_infer": "Keine Bilder zur Inferenz", + "menu.inference_success": "Inferenz abgeschlossen", + "error.inference_failed": "Inferenz fehlgeschlagen", + "error.upload_in_progress": "Upload läuft, bitte warten", + "loading": "Laden...", + "menu.inference_result": "Inferenzergebnis", + "menu.no_image_selected": "Bitte wählen Sie zuerst ein Bild aus", "error.image": "Bild konnte nicht geladen werden", "error.imagevideo": "Fehlendes erforderliches Attribut", "error.upload_in_progress": "Upload bereits im Gange. Bitte warten.", diff --git a/client/src/Localization/translation-en-EN.js b/client/src/Localization/translation-en-EN.js index 61505f3..d2c5500 100644 --- a/client/src/Localization/translation-en-EN.js +++ b/client/src/Localization/translation-en-EN.js @@ -6,6 +6,15 @@ const translationEnEN = { "menu.regions": "Regions", "menu.keyframes": "Keyframes", "menu.images": "Images", + "menu.inference_button": "Run Inference", + "menu.clear_results": "Clear Results", + "menu.no_images_infer": "No images to infer", + "menu.inference_success": "Inference completed", + "error.inference_failed": "Inference failed", + "error.upload_in_progress": "Upload in progress, please wait", + "loading": "Loading...", + "menu.no_image_selected": "Please select an image first", + "menu.inference_result": "Inference Result", "error.image": "Could not load image", "error.imagevideo": "Missing required prop", "error.upload_in_progress": "Upload already in progress. Please wait", diff --git a/server/.env b/server/.env index c924978..ea3cf76 100644 --- a/server/.env +++ b/server/.env @@ -1 +1,5 @@ CLIENT_URL=http://localhost:5173 +MODEL_PATH=./models/best.onnx +CONFIDENCE_THRESHOLD=0.5 +IOU_THRESHOLD=0.45 +INPUT_SIZE=640,640 \ No newline at end of file diff --git a/server/app.py b/server/app.py index fbbd53e..827bdf9 100644 --- a/server/app.py +++ b/server/app.py @@ -2,9 +2,11 @@ from flask import Flask, jsonify, request, url_for, send_from_directory, send_file, after_this_request, current_app from flask_cors import CORS, cross_origin from db.db_handler import Module +import onnxruntime as ort from io import BytesIO import pandas as pd import requests +import cv2 import json from PIL import Image, ImageDraw import os @@ -16,6 +18,7 @@ import re from datetime import datetime from sam_model import SamModel +from handlers.yolo_handler import YOLOHandler from utils import load_image_from_url, format_regions_for_frontend app = Flask(__name__) @@ -36,6 +39,31 @@ # Get the CLIENT_URL environment variable, set a default to 80 client_url = os.getenv("CLIENT_URL", "http://localhost") +MODEL_PATH = os.getenv("MODEL_PATH", "./models/best.onnx") +CONFIDENCE_THRESHOLD = float(os.getenv("CONFIDENCE_THRESHOLD", "0.5")) +IOU_THRESHOLD = float(os.getenv("IOU_THRESHOLD", "0.45")) +INPUT_SIZE = tuple(map(int, os.getenv("INPUT_SIZE", "640,640").split(','))) + + +ort_session = None +handler = None + +def init_model(): + global ort_session, handler + if not os.path.exists(MODEL_PATH): + print(f"ERROR: Model not found at {MODEL_PATH}") + return + ort_session = ort.InferenceSession(MODEL_PATH, providers=['CPUExecutionProvider']) + handler = YOLOHandler( + input_size=INPUT_SIZE, + conf_threshold=CONFIDENCE_THRESHOLD, + iou_threshold=IOU_THRESHOLD + ) + print(f"Model loaded from {MODEL_PATH}") + print(f"YOLOHandler: input_size={INPUT_SIZE}, conf={CONFIDENCE_THRESHOLD}, iou={IOU_THRESHOLD}") + + +init_model() # Set the folder to save uploaded files UPLOAD_FOLDER = "uploads" @@ -129,6 +157,61 @@ def save_annotate_info(): ALLOWED_EXTENSIONS = {"png", "jpg", "jpeg", "gif"} +def allowed_file(filename): + return "." in filename and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSIONS + +@app.route("/infer", methods=["POST"]) +@cross_origin(origin=client_url, headers=["Content-Type"]) +def run_inference(): + + if ort_session is None or handler is None: + return jsonify({"error": "Model not loaded"}), 500 + + data = request.get_json() + if not data or "filenames" not in data: + return jsonify({"error": "Missing 'filenames' in request"}), 400 + + filenames = data["filenames"] + class_names = data.get("labels", []) + if not class_names: + class_names = [f"class_{i}" for i in range(100)] + else: + print(f"Using provided labels: {class_names}", flush=True) + + results = [] + for idx, filename in enumerate(filenames): + file_path = os.path.join(UPLOAD_FOLDER, filename) + + if not os.path.exists(file_path): + results.append([]) + continue + + img = cv2.imread(file_path) + if img is None: + results.append([]) + continue + + original_shape = img.shape[:2] + + input_tensor = handler.preprocess(img) + + ort_inputs = {ort_session.get_inputs()[0].name: input_tensor} + ort_outputs = ort_session.run(None, ort_inputs) + + for i, out in enumerate(ort_outputs): + print(f" Output {i} shape: {out.shape}", flush=True) + + # Optional: print min/max values of raw output + raw_out = ort_outputs[0] + + detections = handler.postprocess(ort_outputs, original_shape, class_names) + if len(detections) > 0: + print(f"First detection: {detections[0]}", flush=True) + + results.append(detections) + return jsonify({"results": results}), 200 + + def allowed_file(filename): return "." in filename and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSIONS diff --git a/server/handlers/__init__.py b/server/handlers/__init__.py new file mode 100644 index 0000000..0b1a670 --- /dev/null +++ b/server/handlers/__init__.py @@ -0,0 +1,3 @@ +from .yolo_handler import YOLOHandler + +#from .ssd_handler import SSDHandler \ No newline at end of file diff --git a/server/handlers/yolo_handler.py b/server/handlers/yolo_handler.py new file mode 100644 index 0000000..b43caa8 --- /dev/null +++ b/server/handlers/yolo_handler.py @@ -0,0 +1,128 @@ +import cv2 +import numpy as np + +class YOLOHandler: + """ + YOLO ONNX model handler (supports YOLOv5, YOLOv8, etc.) + Expects output shape: (1, 84, N) where N = number of predictions. + """ + def __init__(self, input_size=(640, 640), conf_threshold=0.5, iou_threshold=0.45): + self.input_size = input_size + self.conf_threshold = conf_threshold + self.iou_threshold = iou_threshold + + # ---------- Helper methods ---------- + def _letterbox(self, img, color=(114, 114, 114)): + """Resize and pad image to exact self.input_size dimensions.""" + shape = img.shape[:2] # current (height, width) + target_h, target_w = self.input_size + + # Compute scale ratio and new unpad dimensions + r = min(target_h / shape[0], target_w / shape[1]) + new_unpad_h = int(round(shape[0] * r)) + new_unpad_w = int(round(shape[1] * r)) + + # Resize to new unpad dimensions + if (new_unpad_w, new_unpad_h) != (shape[1], shape[0]): + img = cv2.resize(img, (new_unpad_w, new_unpad_h), interpolation=cv2.INTER_LINEAR) + + # Compute padding + dw = target_w - new_unpad_w + dh = target_h - new_unpad_h + left = dw // 2 + right = dw - left + top = dh // 2 + bottom = dh - top + + # Add borders + img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) + + # Final safety check: if dimensions still not exact, force resize + if img.shape[0] != target_h or img.shape[1] != target_w: + img = cv2.resize(img, (target_w, target_h), interpolation=cv2.INTER_LINEAR) + + return img, (r, r), (dw, dh) + + def _scale_coords(self, img1_shape, coords, img0_shape): + """Rescale coordinates (xyxy) from img1_shape to img0_shape.""" + gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) + pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 + coords[:, [0, 2]] -= pad[0] + coords[:, [1, 3]] -= pad[1] + coords[:, :4] /= gain + coords[:, [0, 2]] = coords[:, [0, 2]].clip(0, img0_shape[1]) + coords[:, [1, 3]] = coords[:, [1, 3]].clip(0, img0_shape[0]) + return coords + + def _nms(self, boxes, scores): + """Non‑maximum suppression using OpenCV.""" + indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), self.conf_threshold, self.iou_threshold) + if len(indices) == 0: + return [] + if isinstance(indices, tuple): + indices = indices[0] + return indices.flatten() if hasattr(indices, 'flatten') else indices + + # ---------- Public methods ---------- + def preprocess(self, image_bgr): + img_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) + img, _, _ = self._letterbox(img_rgb, self.input_size) + # img should already be target_h x target_w + img = img.astype(np.float32) / 255.0 + img = np.transpose(img, (2, 0, 1)) # HWC -> CHW + img = np.expand_dims(img, axis=0) # add batch dim + # Assert shape is correct (optional) + assert img.shape == (1, 3, self.input_size[0], self.input_size[1]), \ + f"Expected shape (1,3,{self.input_size[0]},{self.input_size[1]}), got {img.shape}" + return img + + def postprocess(self, raw_output, original_shape, class_names): + """ + Convert raw ONNX output to list of detections. + raw_output: either a list of outputs or a single numpy array. + original_shape: (height, width) of the original image. + class_names: list of strings (index = class id). + Returns list of dicts: [{"bbox": [x1,y1,x2,y2], "class": str, "confidence": float}] + """ + output = raw_output[0] if isinstance(raw_output, list) else raw_output + predictions = np.squeeze(output).T # (N, 84) for YOLOv8 + # For YOLO, first 4 are bbox (cx,cy,w,h), rest are class scores + scores = np.max(predictions[:, 4:], axis=1) # max class score + # Filter by confidence + mask = scores > self.conf_threshold + if not mask.any(): + return [] + predictions = predictions[mask] + scores = scores[mask] + # Extract boxes in xywh (normalized) + xywh = predictions[:, :4] + # Convert to xyxy + x1 = xywh[:, 0] - xywh[:, 2] / 2 + y1 = xywh[:, 1] - xywh[:, 3] / 2 + x2 = xywh[:, 0] + xywh[:, 2] / 2 + y2 = xywh[:, 1] + xywh[:, 3] / 2 + boxes = np.stack([x1, y1, x2, y2], axis=1) + # Class ids + class_ids = np.argmax(predictions[:, 4:], axis=1) + # Apply NMS + keep = self._nms(boxes, scores) + if len(keep) == 0: + return [] + boxes = boxes[keep] + scores = scores[keep] + class_ids = class_ids[keep] + # Scale boxes to original image size + h, w = original_shape + img1_shape = self.input_size + boxes = self._scale_coords(img1_shape, boxes, (h, w)) + # Build result list + results = [] + for box, score, cls_id in zip(boxes, scores, class_ids): + x1, y1, x2, y2 = box.astype(int) + class_name = class_names[cls_id] if cls_id < len(class_names) else f"class_{cls_id}" + results.append({ + "bbox": [int(x1), int(y1), int(x2), int(y2)], + "class": class_name, + "confidence": float(score) + }) + return results \ No newline at end of file diff --git a/server/requirements.txt b/server/requirements.txt index 74ab455..7c42812 100644 --- a/server/requirements.txt +++ b/server/requirements.txt @@ -8,4 +8,5 @@ pandas pillow requests zipp>=3.19.1 # not directly required, pinned by Snyk to avoid a vulnerability -black \ No newline at end of file +black +onnxruntime