ktejeshnaidu commited on
Commit
82d5d02
·
verified ·
1 Parent(s): 7487330

Delete model.py

Browse files
Files changed (1) hide show
  1. model.py +0 -67
model.py DELETED
@@ -1,67 +0,0 @@
1
- import torch
2
- import torch.nn as nn
3
- import torch.nn.functional as F
4
- import pickle
5
- from torchvision import transforms
6
-
7
- import numpy as np
8
- from PIL import Image
9
-
10
-
11
- class FaceClassifier(nn.Module):
12
- def __init__(self, num_classes):
13
- super().__init__()
14
- self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
15
- self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
16
- self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
17
- self.pool = nn.MaxPool2d(2, 2)
18
- self.dropout = nn.Dropout(0.1)
19
- self.fc1 = nn.Linear(128 * 16 * 16, 512)
20
- self.fc2 = nn.Linear(512, num_classes)
21
-
22
- def forward(self, x):
23
- x = self.pool(F.relu(self.conv1(x)))
24
- x = self.pool(F.relu(self.conv2(x)))
25
- x = self.pool(F.relu(self.conv3(x)))
26
- x = x.view(-1, 128 * 16 * 16)
27
- x = self.dropout(F.relu(self.fc1(x)))
28
- return self.fc2(x)
29
-
30
-
31
- class EmotionPredictor:
32
- def __init__(self):
33
- self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
34
-
35
- with open("classes.pkl", "wb") as f:
36
- self.classes = pickle.load(f)
37
-
38
- self.model = FaceClassifier(len(self.classes))
39
- self.model.load_state_dict(
40
- torch.load("face_classifier.pth", map_location=self.device, weights_only=True)
41
- )
42
- self.model.to(self.device).eval()
43
-
44
- self.transform = transforms.Compose([
45
- transforms.Resize((128, 128)),
46
- transforms.ToTensor(),
47
- transforms.Normalize((0.5,), (0.5,)),
48
- ])
49
-
50
- @torch.inference_mode()
51
- def predict(self, image_np: np.ndarray) -> str:
52
- """Return the top predicted emotion label."""
53
- img = Image.fromarray(image_np)
54
- tensor = self.transform(img).unsqueeze(0).to(self.device)
55
- output = self.model(tensor)
56
- return self.classes[output.argmax(1).item()]
57
-
58
- @torch.inference_mode()
59
- def predict_with_confidence(self, image_np: np.ndarray) -> tuple[str, dict]:
60
- """Return (top_label, {label: confidence_float}) using softmax probabilities."""
61
- img = Image.fromarray(image_np)
62
- tensor = self.transform(img).unsqueeze(0).to(self.device)
63
- logits = self.model(tensor)
64
- probs = F.softmax(logits, dim=1).squeeze(0).cpu().tolist()
65
- scores = {cls: round(p, 4) for cls, p in zip(self.classes, probs)}
66
- top = max(scores, key=scores.get)
67
- return top, scores