diff --git a/api/src/main.py b/api/src/main.py index 4917b26..7304312 100644 --- a/api/src/main.py +++ b/api/src/main.py @@ -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(): diff --git a/api/src/server.py b/api/src/server.py index 8d94da8..18df302 100644 --- a/api/src/server.py +++ b/api/src/server.py @@ -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") diff --git a/api/src/yolo/train_yolo_faces.py b/api/src/yolo/train_yolo_faces.py index 4c38e89..e3b20a9 100644 --- a/api/src/yolo/train_yolo_faces.py +++ b/api/src/yolo/train_yolo_faces.py @@ -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) diff --git a/api/src/yolo/yolo_dataset.py b/api/src/yolo/yolo_dataset.py index d368b88..b8a0d36 100644 --- a/api/src/yolo/yolo_dataset.py +++ b/api/src/yolo/yolo_dataset.py @@ -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) diff --git a/website/src/Object_Detection.jsx b/website/src/Object_Detection.jsx index 1dbd73e..80fe29b 100644 --- a/website/src/Object_Detection.jsx +++ b/website/src/Object_Detection.jsx @@ -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 (

Face Detection

- + +
+
) }