86 lines
2.9 KiB
Python
86 lines
2.9 KiB
Python
"""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]
|