added yolo model to api

api/src/server.py

@@ -12,12 +12,15 @@ from PIL import Image
 import io
 import torch.nn.functional as F
 
+from src.yolo.train_yolo_faces import convert_prediction
+from src.yolo.yolo_model import Yolo_model
 from src.bird_cnn.bird_cnn import Bird_CNN
 
 import threading
 
 BUILD_PATH = "./build_models"
-IMAGE_SIZE = 64
+IMAGE_SIZE_CNN = 64
+IMAGE_SIZE_YOLO = 64
 
 sem_ai = threading.Semaphore(1)
 
@@ -30,19 +33,30 @@ class bird_species(Enum):
     Ruddy_Shelduck = 5
     Sarus_Crane = 6
 
-transform = transforms.Compose([
+transform_bird = transforms.Compose([
-    transforms.Resize(IMAGE_SIZE),
+    transforms.Resize(IMAGE_SIZE_CNN),
-    transforms.CenterCrop(IMAGE_SIZE),
+    transforms.CenterCrop(IMAGE_SIZE_CNN),
+    transforms.ToTensor()
+])
+
+transform_face = transforms.Compose([
+    transforms.Resize((IMAGE_SIZE_YOLO, IMAGE_SIZE_YOLO)),
     transforms.ToTensor()
 ])
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-model = Bird_CNN(c_in=3, c_hidden=16, c_out=7)
+model_bird = Bird_CNN(c_in=3, c_hidden=16, c_out=7)
 full_path = os.path.join(BUILD_PATH, "bird_cnn")
-model.load_state_dict(torch.load(full_path, map_location=torch.device(device)))
+model_bird.load_state_dict(torch.load(full_path, map_location=torch.device(device)))
-model.to(device)
+model_bird.to(device)
-model.eval()
+model_bird.eval()
+
+modeL_face = Yolo_model(c_in=3, boxes=1, grid=6, labels=1, c_hidden=16)
+full_path = os.path.join(BUILD_PATH, "face_detection_yolo")
+modeL_face.load_state_dict(torch.load(full_path, map_location=torch.device(device)))
+modeL_face.to(device)
+modeL_face.eval()
 
 app = FastAPI()
 
@@ -65,10 +79,10 @@ async def predict(file: UploadFile = File(...)):
         image_bytes = await file.read()
 
         image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-        image = transform(image).unsqueeze(0).to(device)
+        image = transform_bird(image).unsqueeze(0).to(device)
 
         with torch.no_grad():
-            pred = model(image)
+            pred = model_bird(image)
             probs = F.softmax(pred, dim=1)
             confidence, cls = torch.max(probs, dim=1)
 
@@ -77,6 +91,33 @@ 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()
+
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        C, H, W = image.shape
+        scale_w = W / IMAGE_SIZE_YOLO
+        scale_h = H / IMAGE_SIZE_YOLO
+        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)
+
+            boxes_to_send = []
+            for bbox in bboxes:
+                xmin = int(bbox[0] * scale_w)
+                ymin = int(bbox[1] * scale_h)
+                xmax = int(bbox[2] * scale_w)
+                ymax = int(bbox[3] * scale_h)
+                boxes_to_send.append([xmin, ymin, xmax, ymax])
+
+        return {
+            "bboxes": boxes_to_send
+        }
+
 
 load_dotenv("./database/.env")
 API_KEY = os.getenv("API_KEY")