refactor api

api/src/production/inference.py

@@ -0,0 +1,85 @@
+"""Loading the models and running one prediction. These functions are
+blocking and run on the InferenceGate worker threads."""
+
+import io
+
+from PIL import Image
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
+
+from src.common.boxes import scale_box
+from src.common.checkpoints import load_checkpoint
+from src.models.bird_cnn import BirdSpecies, Bird_CNN
+from src.models.yolo import Yolo_model, convert_prediction
+from src.production import settings
+
+Image.MAX_IMAGE_PIXELS = settings.MAX_IMAGE_PIXELS
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+models = {}
+
+transform_bird = transforms.Compose([
+    transforms.Resize(settings.BIRD_IMAGE_SIZE),
+    transforms.CenterCrop(settings.BIRD_IMAGE_SIZE),
+    transforms.ToTensor()
+])
+
+transform_face = transforms.Compose([
+    transforms.Resize((settings.FACE_IMAGE_SIZE, settings.FACE_IMAGE_SIZE)),
+    transforms.ToTensor()
+])
+
+
+def load_models():
+    bird = Bird_CNN(c_in=3, c_hidden=settings.BIRD_C_HIDDEN, c_out=len(BirdSpecies))
+    models["bird"] = load_checkpoint(bird, settings.BUILD_PATH / settings.BIRD_MODEL_FILE, device)
+
+    face = Yolo_model(c_in=3, boxes=1, grid=settings.FACE_GRID, labels=1, c_hidden=settings.FACE_C_HIDDEN,
+                      size_activation=settings.FACE_SIZE_ACTIVATION)
+    models["face"] = load_checkpoint(face, settings.BUILD_PATH / settings.FACE_MODEL_FILE, device)
+
+
+class InvalidImage(Exception):
+    pass
+
+
+def open_image(data: bytes) -> tuple[Image.Image, tuple[int, int]]:
+    """Open lazily, reject oversized images before decoding, decode JPEGs at
+    reduced scale. Returns the image and its ORIGINAL (width, height)."""
+    try:
+        image = Image.open(io.BytesIO(data))
+        original_size = image.size
+        if original_size[0] * original_size[1] > settings.MAX_IMAGE_PIXELS:
+            raise InvalidImage("Image too large")
+        image.draft("RGB", settings.DECODE_DRAFT_SIZE)
+        return image.convert("RGB"), original_size
+    except InvalidImage:
+        raise
+    except Exception as e:  # PIL raises many different types for bad data
+        raise InvalidImage("Not a valid image") from e
+
+
+def classify_bird(data: bytes) -> dict:
+    image, _ = open_image(data)
+    tensor = transform_bird(image).unsqueeze(0).to(device)
+    with torch.inference_mode():
+        probs = F.softmax(models["bird"](tensor), dim=1)
+        confidence, cls = torch.max(probs, dim=1)
+    return {
+        "class": BirdSpecies(cls.item()).name,
+        "confidence": confidence.item()
+    }
+
+
+def detect_faces(data: bytes) -> list:
+    image, (W, H) = open_image(data)
+    scale_w = W / settings.FACE_IMAGE_SIZE
+    scale_h = H / settings.FACE_IMAGE_SIZE
+
+    tensor = transform_face(image).unsqueeze(0).to(device)
+    with torch.inference_mode():
+        pred = models["face"](tensor)
+        bboxes, _, _ = convert_prediction(pred.squeeze(0), tensor.squeeze(0), threshold=settings.FACE_THRESHOLD)
+
+    return [scale_box(b, scale_w, scale_h) for b in bboxes]