Update app.py

app.py

@@ -13,32 +13,62 @@ import bz2
 import shutil
 from efficientnet_pytorch import EfficientNet
 
-# Define paths
-SHAPE_PREDICTOR_URL = "http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2"
+# Define paths and URLs
+SHAPE_PREDICTOR_URL = "https://github.com/italojs/facial-landmarks-recognition/raw/master/shape_predictor_68_face_landmarks.dat.bz2"
 SHAPE_PREDICTOR_PATH = "shape_predictor_68_face_landmarks.dat"
-MODEL_WEIGHTS_PATH = "quad_stream_model_rafdb.pth"  # Update if weights are in a different path
+MODEL_WEIGHTS_URL = "https://huggingface.co/Sakibrumu/Quad_Stream_Face_Emotion_Classifier/resolve/main/quad_stream_model_rafdb.pth"
+MODEL_WEIGHTS_PATH = "quad_stream_model_rafdb.pth"
 
-# Download and extract shape predictor if not present
+# Download shape predictor if not present
 def download_shape_predictor():
     if not os.path.exists(SHAPE_PREDICTOR_PATH):
         print("Downloading shape predictor...")
-        response = requests.get(SHAPE_PREDICTOR_URL, stream=True)
-        with open("shape_predictor_68_face_landmarks.dat.bz2", "wb") as f:
-            f.write(response.content)
-        print("Extracting shape predictor...")
-        with bz2.BZ2File("shape_predictor_68_face_landmarks.dat.bz2", "rb") as f_in:
-            with open(SHAPE_PREDICTOR_PATH, "wb") as f_out:
-                shutil.copyfileobj(f_in, f_out)
-        os.remove("shape_predictor_68_face_landmarks.dat.bz2")
-        print("Shape predictor ready.")
+        try:
+            response = requests.get(SHAPE_PREDICTOR_URL, stream=True, timeout=30)
+            response.raise_for_status()
+            with open("shape_predictor_68_face_landmarks.dat.bz2", "wb") as f:
+                f.write(response.content)
+            print("Extracting shape predictor...")
+            with bz2.BZ2File("shape_predictor_68_face_landmarks.dat.bz2", "rb") as f_in:
+                with open(SHAPE_PREDICTOR_PATH, "wb") as f_out:
+                    shutil.copyfileobj(f_in, f_out)
+            os.remove("shape_predictor_68_face_landmarks.dat.bz2")
+            print("Shape predictor ready.")
+        except Exception as e:
+            print(f"Failed to download or extract shape predictor: {e}")
+            raise RuntimeError("Shape predictor download failed.")
     else:
         print("Shape predictor already exists.")
 
 download_shape_predictor()
 
+# Download model weights from Hugging Face Model Hub
+def download_model_weights():
+    if not os.path.exists(MODEL_WEIGHTS_PATH):
+        print(f"Downloading model weights from {MODEL_WEIGHTS_URL}...")
+        try:
+            response = requests.get(MODEL_WEIGHTS_URL, stream=True, timeout=30)
+            response.raise_for_status()
+            with open(MODEL_WEIGHTS_PATH, "wb") as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    if chunk:
+                        f.write(chunk)
+            print("Model weights downloaded successfully.")
+        except Exception as e:
+            print(f"Failed to download model weights: {e}")
+            raise RuntimeError("Model weights download failed.")
+    else:
+        print("Model weights already exist locally.")
+
+download_model_weights()
+
 # Initialize Dlib detector and predictor
-detector = dlib.get_frontal_face_detector()
-predictor = dlib.shape_predictor(SHAPE_PREDICTOR_PATH)
+try:
+    detector = dlib.get_frontal_face_detector()
+    predictor = dlib.shape_predictor(SHAPE_PREDICTOR_PATH)
+except Exception as e:
+    print(f"Error initializing Dlib: {e}")
+    raise RuntimeError("Failed to initialize Dlib.")
 
 # Class mapping for RAF-DB
 class_mapping = {
@@ -119,7 +149,7 @@ def extract_landmark_features(image):
     features.append(angle)
     
     mouth_center = ((key_points['mouth_left'][0] + key_points['mouth_right'][0]) / 2, 
-                    (key_points['mouth_left'][1] + key_points['mouth_right'][1]) / 2)
+                    (key_points['mouth_left'][1] - key_points['mouth_right'][1]) / 2)
     mouth_to_left_eye = np.sqrt((mouth_center[0] - key_points['left_eye'][0])**2 + 
                                 (mouth_center[1] - key_points['left_eye'][1])**2)
     mouth_to_right_eye = np.sqrt((mouth_center[0] - key_points['right_eye'][0])**2 + 
@@ -182,7 +212,7 @@ def get_landmark_mask(image, target_size=(7, 7)):
     mask = np.clip(mask, 0, 1)
     return mask
 
-# Model definitions
+# Model definitions (unchanged)
 class EfficientNetBackbone(nn.Module):
     def __init__(self):
         super(EfficientNetBackbone, self).__init__()
@@ -362,14 +392,12 @@ class QuadStreamHLAViT(nn.Module):
 
 # Load model
 model = QuadStreamHLAViT(num_classes=7)
-if os.path.exists(MODEL_WEIGHTS_PATH):
-    try:
-        model.load_state_dict(torch.load(MODEL_WEIGHTS_PATH, map_location=torch.device('cpu'), weights_only=True))
-        print("Model weights loaded successfully.")
-    except Exception as e:
-        print(f"Error loading model weights: {e}")
-else:
-    print(f"Model weights not found at {MODEL_WEIGHTS_PATH}. Please upload the weights.")
+try:
+    model.load_state_dict(torch.load(MODEL_WEIGHTS_PATH, map_location=torch.device('cpu'), weights_only=True))
+    print("Model weights loaded successfully.")
+except Exception as e:
+    print(f"Error loading model weights: {e}")
+    raise RuntimeError("Failed to load model weights.")
 model.eval()
 
 # Inference function