Upload 3 files

face_facts.py

@@ -0,0 +1,155 @@
+import cv2
+import numpy as np
+import tempfile
+import time
+import streamlit as st
+from PIL import Image 
+from io import BytesIO
+import torch
+from ultralytics import YOLO
+from utils import Ultimate_Lightning, age_lightning, gender_lightning, race_lightning
+import pandas as pd
+import mediapipe as mp
+from torch.cuda import is_available as gpu_ready
+from mediapipe.tasks import python
+from mediapipe.tasks.python import vision
+
+torch.manual_seed(42)
+
+DEMO_IMAGE = 'demo.JPG'
+GENDER_DICT = {
+    1: 'Female',
+    0: 'Male'
+}
+RACE_DICT = {
+    0: 'White', 
+    1: 'Black',
+    2: 'Asian',
+    3: 'Indian',
+    4: 'Others'
+}
+device = 'cuda' if gpu_ready() else 'cpu'
+
+def load_face_detector():
+  base_options = python.BaseOptions(model_asset_path='models/detector.tflite')
+  options = vision.FaceDetectorOptions(base_options=base_options)
+  detector = vision.FaceDetector.create_from_options(options)
+  return detector
+
+@st.cache_data
+def load_model():
+  joint_model = Ultimate_Lightning()
+  age_model = age_lightning()
+  age_model.load_state_dict(torch.load('models/age_model.pth'))
+  race_model = race_lightning()
+  race_model.load_state_dict(torch.load('models/race_model.pth'))  
+  gender_model = gender_lightning()
+  gender_model.load_state_dict(torch.load('models/gender_model.pth'))
+  joint_model.load_state_dict(torch.load('models/joint_model.pth'))
+
+  return age_model, gender_model, race_model, joint_model
+
+st.set_page_config(page_title="Face-Facts")
+st.title('Face-Facts')
+
+app_mode = st.sidebar.selectbox('Choose Page', ['About the App', 'Run Face Facts'])
+
+st.markdown(
+      """
+      <style>
+          [data-testid = 'stSidebar'][aria-expanded = 'true'] > div:first-child{
+            width: 350px
+          }
+          [data-testid = 'stSidebar'][aria-expanded = 'false'] > div:first-child{
+            width: 350px
+            margin-left: -350px
+          }
+      </style>
+      """, unsafe_allow_html = True
+)
+
+if app_mode == 'About the App':
+  st.markdown('')
+
+elif app_mode == 'Run Face Facts':
+  
+  age_model, gender_model, race_model, joint_model = load_model()
+  detector = load_face_detector()
+
+  st.sidebar.markdown('---')
+  use_single_model = st.sidebar.checkbox('Use single model', value = False)
+  
+  kpi1, age_col, gender_col, race_col, kpi5 = st.columns(5)
+  with age_col:
+    st.markdown('**Age**')
+    age_text = st.markdown('0')
+  with gender_col:
+    st.markdown('**Gender**')
+    gender_text = st.markdown('0')
+  with race_col:
+    st.markdown('**Race**')
+    race_text = st.markdown('0')
+
+  img_file_buffer = st.sidebar.file_uploader('Upload an Image', type = ['jpg', 'png', 'jpeg'])
+  if img_file_buffer:
+    buffer = BytesIO(img_file_buffer.read())
+    data = np.frombuffer(buffer.getvalue(), dtype=np.uint8)
+    image_orig = cv2.imdecode(data, cv2.IMREAD_COLOR)
+  else:
+    demo_image = DEMO_IMAGE
+    image_orig = cv2.imread(demo_image, cv2.IMREAD_COLOR)
+  
+  st.sidebar.text('Original Image')
+  st.sidebar.image(image_orig, channels = 'BGR')
+
+  image = cv2.cvtColor(image_orig, cv2.COLOR_BGR2RGB)
+  image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image)
+  detection_result = detector.detect(image)
+  if detection_result.detections != []:
+    res = detection_result.detections[0].bounding_box
+    x, y, w, h = res.origin_x, res.origin_y, res.width, res.height
+    image = image.numpy_view()[y:y+h, x:x+w]
+    image = cv2.resize(image, (200, 200))
+    image = torch.from_numpy(image).permute(2, 0, 1) / 255.
+    age_model.eval()
+    gender_model.eval()
+    race_model.eval()
+    joint_model.eval()
+    with torch.no_grad():
+      if use_single_model:
+        age_pred, gender_pred, race_pred = joint_model(image.unsqueeze(0))
+      else:
+        age_pred, gender_pred, race_pred = age_model(image.unsqueeze(0)), gender_model(image.unsqueeze(0)), race_model(image.unsqueeze(0))
+    age = int(age_pred.item())
+    gender = GENDER_DICT[int(gender_pred.item() > 0.5)]
+    race = RACE_DICT[race_pred.argmax(dim = 1).item()] 
+    gender_emoji = '♂️' if gender == 'Male' else '♀️'
+    gender_color = 'blue' if gender == 'Male' else 'pink'
+  else:
+    age = '-'
+    gender = '-'
+    gender_emoji = '-'
+    gender_color = ''  
+    race = '-'
+    st.error('No face detected in the image')
+  age_text.write(f'<h1> {age} </h1>', unsafe_allow_html = True)
+  gender_text.write(f"<h1 style='color: {gender_color};'>{gender_emoji}</h1>", unsafe_allow_html=True)
+  race_text.write(f"<h1> {race} </h1>", unsafe_allow_html = True)
+
+  st.markdown('---')
+
+  if not age == '-':
+    with st.expander('🔻 More Details 🔻'):
+      gender_decimal = gender_pred.item() if gender == 'Female' else abs(1 - gender_pred.item())
+      gender_precentage = f'{100 * gender_decimal:.2f}%'
+      st.write('---')
+      cols = st.columns(2)
+      with cols[0]:
+        st.markdown("<h3 style='color: gray;'>Face of Interest</h3>", unsafe_allow_html=True)
+        st.image(image.permute(1, 2, 0).numpy(), channels = 'RGB', use_column_width = True)
+      with cols[1]:
+        st.write(f"<h3 style = 'color: {gender_color};'> {gender}; {gender_precentage} Probability </h3>", unsafe_allow_html = True)
+        st.progress(gender_decimal)
+        st.write('---')
+        st.bar_chart(pd.DataFrame({'Probability': race_pred.squeeze().numpy(), 'Race': list(RACE_DICT.values())}), x = 'Race', y = 'Probability')
+      st.write('---')

