marvin/CNN_Website
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

changed http request to websocket

Browse Source
This commit is contained in:
marvin
2026-05-29 17:50:46 +02:00
parent 60c29d57ec
commit b49fc5ac24
5 changed files with 81 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/src/main.py
+8 -2
View File
@@ -1,5 +1,11 @@
from r_cnn.r_cnn_test import train_cnn_test
from yolo.train_yolo_faces import train_yolo
import sys
sys.path.insert(1, './src/yolo')
from yolo_model_production import convert_prediction, Yolo_model
from train_yolo_faces import train_yolo
sys.path.insert(2, './src/bird_cnn')
from bird_cnn import Bird_CNN
def main():
api/src/server.py
+18 -10
View File
@@ -2,11 +2,13 @@ from enum import Enum
import os
import sqlite3
import cv2
from dotenv import load_dotenv
import numpy as np
from pydantic import BaseModel
import torch
from torchvision import transforms
from fastapi import Depends, FastAPI, File, HTTPException, Header, UploadFile
from fastapi import Depends, FastAPI, File, HTTPException, Header, UploadFile, WebSocket
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
import io
@@ -95,17 +97,23 @@ async def predict(file: UploadFile = File(...)):
"confidence": confidence.item()
}
@app.post("/predict_face")
async def predict_face(file: UploadFile = File(...)):
with sem_ai:
image_bytes = await file.read()
@app.websocket("/predict_face")
async def predict_face(websocket: WebSocket):
await websocket.accept()
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
H, W = image.size
while True:
jpg_bytes = await websocket.receive_bytes()
np_arr = np.frombuffer(jpg_bytes, np.uint8)
frame = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
H, W, _ = frame.shape
scale_w = W / IMAGE_SIZE_YOLO
scale_h = H / IMAGE_SIZE_YOLO
image = transform_face(image).unsqueeze(0).to(device)
image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
image = transform_face(image).unsqueeze(0).to(device)
with torch.no_grad():
pred = modeL_face(image)
bboxes, _, _ = convert_prediction(pred.squeeze(0), image.squeeze(0), threshold=0.9)
@@ -118,9 +126,9 @@ async def predict_face(file: UploadFile = File(...)):
ymax = int(bbox[3] * scale_h)
boxes_to_send.append([xmin, ymin, xmax, ymax])
return {
await websocket.send_json({
"bboxes": boxes_to_send
}
})
load_dotenv("./database/.env")
api/src/yolo/train_yolo_faces.py
+3 -40
View File
@@ -2,25 +2,17 @@ import os
import cv2
from tqdm import tqdm
from yolo_dataset import YoloDataset, turn_image_centered
from yolo_dataset import YoloDataset
from yolo_model import train, Yolo_model, sample
from yolo_loss import YoloLoss
from src.util import TransformedSubset, test_workers_speed, visualizeImage
from util import TransformedSubset, test_workers_speed, visualizeImage
from torch.utils.data import DataLoader, random_split
from torchvision import datasets
import torch
from torchvision.utils import draw_bounding_boxes
import torch.nn.functional as F
def xy_center_to_edges(xcenter, ycenter, width, height):
width = max(width, 1)
height = max(height, 1)
x = xcenter - (width / 2)
y = ycenter - (height / 2)
return [x, y, x + width, y + height]
from yolo_model_production import convert_prediction
def view_data(dataset):
dataloader = DataLoader(dataset=dataset, batch_size=1, shuffle=False)
@@ -51,35 +43,6 @@ def visualize_boxes(label, image, threshold=0.95):
image = draw_bounding_boxes(image, boxes_to_draw, colors=(255, 0, 0))
visualizeImage(image)
def convert_prediction(label, image, threshold=0.9):
image = image.clone().detach()
label = label.clone().detach()
image_size = image.shape[1]
grid_number = label.shape[0]
grid_size = image_size / grid_number
boxes_to_draw = []
grids_to_draw_obj = []
grids_to_draw_noobj = []
for x in range(label.shape[0]):
for y in range(label.shape[1]):
if label[x, y, 4].item() < threshold:
grids_to_draw_noobj.append([x*grid_size, y*grid_size, (x+1)*grid_size, (y+1)*grid_size]) #xmin, ymin, xmax, ymax
continue
grids_to_draw_obj.append([x*grid_size, y*grid_size, (x+1)*grid_size, (y+1)*grid_size]) #xmin, ymin, xmax, ymax
boxx = label[x, y, 0] * (image_size / grid_number)
boxy = label[x, y, 1] * (image_size / grid_number)
boxw = label[x, y, 2] * image_size
boxh = label[x, y, 3] * image_size
boxx, boxy = turn_image_centered(x=boxx, y=boxy, img_w=image_size, img_h=image_size, S=grid_number, cell_i=x, cell_j=y)
boxes_to_draw.append(xy_center_to_edges(boxx, boxy, boxw, boxh)) #xmin, ymin, xmax, ymax
return boxes_to_draw, grids_to_draw_obj, grids_to_draw_noobj
def use_webcam(grid, img_size):
# 0 = default webcam
cap = cv2.VideoCapture(0)
api/src/yolo/yolo_dataset.py
-9
View File
@@ -37,15 +37,6 @@ def turn_grid_centered(x, y, img_w, img_h, S, cell_i, cell_j):
return x - cell_border_w, y - cell_border_h
def turn_image_centered(x, y, img_w, img_h, S, cell_i, cell_j):
cell_w = img_w / S
cell_h = img_h / S
cell_border_w = cell_w * cell_i
cell_border_h = cell_h * cell_j
return x + cell_border_w, y + cell_border_h
def flip_bbox_horizontal(box, image_width):
x, y, w, h = box
return (image_width - x - w, y, w, h)
website/src/Object_Detection.jsx
+52 -6
View File
@@ -1,8 +1,12 @@
import styles from './Object_Detection.module.css';
import { useRef, useEffect } from 'react'
import { useRef, useEffect, useState } from 'react'
function Object_Detection() {
const videoRef = useRef(null);
const canvasRef = useRef(null);
const [error, setError] = useState(false);
const [loading, setLoading] = useState(false);
const [bboxes, setBboxes] = useState(null)
useEffect(() => {
async function startWebcam() {
@@ -44,20 +48,62 @@ function Object_Detection() {
ctx.drawImage(video, 0, 0, width, height);
const dataUrl = canvas.toDataURL("image/png");
ctx.strokeStyle = "red";
ctx.lineWidth = 4;
setImage(dataUrl);
if (bboxes) {
bboxes.forEach(element => {
ctx.strokeRect(element[0], element[0], element[0], element[0]);
});
}
canvas.toBlob(async (blob) => {
if (!blob) return;
const formData = new FormData();
formData.append("file", blob, "frame.jpg");
setLoading(true)
try {
const response = await fetch("https://api.marvinkrausser.com/predict_face", {
method: "POST",
body: formData,
});
if (!response.ok) {
setError(true);
return;
}
else {
setError(false);
}
const result = await response.json();
setBboxes(result["bboxes"])
} catch (e) {
setError(true);
}
finally {
setLoading(false);
}
}, "image/jpeg", 0.9);
};
setInterval(captureImage, 100);
return (
<div className='site-box'>
<h1 className='site-headline'>Face Detection</h1>
<div id={styles.container}>
<video autoPlay={true} id={styles.videoElement} ref={videoRef}>
</video>
<video autoPlay={true} ref={videoRef} style={{ display: "none" }}></video>
<canvas ref={canvasRef} width={640} height={480} />
</div>
<button onClick={captureImage}>Click</button>
</div>
)
}