changed http request to websocket

api/src/main.py

@@ -1,5 +1,11 @@
-from r_cnn.r_cnn_test import train_cnn_test
+import sys
-from yolo.train_yolo_faces import train_yolo
+
+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

@@ -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")
+@app.websocket("/predict_face")
-async def predict_face(file: UploadFile = File(...)):
+async def predict_face(websocket: WebSocket):
-    with sem_ai:
+    await websocket.accept()
-        image_bytes = await file.read()
 
-        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    while True:
-        H, W = image.size
+        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

@@ -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):
+from yolo_model_production import convert_prediction
-    width = max(width, 1)
-    height = max(height, 1)
-
-    x = xcenter - (width / 2)
-    y = ycenter - (height / 2)
-
-    return [x, y, x + width, y + height]
-
 
 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

@@ -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

@@ -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 autoPlay={true} ref={videoRef} style={{ display: "none" }}></video>
-
+                <canvas ref={canvasRef} width={640} height={480} />
-                </video>
             </div>
+            <button onClick={captureImage}>Click</button>
         </div>
     )
 }