Update app.py

app.py

@@ -3,38 +3,56 @@ import numpy as np
 import gradio as gr
 
 # =========================
-# TensorFlow / Keras
+# TensorFlow / Keras Setup
 # =========================
 import tensorflow as tf
 from tensorflow.keras.applications.resnet50 import preprocess_input as resnet_preprocess
 
-# GPU config for TensorFlow (enable if available)
-gpus = tf.config.experimental.list_physical_devices("GPU")
-if gpus:
-    try:
-        for g in gpus:
-            tf.config.experimental.set_memory_growth(g, True)
-        # Mixed precision can speed up on modern GPUs
-        try:
-            tf.keras.mixed_precision.set_global_policy("mixed_float16")
-            print("[INFO] TF mixed precision enabled.")
-        except Exception as e:
-            print(f"[WARN] Could not enable mixed precision: {e}")
-    except Exception as e:
-        print(f"[WARN] Could not set memory growth: {e}")
+# Limit TensorFlow thread usage (prevents OOM)
+os.environ["OMP_NUM_THREADS"] = "2"
+os.environ["MKL_NUM_THREADS"] = "2"
+os.environ["OPENBLAS_NUM_THREADS"] = "2"
+os.environ["NUMEXPR_NUM_THREADS"] = "2"
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
+
+# Try to enable GPU safely for TensorFlow
+try:
+    gpus = tf.config.experimental.list_physical_devices("GPU")
+    if gpus:
+        for gpu in gpus:
+            tf.config.experimental.set_memory_growth(gpu, True)
+        tf.keras.mixed_precision.set_global_policy("mixed_float16")
+        print("[INFO] ✅ TensorFlow GPU detected & configured (mixed precision ON).")
+    else:
+        print("[WARN] ⚠️ No TensorFlow GPU detected, running on CPU.")
+except Exception as e:
+    print(f"[WARN] TensorFlow GPU init skipped: {e}")
 
 # =========================
-# Vision stack
+# PyTorch / MTCNN Setup
 # =========================
-import cv2
 import torch
 from facenet_pytorch import MTCNN
+import cv2
 import dlib
 from imutils import face_utils
 from scipy.spatial import distance as dist
 
 warnings.filterwarnings("ignore")
 
+# Use GPU for MTCNN if available
+if torch.cuda.is_available():
+    _torch_device = torch.device("cuda")
+    torch.backends.cudnn.benchmark = True
+    try:
+        torch.set_float32_matmul_precision("medium")
+    except Exception:
+        pass
+    print("[INFO] ✅ PyTorch GPU detected — MTCNN will use CUDA.")
+else:
+    _torch_device = torch.device("cpu")
+    print("[WARN] ⚠️ No GPU detected — MTCNN running on CPU.")
+
 # =========================
 # Paths / Config
 # =========================
@@ -42,33 +60,18 @@ VIDEO_MODEL_PATH = "models/video_model.h5"
 DLIB_LANDMARK_MODEL = "shape_predictor_68_face_landmarks.dat"
 
 IMG_SIZE = (224, 224)
-FRAME_STEP = 5             # every 5th frame
-NUM_MAX_FACES = 300        # cap on faces collected
-
-# Blink (EAR) features
+FRAME_STEP = 5
+NUM_MAX_FACES = 300
 EAR_THRESHOLD = 0.25
 EAR_CONSEC_FRAMES = 3
-
-# Prediction threshold
 PRED_THRESHOLD = 0.5
 
-# =========================
-# Lazy-loaded state
-# =========================
+# Lazy-loaded globals
 _video_model = None
 _mtcnn = None
 _dlib_detector = None
 _dlib_predictor = None
 
-# Torch / MTCNN device selection
-_torch_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-if torch.cuda.is_available():
-    torch.backends.cudnn.benchmark = True
-    try:
-        torch.set_float32_matmul_precision("medium")
-    except Exception:
-        pass
-print(f"[INFO] PyTorch device: {_torch_device}")
 
 def lazy_load():
     global _video_model, _mtcnn, _dlib_detector, _dlib_predictor