predictor/predict.py

import torch
from torchvision import transforms
from config import DEVICE
import gradio as gr


def predict_faces(model, faces, progress=gr.Progress(track_tqdm=True)):
    transform = transforms.Compose([
    transforms.Resize((299, 299)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

    predictions = []
    for face in progress.tqdm(faces, desc="Running predictions"):
        face = transform(face).unsqueeze(0).to(DEVICE)
        with torch.no_grad():
            pred = model(face).item()
        predictions.append(pred)
    return predictions