changed http request to websocket

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():

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")

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)

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)