Update app.py

app.py

@@ -21,7 +21,8 @@ required_packages = {
     "numpy": None,
     "gradio": None,
     "mediapipe": None,
-    "tensorflow": None
+    "tensorflow": None,
+    "gitpython": None  # For git operations
 }
 
 installed_packages = {pkg.key for pkg in pkg_resources.working_set}
@@ -38,9 +39,9 @@ import tensorflow as tf
 from tensorflow.keras.preprocessing.image import img_to_array
 from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
 import time
-import os
 from pathlib import Path
 import tempfile
+import git
 
 # Set TensorFlow to use memory growth to avoid consuming all GPU memory
 physical_devices = tf.config.list_physical_devices('GPU')
@@ -58,22 +59,58 @@ mp_drawing = mp.solutions.drawing_utils
 # Global variable for model
 mask_model = None
 
-def load_mask_model(model_path="mask_recog.h5"):
+def download_model_repo():
+    """Download the face mask detection model from GitHub"""
+    repo_url = "https://github.com/misbah4064/face_mask_detection.git"
+    repo_dir = "face_mask_detection"
+    model_path = os.path.join(repo_dir, "mask_recog.h5")
+    
+    # Check if model already exists
+    if os.path.exists(model_path):
+        print(f"Model already exists at {model_path}")
+        return model_path
+    
+    # Check if repository directory exists
+    if os.path.exists(repo_dir):
+        print(f"Repository directory already exists at {repo_dir}")
+    else:
+        print(f"Cloning repository from {repo_url}...")
+        try:
+            git.Repo.clone_from(repo_url, repo_dir)
+            print("Repository cloned successfully")
+        except Exception as e:
+            print(f"Error cloning repository: {e}")
+            # Try alternative method with subprocess
+            try:
+                subprocess.check_call(["git", "clone", repo_url])
+                print("Repository cloned with subprocess")
+            except Exception as sub_e:
+                print(f"Error with subprocess clone: {sub_e}")
+                return None
+    
+    # Verify model file exists
+    if os.path.exists(model_path):
+        print(f"Model file found at {model_path}")
+        return model_path
+    else:
+        print(f"Model file not found at {model_path}")
+        return None
+
+def load_mask_model():
     """Load the mask detection model once and cache it"""
     global mask_model
     if mask_model is None:
         try:
-            # Use TensorFlow Lite if available for better performance
-            if os.path.exists("mask_recog.tflite"):
-                interpreter = tf.lite.Interpreter(model_path="mask_recog.tflite")
-                interpreter.allocate_tensors()
-                mask_model = interpreter
-                print("Loaded TFLite model")
+            # First try to download/access the model from GitHub
+            model_path = download_model_repo()
+            if model_path and os.path.exists(model_path):
+                # Use standard TF model
+                mask_model = tf.keras.models.load_model(model_path)
+                print(f"Loaded {model_path} successfully")
                 return True
-            # Otherwise use standard TF model
-            mask_model = tf.keras.models.load_model(model_path)
-            print(f"Loaded {model_path} successfully")
-            return True
+            else:
+                print("Failed to download or find the model")
+                return False
         except Exception as e:
             print(f"Error loading model: {e}")
             return False
@@ -126,23 +163,8 @@ def predict_mask(face_img):
     face_array = np.expand_dims(face_array, axis=0)
     face_array = preprocess_input(face_array)
     
-    # Check if we're using TFLite
-    if isinstance(mask_model, tf.lite.Interpreter):
-        # Get input and output tensors
-        input_details = mask_model.get_input_details()
-        output_details = mask_model.get_output_details()
-        
-        # Set input tensor
-        mask_model.set_tensor(input_details[0]['index'], face_array.astype(np.float32))
-        
-        # Run inference
-        mask_model.invoke()
-        
-        # Get output
-        preds = mask_model.get_tensor(output_details[0]['index'])
-    else:
-        # Use standard TF model
-        preds = mask_model.predict(face_array, verbose=0)
+    # Use standard TF model
+    preds = mask_model.predict(face_array, verbose=0)
     
     mask_prob = float(preds[0][0])
     return mask_prob > 0.5, mask_prob
@@ -223,7 +245,7 @@ def analyze_frame(frame, face_detector, min_detection_confidence=0.5, blur_thres
 def process_video(video_path, progress=gr.Progress(), min_detection_confidence=0.5, blur_threshold=100):
     """Process video file and return the path to the processed video"""
     if not load_mask_model():
-        return None, "Error: Could not load the mask detection model."
+        return None, "Error: Could not load the mask detection model. Please check the console for details."
     
     # Create a temporary file for the output
     with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as temp_file:
@@ -283,7 +305,7 @@ def process_video(video_path, progress=gr.Progress(), min_detection_confidence=0
 def process_webcam_frame(frame, min_detection_confidence, blur_threshold):
     """Process a single webcam frame"""
     if not load_mask_model():
-        return None
+        return frame  # Return original frame if model couldn't be loaded
     
     # Initialize face detector for each frame in webcam mode
     # This is less efficient but necessary for the Gradio webcam interface
@@ -365,5 +387,14 @@ with gr.Blocks(title="Enhanced Face Analysis System") as demo:
     - Higher blur threshold means more tolerance for blurry video
     """)
 
+# Ensure the model is downloaded when the app starts
+def initialize_app():
+    print("Initializing app and downloading model...")
+    if load_mask_model():
+        print("Model loaded successfully!")
+    else:
+        print("Failed to load model, some features may not work.")
+
 if __name__ == "__main__":
+    initialize_app()
     demo.launch()