DeepActionPotential/DrowSeeAi

1import torch
2import torchvision.transforms as transforms
3from PIL import Image
4
5
6val_test_transform = transforms.Compose([
7    transforms.Resize((224, 224)),
8    transforms.ToTensor(),
9  
10])
11
12
13
14def load_model(model_path: str):
15    """
16    Load a trained PyTorch model from disk (saved via torch.save(model, path))
17    and set it to eval() mode.
18
19    Args:
20        model_path (str): Path to the .pth or .pt file containing your trained model.
21
22    Returns:
23        torch.nn.Module: The loaded model in eval mode (on CPU).
24    """
25
26
27    model = torch.load(
28        model_path,
29        map_location=torch.device("cpu"),
30        weights_only=False,   # Allow loading the entire saved model object
31    )
32    model.eval()
33    return model
34
35
36
37def predict(model: torch.nn.Module, image: Image.Image) -> int:
38    """
39    Given a loaded model and a PIL.Image, return 0 (not drowsy) or 1 (drowsy).
40
41    Args:
42        model (torch.nn.Module): Your trained PyTorch model in eval() mode.
43        image (PIL.Image.Image): A PIL image (RGB) of a human face.
44
45    Returns:
46        int: 0 if non-drowsy, 1 if drowsy.
47    """
48    # Apply the validation/test transform:
49    image_tensor = val_test_transform(image)        # [3, 224, 224]
50    image_tensor = image_tensor.unsqueeze(0)        # [1, 3, 224, 224]
51
52    with torch.no_grad():
53        outputs = model(image_tensor)               # assume shape [1, 2] or [1, 1]
54        # If your model outputs two logits (for classes 0 vs 1):
55        if outputs.dim() == 2 and outputs.shape[1] == 2:
56            # e.g. softmax‐based two‐class output
57            _, predicted = torch.max(outputs, dim=1)
58            return int(predicted.item())
59        else:
60            # If your model outputs a single logit (e.g. using `nn.Linear(…) -> [1, 1]`):
61            # apply a sigmoid threshold of 0.5
62            prob = torch.sigmoid(outputs).item()
63            return 1 if prob >= 0.5 else 0
64