diff --git a/api/build_models/face_detection_yolo b/api/build_models/face_detection_yolo new file mode 100644 index 0000000..2a8edf8 Binary files /dev/null and b/api/build_models/face_detection_yolo differ diff --git a/api/src/server.py b/api/src/server.py index 3a6c4f9..b78f531 100644 --- a/api/src/server.py +++ b/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([ - transforms.Resize(IMAGE_SIZE), - transforms.CenterCrop(IMAGE_SIZE), +transform_bird = transforms.Compose([ + transforms.Resize(IMAGE_SIZE_CNN), + 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.to(device) -model.eval() +model_bird.load_state_dict(torch.load(full_path, map_location=torch.device(device))) +model_bird.to(device) +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) @@ -76,6 +90,33 @@ async def predict(file: UploadFile = File(...)): "class": bird_species(cls.item()).name, "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")