utils.py

@@ -0,0 +1,180 @@
+import torch.nn as nn
+import torch
+import pytorch_lightning as pl
+import torch.nn.functional as F
+class CustomModelMain(nn.Module):
+    def __init__(self, problem_type, n_classes):
+        super().__init__()
+        if problem_type == 'Classification' and n_classes == 1:
+            output = nn.Sigmoid()
+        elif problem_type == 'Regression' and n_classes == 1:
+            output = nn.ReLU()
+        elif problem_type == 'Classification' and n_classes > 1:
+            output = nn.Softmax(dim = 1)
+        self.conv1 = nn.Conv2d(in_channels = 3, out_channels = 32, kernel_size = 3, padding = 1)
+        self.pool1 = nn.MaxPool2d(kernel_size = 2, stride = 2)
+        self.conv2 = nn.Conv2d(in_channels = 32, out_channels = 64, kernel_size = 3, padding = 1)
+        self.pool2 = nn.MaxPool2d(kernel_size = 2, stride = 2)
+        self.conv3 = nn.Conv2d(in_channels = 64, out_channels = 64, kernel_size = 3, padding = 1)
+        self.pool3 = nn.MaxPool2d(kernel_size = 2, stride = 2)
+        self.flatten = nn.Flatten()
+        self.fc1 = nn.Linear(64 * 25 * 25 , 128)
+        self.relu = nn.ReLU()
+        self.dropout = nn.Dropout(p = 0.5)
+        self.fc2 = nn.Linear(128, n_classes)
+        self.output = output
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.pool1(x)
+        x = self.relu(x)
+        x = self.conv2(x)
+        x = self.pool2(x)
+        x = self.relu(x)
+        x = self.conv3(x)
+        x = self.pool3(x)
+        x = self.relu(x)
+        x = self.flatten(x)
+        x = self.fc1(x)
+        x = self.relu(x)
+        x = self.dropout(x)
+        x = self.fc2(x)
+        x = self.output(x)
+        return x
+class age_lightning(pl.LightningModule):
+    def __init__(self):
+        super().__init__()
+        self.model = CustomModelMain('Regression', 1)
+    def forward(self, x):
+        return self.model(x)
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y = y[:, 0]
+        y_hat = self(x)
+        loss = F.mse_loss(y_hat, y.unsqueeze(-1).float())
+        acc = torch.eq((y_hat > 0.5).int().to(torch.int64), y.unsqueeze(-1).int()).all(dim=1).sum() / len(y)
+        self.log('train loss', loss, prog_bar = True)
+        return loss
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        y_val = y[:, 0]
+        y_hat = self(x)
+        loss = F.mse_loss(y_hat, y_val.unsqueeze(-1).float())
+        self.log('valid loss', loss, prog_bar = True)
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=1e-4)
+class gender_lightning(pl.LightningModule):
+    def __init__(self):
+        super().__init__()
+        self.model = CustomModelMain('Classification', 1)
+    def forward(self, x):
+        return self.model(x)
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y = y[:, 1]
+        y_hat = self(x)
+        loss = F.binary_cross_entropy(y_hat, y.unsqueeze(-1).float())
+        acc = torch.eq((y_hat > 0.5).int().to(torch.int64), y.unsqueeze(-1).int()).all(dim=1).sum() / len(y)
+        self.log('train loss', loss, prog_bar = True)
+        self.log('accuracy', acc, prog_bar = True)
+        return loss
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        y_val = y[:, 1]
+        y_hat = self(x)
+        loss = F.binary_cross_entropy_with_logits(y_hat, y_val.unsqueeze(-1).float())
+        acc =  torch.eq((y_hat > 0.5).int().to(torch.int64), y_val.unsqueeze(-1).int()).all(dim=1).sum() / len(y_val)
+        self.log('valid loss', loss, prog_bar = True)
+        self.log('val accuracy', acc, prog_bar = True)
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=1e-4)
+class race_lightning(pl.LightningModule):
+    def __init__(self):
+        super().__init__()
+        self.model = CustomModelMain('Classification', 5)
+    def forward(self, x):
+        return self.model(x)
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y = y[:, 2]
+        y_hat = self(x)
+        y_oh = F.one_hot(y, num_classes = 5)
+        loss = F.cross_entropy(y_hat.log(), y_oh.float())
+        preds = y_hat.argmax(dim = 1)
+        acc = torch.eq(y, preds).float().mean()
+        self.log('train loss', loss, prog_bar = True)
+        self.log('accuracy', acc, prog_bar = True)
+        return loss
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        y_val = y[:, 2]
+        y_hat = self(x)
+        y_oh = F.one_hot(y_val, num_classes = 5)
+        loss = F.cross_entropy(y_hat, y_oh.float())
+        preds = y_hat.argmax(dim = 1)
+        acc = torch.eq(y_val, preds).float().mean()
+        self.log('valid loss', loss, prog_bar = True)
+        self.log('val accuracy', acc, prog_bar = True)
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=1e-4)
+class Ultimate_Lightning(pl.LightningModule):
+  def __init__(self):
+    super().__init__()
+    self.age_model = CustomModelMain('Regression', 1)
+    self.gender_model = CustomModelMain('Classification', 1)
+    self.race_model = CustomModelMain('Classification', 5)
+  def forward(self, x):
+    return self.age_model(x), self.gender_model(x), self.race_model(x)
+  def training_step(self, batch, batch_idx):
+    x, y = batch
+    y_age, y_gender, y_race = y[:, 0], y[:, 1], y[:, 2]
+    y_hat_age, y_hat_gender, y_hat_race = self(x)
+    
+    age_loss = F.mse_loss(y_hat_age, y_age.unsqueeze(-1).float())
+    age_acc = torch.eq((y_hat_age > 0.5).int().to(torch.int64), y_age.unsqueeze(-1).int()).all(dim=1).sum() / len(y_age)
+    
+    gender_loss = F.binary_cross_entropy(y_hat_gender, y_gender.unsqueeze(-1).float())
+    gender_acc = torch.eq((y_hat_gender > 0.5).int().to(torch.int64), y_gender.unsqueeze(-1).int()).all(dim=1).sum() / len(y_gender)
+
+    y_race_oh = F.one_hot(y_race, num_classes = 5)
+    race_loss = F.cross_entropy(y_hat_race.log(), y_race_oh.float())
+    race_preds = y_hat_race.argmax(dim = 1)
+    race_acc = torch.eq(y_race, race_preds).float().mean()
+    
+    total_loss = (0.001 * age_loss) + gender_loss + race_loss
+
+    self.log('age loss', age_loss, prog_bar = True)
+    self.log('gender loss', gender_loss, prog_bar = True)
+    self.log('race loss', race_loss, prog_bar = True)
+    self.log('gender acc', gender_acc, prog_bar = True)
+    self.log('race acc', race_acc, prog_bar = True)
+    self.log('total loss', total_loss, prog_bar = True)
+    
+    return total_loss
+      
+  def validation_step(self, batch, batch_idx):
+    x, y = batch
+    y_age, y_gender, y_race = y[:, 0], y[:, 1], y[:, 2]
+    y_hat_age, y_hat_gender, y_hat_race = self(x)
+    
+    age_loss = F.mse_loss(y_hat_age, y_age.unsqueeze(-1).float())
+    age_acc = torch.eq((y_hat_age > 0.5).int().to(torch.int64), y_age.unsqueeze(-1).int()).all(dim=1).sum() / len(y_age)
+    
+    gender_loss = F.binary_cross_entropy(y_hat_gender, y_gender.unsqueeze(-1).float())
+    gender_acc = torch.eq((y_hat_gender > 0.5).int().to(torch.int64), y_gender.unsqueeze(-1).int()).all(dim=1).sum() / len(y_gender)
+
+    y_race_oh = F.one_hot(y_race, num_classes = 5)
+    race_loss = F.cross_entropy(y_hat_race.log(), y_race_oh.float())
+    race_preds = y_hat_race.argmax(dim = 1)
+    race_acc = torch.eq(y_race, race_preds).float().mean()
+    
+    total_loss = (0.001 * age_loss) + gender_loss + race_loss
+
+    self.log('val age loss', age_loss, prog_bar = True)
+    
+    self.log('val gender acc', gender_acc, prog_bar = True)
+    
+    self.log('val race acc', race_acc, prog_bar = True)
+    
+  def configure_optimizers(self):
+    return torch.optim.Adam(self.parameters(), lr=1e-4)
+