diff --git a/api/src/server.py b/api/src/server.py
index a2eb327..9fb53b4 100644
--- a/api/src/server.py
+++ b/api/src/server.py
@@ -12,12 +12,16 @@ from PIL import Image
 import io
 import torch.nn.functional as F
 
-from app.src.yolo.train_yolo_faces import convert_prediction
-from app.src.yolo.yolo_model import Yolo_model
-from app.src.bird_cnn.bird_cnn import Bird_CNN
-
 import threading
 
+import sys
+sys.path.insert(1, './src/yolo')
+from train_yolo_faces import convert_prediction
+from yolo_model import Yolo_model
+
+sys.path.insert(1, './src/bird_cnn')
+from bird_cnn import Bird_CNN
+
 BUILD_PATH = "./build_models"
 IMAGE_SIZE_CNN = 64
 IMAGE_SIZE_YOLO = 64
@@ -97,7 +101,7 @@ async def predict_face(file: UploadFile = File(...)):
         image_bytes = await file.read()
 
         image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-        C, H, W = image.shape
+        H, W = image.size
         scale_w = W / IMAGE_SIZE_YOLO
         scale_h = H / IMAGE_SIZE_YOLO
         image = transform_face(image).unsqueeze(0).to(device)
diff --git a/api/src/yolo/train_yolo_faces.py b/api/src/yolo/train_yolo_faces.py
index 63d9a5b..4c38e89 100644
--- a/api/src/yolo/train_yolo_faces.py
+++ b/api/src/yolo/train_yolo_faces.py
@@ -2,10 +2,10 @@ import os
 
 import cv2
 from tqdm import tqdm
-from yolo.yolo_dataset import YoloDataset, turn_image_centered
-from yolo.yolo_model import train, Yolo_model, sample
-from yolo.yolo_loss import YoloLoss
-from util import TransformedSubset, test_workers_speed, visualizeImage
+from yolo_dataset import YoloDataset, turn_image_centered
+from yolo_model import train, Yolo_model, sample
+from yolo_loss import YoloLoss
+from src.util import TransformedSubset, test_workers_speed, visualizeImage
 from torch.utils.data import DataLoader, random_split
 from torchvision import datasets
 import torch