Sebahadin1234/Facial_Expression_Recognition_Deep_Learning
0
1import streamlit as st2import torch3import torch.nn as nn4import torch.nn.functional as F5from torchvision import models, transforms6from PIL import Image7 8 9 10 11from pytorch_grad_cam import GradCAM12from pytorch_grad_cam.utils.image import show_cam_on_image13 14import numpy as np15 16 17 18 19 20# -------------------------------21# Label Dictionary (1-indexed)22# -------------------------------23label_dict = {24 1: 'Surprise',25 2: 'Disgust',26 3: 'Happiness',27 4: 'Sadness',28 5: 'Anger',29 6: 'Neutral'30}31 32 33 34import torch35from torch.utils.data import DataLoader36from torchvision import datasets, transforms37 38 39 40 41 42# Parameters43batch_size = 6444 45 46 47 48img_size = 100 # Updated from 48 to 10049 50# Transforms for CNN51transform_train = transforms.Compose([52 transforms.Resize((img_size, img_size)), # Resize to 100x10053 transforms.RandomHorizontalFlip(), # Data augmentation54 transforms.RandomRotation(degrees=10), # Data augmentation55 transforms.ToTensor(), # Convert to tensor56 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) # Normalize RGB channels57])58 59transform_test = transforms.Compose([60 transforms.Resize((img_size, img_size)),61 transforms.ToTensor(),62 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))63])64 65 66 67 68# Transforms for VGG and ViT69transform = transforms.Compose([70 transforms.Resize((224, 224)), # Resize to 224x22471 transforms.ToTensor(), # Convert to tensor [0,1]72 # transforms.RandomRotation(9),73 transforms.Normalize( # Normalize using ImageNet stats74 mean=[0.485, 0.456, 0.406],75 std=[0.229, 0.224, 0.225]76 )77])78 79 80 81 82 83 84# Then unzip it85import zipfile86import os87 88with zipfile.ZipFile('dataset_final.zip', 'r') as zip_ref:89 zip_ref.extractall('Dataset_final')90 91 92#datasets93train_dataset_cnn = datasets.ImageFolder(root='Dataset_final/train', transform=transform_train)94test_dataset_cnn = datasets.ImageFolder(root='Dataset_final/test', transform=transform_test)95 96 97train_dataset_v = datasets.ImageFolder(root='Dataset_final/train', transform=transform)98test_dataset_v = datasets.ImageFolder(root='Dataset_final/test', transform=transform)99 100 101 102 103 104# DataLoaders105train_loader_cnn = DataLoader(train_dataset_cnn, batch_size=batch_size, shuffle=True, num_workers=2)106test_loader_cnn = DataLoader(test_dataset_cnn, batch_size=batch_size, shuffle=False, num_workers=2)107 108 109train_loader_v = DataLoader(train_dataset_v, batch_size=batch_size, shuffle=True, num_workers=2)110test_loader_v = DataLoader(test_dataset_v, batch_size=batch_size, shuffle=False, num_workers=2)111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141# -------------------------------142# Model: CNN (your custom model)143# -------------------------------144class FacialReaction(nn.Module):145 def __init__(self, num_classes=7):146 super(FacialReaction, self).__init__()147 self.conv1 = nn.Conv2d(3, 64, kernel_size=4, padding=1)148 self.pool = nn.MaxPool2d(2, 2)149 self.conv2 = nn.Conv2d(64, 64, kernel_size=4, padding=1)150 self.fc1 = nn.Linear(64 * 24 * 24, 128)151 self.fc2 = nn.Linear(128, num_classes)152 153 def forward(self, x):154 x = self.pool(F.relu(self.conv1(x)))155 x = self.pool(F.relu(self.conv2(x)))156 x = x.view(x.size(0), -1)157 x = F.relu(self.fc1(x))158 return self.fc2(x)159 160 161 162 163 164 165 166 167 168 169 170 171 172class PatchEmbed(nn.Module):173 def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=192):174 super(PatchEmbed, self).__init__()175 self.img_size = img_size176 self.patch_size = patch_size177 self.embed_dim = embed_dim178 179 # Create a convolutional layer for patch embedding180 self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)181 182 def forward(self, x):183 x = self.proj(x) # (B, embed_dim, H/patch_size, W/patch_size)184 x = x.flatten(2) # Flatten (B, embed_dim, N_patches)185 x = x.transpose(1, 2) # (B, N_patches, embed_dim)186 return x187 188class MultiHeadSelfAttention(nn.Module):189 def __init__(self, embed_dim, num_heads):190 super(MultiHeadSelfAttention, self).__init__()191 self.num_heads = num_heads192 self.embed_dim = embed_dim193 self.head_dim = embed_dim // num_heads194 195 assert self.head_dim * num_heads == embed_dim, "Embedding dimension must be divisible by num_heads"196 197 self.qkv = nn.Linear(embed_dim, embed_dim * 3)198 self.fc_out = nn.Linear(embed_dim, embed_dim)199 200 def forward(self, x):201 B, N, E = x.shape202 203 # Linear transformation to get queries, keys and values204 qkv = self.qkv(x)205 qkv = qkv.reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)206 207 q, k, v = qkv[0], qkv[1], qkv[2]208 209 # Attention calculation210 energy = torch.einsum("nqhd,nkhd->nhqk", [q, k])211 attention = torch.softmax(energy / (self.embed_dim ** 0.5), dim=-1)212 213 out = torch.einsum("nhql,nlhd->nqhd", [attention, v]).reshape(B, N, E)214 out = self.fc_out(out)215 return out216 217class FeedForward(nn.Module):218 def __init__(self, embed_dim, hidden_dim=768):219 super(FeedForward, self).__init__()220 self.fc1 = nn.Linear(embed_dim, hidden_dim)221 self.fc2 = nn.Linear(hidden_dim, embed_dim)222 self.dropout = nn.Dropout(0.1)223 224 def forward(self, x):225 x = self.fc1(x)226 x = F.gelu(x)227 x = self.dropout(x)228 x = self.fc2(x)229 return x230 231class TransformerBlock(nn.Module):232 def __init__(self, embed_dim, num_heads, hidden_dim=768):233 super(TransformerBlock, self).__init__()234 self.attn = MultiHeadSelfAttention(embed_dim, num_heads)235 self.ffn = FeedForward(embed_dim, hidden_dim)236 self.layernorm1 = nn.LayerNorm(embed_dim)237 self.layernorm2 = nn.LayerNorm(embed_dim)238 239 def forward(self, x):240 attn_out = self.attn(x)241 x = self.layernorm1(x + attn_out) # Add & Norm242 ffn_out = self.ffn(x)243 x = self.layernorm2(x + ffn_out) # Add & Norm244 return x245 246class VisionTransformer(nn.Module):247 def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=192, num_heads=3, num_layers=12, num_classes=6):248 super(VisionTransformer, self).__init__()249 self.embed_dim = embed_dim250 251 # Patch Embedding252 self.patch_embed = PatchEmbed(img_size, patch_size, in_channels, embed_dim)253 254 # Positional Encoding255 self.pos_embed = nn.Parameter(torch.randn(1, (img_size // patch_size) ** 2 + 1, embed_dim))256 self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim))257 258 # Transformer Blocks259 self.blocks = nn.ModuleList([260 TransformerBlock(embed_dim, num_heads) for _ in range(num_layers)261 ])262 263 # MLP Head for classification264 self.fc_out = nn.Linear(embed_dim, num_classes)265 266 def forward(self, x):267 # Embed the image into patches268 x = self.patch_embed(x)269 270 # Add class token to the sequence271 batch_size = x.size(0)272 cls_tokens = self.cls_token.expand(batch_size, -1, -1)273 x = torch.cat((cls_tokens, x), dim=1) # (B, N_patches+1, embed_dim)274 275 # Add positional encoding276 x = x + self.pos_embed277 278 # Pass through transformer blocks279 for block in self.blocks:280 x = block(x)281 282 # Classification head283 cls_output = x[:, 0] # Extract the class token output284 out = self.fc_out(cls_output)285 286 return out287 288 289 290 291 292 293 294 295# -------------------------------296# Load Model Functions (correct filenames)297# -------------------------------298@st.cache_resource299def load_cnn_model():300 model = FacialReaction(num_classes=6)301 302 # Load full checkpoint303 checkpoint = torch.load('CNN_facial_reaction.pth',map_location='cpu')304 305 # Load only the model weights306 model.load_state_dict(checkpoint['model_state_dict'])307 308 model.eval()309 return model310 311@st.cache_resource312def load_vgg_model():313 model = models.vgg16(pretrained=False)314 model.classifier[6] = nn.Linear(4096, 6)315 model.load_state_dict(torch.load("vgg_dataset2_84_74.pth", map_location='cpu'))316 model.eval()317 return model318 319@st.cache_resource320def load_vit_model():321 model = VisionTransformer()322 323 model.heads = nn.Sequential(nn.Linear(192,6))324 325 326 checkpoint = torch.load('vit_70_67.pth', map_location='cpu')327 model.load_state_dict(checkpoint['model_state_dict'])328 329 330 331 model.eval()332 return model333 334# -------------------------------335# Preprocess Webcam Image336# -------------------------------337def preprocess_image(img: Image.Image,model_type='CNN'):338 339 if model_type == 'CNN':340 transform = transforms.Compose([341 transforms.Resize((100, 100)),342 transforms.ToTensor(),343 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))344 ])345 return transform(img).unsqueeze(0) # [1, 3, 100, 100]346 347 else:348 transform = transforms.Compose([349 transforms.Resize((224, 224)), # Resize to 224x224350 transforms.ToTensor(), # Convert to tensor [0,1]351 # transforms.RandomRotation(9),352 transforms.Normalize( # Normalize using ImageNet stats353 mean=[0.485, 0.456, 0.406],354 std=[0.229, 0.224, 0.225]355 )356 ])357 return transform(img).unsqueeze(0) # [1, 3, 224, 224]358 359# ------------------------------------------360# ✂️ Tightly crop center of image361# ------------------------------------------362def tight_center_crop(img: Image.Image, crop_ratio: float = 0.7) -> Image.Image:363 """364 Crops a tighter square from the center of the image.365 crop_ratio defines the portion to keep (e.g., 0.7 means 70% of the smaller side).366 """367 width, height = img.size368 side = min(width, height)369 crop_size = int(side * crop_ratio)370 371 left = (width - crop_size) // 2372 top = (height - crop_size) // 2373 right = left + crop_size374 bottom = top + crop_size375 376 return img.crop((left, top, right, bottom))377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394def apply_gradcam_streamlit(model, input_tensor, target_layer, class_names=None, true_label=None):395 """396 Applies Grad-CAM on a given image tensor and returns:397 - Original image398 - Original + Grad-CAM overlay399 400 Args:401 - model: Trained CNN/VGG/ViT model.402 - input_tensor: A single image tensor (1, 3, H, W).403 - target_layer: Target layer for Grad-CAM.404 - class_names: Optional dict mapping class indices to names.405 - true_label: Optional integer ground-truth label (1-indexed).406 407 Returns:408 - Tuple of original image and Grad-CAM overlay (both as NumPy arrays)409 """410 411 model.eval()412 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')413 model = model.to(device)414 input_tensor = input_tensor.to(device)415 416 # GradCAM setup417 cam = GradCAM(model=model, target_layers=[target_layer])418 targets = None419 420 # Run inference421 outputs = model(input_tensor)422 _, predicted = outputs.max(1)423 predicted_label = predicted.item() + 1 # shift from 0–5 to 1–6424 425 # Grad-CAM computation426 grayscale_cam = cam(input_tensor=input_tensor)[0, :] # (H, W)427 428 # Unnormalize and prepare original image429 img_disp = input_tensor.squeeze(0).cpu()430 img_disp = img_disp * 0.5 + 0.5 # Assuming normalization was [-1, 1]431 img_disp = img_disp.permute(1, 2, 0).numpy() # (H, W, C)432 433 # Create heatmap image434 heatmap_image = show_cam_on_image(img_disp, grayscale_cam, use_rgb=True)435 436 # Return both images for display437 return img_disp, heatmap_image, predicted_label438 439