FEAT: Finalize code for Hugging Face deployment

.gitignore

@@ -205,3 +205,4 @@ cython_debug/
 marimo/_static/
 marimo/_lsp/
 __marimo__/
+aws-key.pem

app.py

@@ -9,53 +9,49 @@ import tempfile
 import time
 from streamlit_option_menu import option_menu
 
-# --- Page Config (Set once at the top) ---
+# --- Page Config ---
 st.set_page_config(page_title="Facial Analysis", page_icon="👤", layout="wide", initial_sidebar_state="expanded")
 
-# --- Backend Loading (Robust and Unchanged) ---
+# --- Path Setup & Model Loading ---
 try:
+    # This works for local development
     src_path = os.path.abspath(os.path.join(os.path.dirname(__file__), 'src'))
     if src_path not in sys.path: sys.path.append(src_path)
     from cnnClassifier.pipeline.prediction import PredictionPipeline
 except ImportError:
-    st.error("FATAL: Prediction pipeline not found. Check project structure.")
-    st.stop()
+    # This is a fallback for Hugging Face Spaces
+    from src.cnnClassifier.pipeline.prediction import PredictionPipeline
+
+# --- TF Config (for MTCNN in Image/Video modes) ---
 try:
     gpus = tf.config.list_physical_devices('GPU')
     if gpus:
         for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True)
 except Exception: pass
+
 @st.cache_resource
 def load_pipeline():
     return PredictionPipeline()
+
 pipeline = load_pipeline()
 
-# --- Session State for Webcam Control ---
 if 'webcam_running' not in st.session_state: st.session_state.webcam_running = False
 def start_webcam(): st.session_state.webcam_running = True
 def stop_webcam(): st.session_state.webcam_running = False
 
-# --- Sidebar UI (Clean and Themed) ---
+# --- UI ---
 with st.sidebar:
     st.markdown("## ⚙️ Controls")
-    app_mode = option_menu(
-        menu_title=None,
-        options=["Image", "Video", "Live Feed"],
-        icons=["image", "film", "camera-video"], 
-        menu_icon="cast",
-        default_index=0,
-    )
-    st.divider()
-    st.info("This app uses a multi-task EfficientNet model to predict age and gender.")
-
-# --- Main Page Content ---
-st.title(f"👤 Facial Demographics Analysis")
-st.markdown(f"### Mode: {app_mode}")
-st.divider()
+    app_mode = option_menu(None, ["Image", "Video", "Live Feed"], 
+        icons=['image', 'film', 'camera-video'], menu_icon="cast", default_index=0)
 
 if not pipeline:
     st.error("AI Pipeline failed to load. Please check the terminal for errors.")
 else:
+    st.title("👤 Facial Demographics Analysis")
+    st.header(f"Mode: {app_mode}")
+    st.divider()
+
     if app_mode == "Image":
         uploaded_file = st.file_uploader("Upload an image for analysis", type=["jpg", "jpeg", "png"])
         if uploaded_file:
@@ -63,8 +59,9 @@ else:
             col1, col2 = st.columns(2)
             with col1: st.image(image, caption='Original Image', use_column_width=True)
             with col2:
-                with st.spinner('🔬 Analyzing...'):
-                    annotated_image, predictions = pipeline.predict_image(np.array(image))
+                with st.spinner('🔬 Analyzing with high-quality detector...'):
+                    # --- THE FIX: Call the HQ method ---
+                    annotated_image, predictions = pipeline.predict_hq(np.array(image))
                 st.image(annotated_image, caption='Processed Image', use_column_width=True)
                 if predictions:
                     with st.expander("View Details", expanded=True):
@@ -79,18 +76,17 @@ else:
             tfile.write(uploaded_file.read())
             cap = cv2.VideoCapture(tfile.name)
             frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-            st.info(f"Video has {frame_count} frames.")
-            if st.button("Start Video Processing", type="primary", use_container_width=True):
+            st.info(f"Video has {frame_count} frames. This will be slow but high-quality.")
+            if st.button("Process Video", type="primary", use_container_width=True):
                 progress_bar = st.progress(0, text="Initializing...")
                 out_tfile = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
                 h, w = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                 out = cv2.VideoWriter(out_tfile.name, cv2.VideoWriter_fourcc(*'mp4v'), cap.get(cv2.CAP_PROP_FPS), (w, h))
-                def frame_generator():
-                    for _ in range(frame_count):
-                        ret, frame = cap.read()
-                        if not ret: break
-                        yield cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-                for i, annotated_frame_rgb in enumerate(pipeline.process_video_stream(frame_generator())):
+                for i in range(frame_count):
+                    ret, frame = cap.read()
+                    if not ret: break
+                    # --- THE FIX: Call the HQ method ---
+                    annotated_frame_rgb, _ = pipeline.predict_hq(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                     out.write(cv2.cvtColor(annotated_frame_rgb, cv2.COLOR_RGB2BGR))
                     progress_bar.progress((i + 1) / frame_count, text=f"Processing Frame {i+1}/{frame_count}")
                 cap.release(), out.release()
@@ -100,15 +96,14 @@ else:
                     st.download_button("Download Processed Video", f, "output.mp4", "video/mp4", use_container_width=True)
 
     elif app_mode == "Live Feed":
+        st.info("Live feed uses a lightweight face detector for higher FPS.")
         col1, col2 = st.columns(2)
         with col1: st.button("Start Feed", on_click=start_webcam, use_container_width=True, type="primary")
         with col2: st.button("Stop Feed", on_click=stop_webcam, use_container_width=True)
-        
         _, center_col, _ = st.columns([1, 2, 1])
         with center_col:
             FRAME_WINDOW = st.image([])
             fps_display = st.empty()
-        
         if st.session_state.webcam_running:
             cap = cv2.VideoCapture(0)
             while st.session_state.webcam_running:
@@ -116,7 +111,8 @@ else:
                 ret, frame = cap.read()
                 if not ret: break
                 frame = cv2.flip(frame, 1)
-                annotated_frame = pipeline.process_live_frame(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+                # --- THE FIX: Call the LQ method ---
+                annotated_frame, _ = pipeline.predict_lq(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                 FRAME_WINDOW.image(annotated_frame, channels="RGB")
                 fps = 1.0 / (time.time() - start_time) if (time.time() - start_time) > 0 else 0
                 fps_display.markdown(f"<p style='text-align: center;'><b>FPS: {fps:.2f}</b></p>", unsafe_allow_html=True)

packages.txt

@@ -0,0 +1,2 @@
+libgl1-mesa-glx
+libglib2.0-0

requirements.txt

@@ -1,36 +1,17 @@
-# For PyTorch with CUDA 11.8 - MUST be installed with the extra index URL
---extra-index-url https://download.pytorch.org/whl/cu118
-torch==2.1.0+cu118
-torchvision==0.16.0+cu118
+torch==2.1.0
+torchvision==0.16.0
 torchaudio==2.1.0
-
-# Pin NumPy to a version compatible with Torch 2.1.0
-numpy>=1.23,<2.0
-
-# Hugging Face
+numpy<2.0
 transformers==4.36.2
 tokenizers==0.15.0
-datasets>=2.14.5
-evaluate
-accelerate>=0.25
-
-# MLOps and Utilities
-mlflow
-dvc[s3] # Assuming you might use S3 with DVC for AWS
+safetensors
 python-box
 PyYAML
-ensure
 pandas
 scikit-learn
 Pillow
-tqdm
-imblearn
-seaborn
-# Frontend and Real-time Processing
 streamlit
-opencv-python
-mtcnn
-tensorflow==2.15.0 
 streamlit-option-menu
-# AWS Deployment
-boto3
+opencv-python-headless
+mtcnn
+tensorflow==2.15.0

src/cnnClassifier/pipeline/prediction.py

@@ -1,110 +1,58 @@
 import torch
-import pandas as pd
 import numpy as np
 from PIL import Image
 from transformers import AutoImageProcessor
 import cv2
-from mtcnn import MTCNN
+from mtcnn import MTCNN  # For high-quality
 from pathlib import Path
 import sys
 import os
 from torchvision.transforms import Compose, Resize, ToTensor, Normalize
 from safetensors.torch import load_file as load_safetensors
-from collections import OrderedDict
-from scipy.spatial import distance as dist
 
 try:
     src_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
     if src_path not in sys.path: sys.path.append(src_path)
     from components.multi_task_model_trainer import MultiTaskEfficientNet
     from utils.common import read_yaml
-except ImportError as e:
-    print(f"Could not import custom modules: {e}.")
-    sys.exit(1)
-
-class CentroidTracker:
-    def __init__(self, max_disappeared=20):
-        self.next_object_id = 0
-        self.objects = OrderedDict()
-        self.disappeared = OrderedDict()
-        self.max_disappeared = max_disappeared
-
-    def register(self, centroid, box):
-        self.objects[self.next_object_id] = {'centroid': centroid, 'box': box, 'labels': {}, 'ema_preds': {}}
-        self.disappeared[self.next_object_id] = 0
-        self.next_object_id += 1
-
-    def deregister(self, object_id):
-        del self.objects[object_id]
-        del self.disappeared[object_id]
-
-    def update(self, boxes):
-        if len(boxes) == 0:
-            for object_id in list(self.disappeared.keys()):
-                self.disappeared[object_id] += 1
-                if self.disappeared[object_id] > self.max_disappeared:
-                    self.deregister(object_id)
-            return self.objects
-
-        input_centroids = np.array([(x + w // 2, y + h // 2) for (x, y, w, h) in boxes])
-
-        if len(self.objects) == 0:
-            for i in range(len(input_centroids)):
-                self.register(input_centroids[i], boxes[i])
-        else:
-            object_ids = list(self.objects.keys())
-            object_centroids = np.array([v['centroid'] for v in self.objects.values()])
-            D = dist.cdist(object_centroids, input_centroids)
-            rows = D.min(axis=1).argsort()
-            cols = D.argmin(axis=1)[rows]
-            used_rows, used_cols = set(), set()
-            for row, col in zip(rows, cols):
-                if row in used_rows or col in used_cols: continue
-                object_id = object_ids[row]
-                self.objects[object_id]['centroid'] = input_centroids[col]
-                self.objects[object_id]['box'] = boxes[col]
-                self.disappeared[object_id] = 0
-                used_rows.add(row)
-                used_cols.add(col)
-
-            unused_rows = set(range(D.shape[0])).difference(used_rows)
-            unused_cols = set(range(D.shape[1])).difference(used_cols)
-
-            if D.shape[0] >= D.shape[1]:
-                for row in unused_rows:
-                    object_id = object_ids[row]
-                    self.disappeared[object_id] += 1
-                    if self.disappeared[object_id] > self.max_disappeared:
-                        self.deregister(object_id)
-            else:
-                for col in unused_cols:
-                    self.register(input_centroids[col], boxes[col])
-        return self.objects
+except ImportError:
+    # Fallback for Hugging Face Spaces
+    from src.cnnClassifier.components.multi_task_model_trainer import MultiTaskEfficientNet
+    from src.cnnClassifier.utils.common import read_yaml
 
 class PredictionPipeline:
-    def __init__(self, model_path: str = "artifacts/multi_task_model_trainer/checkpoint-26873"):
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+    def __init__(self, model_path: str = "model/checkpoint-26873"):
+        self.device = "cpu"  # Force CPU for deployment
         self.model_path = Path(model_path)
         self.base_model_name = "google/efficientnet-b2"
-        params = read_yaml(Path("params.yaml"))
+        self.params = read_yaml(Path("model/params.yaml"))
+        
+        self.label_maps = {
+            'age_id2label': {'0': '0-2', '1': '3-9', '2': '10-19', '3': '20-29', '4': '30-39', '5': '40-49', '6': '50-59', '7': '60-69', '8': 'more than 70'},
+            'gender_id2label': {'0': 'Male', '1': 'Female'}
+        }
+        
+        print("--- Initializing Prediction Pipeline ---")
         self.processor = AutoImageProcessor.from_pretrained(self.base_model_name)
-        self.transforms = Compose([Resize((params.IMAGE_SIZE, params.IMAGE_SIZE)), ToTensor(), Normalize(mean=self.processor.image_mean, std=self.processor.image_std)])
-        self.label_maps = self._load_label_maps()
+        self.transforms = Compose([Resize((self.params.IMAGE_SIZE, self.params.IMAGE_SIZE)), ToTensor(), Normalize(mean=self.processor.image_mean, std=self.processor.image_std)])
         self.model = self._load_model()
-        self.face_detector = MTCNN()
-        self.tracker = CentroidTracker()
-        print(f"--- Pipeline Initialized on device: {self.device} ---")
-
-    def _load_label_maps(self):
-        maps = {'age_id2label': {'0': '0-2', '1': '3-9', '2': '10-19', '3': '20-29', '4': '30-39', '5': '40-49', '6': '50-59', '7': '60-69', '8': 'more than 70'},
-                'gender_id2label': {'0': 'Male', '1': 'Female'}}
-        return maps
+        
+        # --- THE FIX: LOAD BOTH DETECTORS ---
+        # High-quality detector for offline tasks
+        self.hq_face_detector = MTCNN() 
+        # Lightweight detector for live feed
+        haar_cascade_path = cv2.data.haarcascades + 'haarcascade_frontalface_default.xml'
+        self.lq_face_detector = cv2.CascadeClassifier(haar_cascade_path)
+        # --- END FIX ---
+        
+        print(f"--- Pipeline Initialized Successfully on device: {self.device} ---")
     
     def _load_model(self):
         num_age, num_gender, num_race = len(self.label_maps['age_id2label']), len(self.label_maps['gender_id2label']), 7
         model = MultiTaskEfficientNet(self.base_model_name, num_age, num_gender, num_race)
         weight_file = self.model_path / 'model.safetensors'
         if not weight_file.exists(): weight_file = self.model_path / 'pytorch_model.bin'
+        if not weight_file.exists(): raise FileNotFoundError(f"Weights not found in {self.model_path}")
         state_dict = load_safetensors(weight_file, device="cpu") if weight_file.suffix == ".safetensors" else torch.load(weight_file, map_location="cpu")
         model.load_state_dict(state_dict)
         model.to(self.device)
@@ -126,66 +74,48 @@ class PredictionPipeline:
         for i, line in enumerate(text_lines):
             y_text = y - total_height + (i * line_height) + 18
             cv2.putText(image, line, (x + 5, y_text), font, font_scale, text_color, font_thickness, cv2.LINE_AA)
-    
-    def _predict_for_box(self, frame, box):
-        x, y, w, h = [int(c) for c in box]
-        face_img = frame[max(0,y):min(frame.shape[0],y+h), max(0,x):min(frame.shape[1],x+w)]
-        if face_img.size == 0: return None
-        pixel_values = self.transforms(Image.fromarray(face_img)).unsqueeze(0).to(self.device)
-        with torch.no_grad(): outputs = self.model(pixel_values=pixel_values)
-        return outputs
-    
-    def predict_image(self, image_array):
+
+    def predict_hq(self, image_array: np.ndarray) -> (np.ndarray, list):
+        """High-quality prediction using MTCNN for images and videos."""
         annotated_image, predictions = image_array.copy(), []
-        face_results = self.face_detector.detect_faces(image_array)
+        face_results = self.hq_face_detector.detect_faces(image_array)
         if not face_results: return annotated_image, predictions
+
         for face in face_results:
-            if face['confidence'] < 0.9: continue
-            box = face['box']
-            outputs = self._predict_for_box(annotated_image, box)
-            if outputs:
-                age_label = self.label_maps['age_id2label'][str(outputs['age_logits'].argmax(1).item())]
-                gender_label = self.label_maps['gender_id2label'][str(outputs['gender_logits'].argmax(1).item())]
-                prediction_labels = {"age": age_label, "gender": gender_label}
-                predictions.append({**prediction_labels, 'box': box})
-                self._draw_predictions(annotated_image, box, prediction_labels)
+            if face['confidence'] < 0.95: continue
+            x, y, w, h = face['box']
+            face_img = image_array[max(0,y):min(image_array.shape[0],y+h), max(0,x):min(image_array.shape[1],x+w)]
+            if face_img.size == 0: continue
+            pil_face = Image.fromarray(face_img)
+            pixel_values = self.transforms(pil_face).unsqueeze(0).to(self.device)
+            with torch.no_grad(): outputs = self.model(pixel_values=pixel_values)
+            pred_id_age = str(outputs['age_logits'].argmax(1).item())
+            pred_id_gender = str(outputs['gender_logits'].argmax(1).item())
+            age_label = self.label_maps['age_id2label'].get(pred_id_age, "N/A")
+            gender_label = self.label_maps['gender_id2label'].get(pred_id_gender, "N/A")
+            prediction_labels = {"age": age_label, "gender": gender_label}
+            predictions.append({**prediction_labels, 'box': (x, y, w, h)})
+            self._draw_predictions(annotated_image, (x, y, w, h), prediction_labels)
         return annotated_image, predictions
 
-    def process_video_stream(self, frame_generator):
-        self.tracker = CentroidTracker()
-        for frame in frame_generator:
-            face_results = self.face_detector.detect_faces(frame)
-            boxes = [tuple(face['box']) for face in face_results if face['confidence'] > 0.9]
-            tracked_objects = self.tracker.update(boxes)
-            
-            for obj_id, data in tracked_objects.items():
-                # Predict only for new tracks or tracks that have just been re-found
-                if 'labels' not in data or self.tracker.disappeared[obj_id] == 0:
-                    outputs = self._predict_for_box(frame, data['box'])
-                    if outputs:
-                        alpha = 0.3
-                        current_probs = {
-                            'age': outputs['age_logits'].softmax(1).cpu().numpy()[0],
-                            'gender': outputs['gender_logits'].softmax(1).cpu().numpy()[0]
-                        }
-                        # Apply EMA smoothing
-                        if not data.get('ema_preds'): data['ema_preds'] = current_probs
-                        else:
-                            for task in ['age', 'gender']:
-                                data['ema_preds'][task] = alpha * current_probs[task] + (1 - alpha) * data['ema_preds'][task]
-                
-                # Always update the label from the latest smoothed probabilities
-                if data.get('ema_preds'):
-                    age_label = self.label_maps['age_id2label'][str(np.argmax(data['ema_preds']['age']))]
-                    gender_label = self.label_maps['gender_id2label'][str(np.argmax(data['ema_preds']['gender']))]
-                    data['labels'] = {"age": age_label, "gender": gender_label}
-            
-            annotated_frame = frame.copy()
-            for obj_id, data in tracked_objects.items():
-                if 'labels' in data:
-                    self._draw_predictions(annotated_frame, data['box'], data['labels'])
-            yield annotated_frame
-            
-    def process_live_frame(self, frame):
-        annotated_frame, _ = self.predict_image(frame)
-        return annotated_frame
+    def predict_lq(self, image_array: np.ndarray) -> (np.ndarray, list):
+        """Lightweight prediction using Haar Cascade for live feed."""
+        annotated_image, predictions = image_array.copy(), []
+        gray_image = cv2.cvtColor(image_array, cv2.COLOR_RGB2GRAY)
+        faces = self.lq_face_detector.detectMultiScale(gray_image, scaleFactor=1.1, minNeighbors=5, minSize=(60, 60))
+        if len(faces) == 0: return annotated_image, predictions
+
+        for (x, y, w, h) in faces:
+            face_img = image_array[y:y+h, x:x+w]
+            if face_img.size == 0: continue
+            pil_face = Image.fromarray(face_img)
+            pixel_values = self.transforms(pil_face).unsqueeze(0).to(self.device)
+            with torch.no_grad(): outputs = self.model(pixel_values=pixel_values)
+            pred_id_age = str(outputs['age_logits'].argmax(1).item())
+            pred_id_gender = str(outputs['gender_logits'].argmax(1).item())
+            age_label = self.label_maps['age_id2label'].get(pred_id_age, "N/A")
+            gender_label = self.label_maps['gender_id2label'].get(pred_id_gender, "N/A")
+            prediction_labels = {"age": age_label, "gender": gender_label}
+            predictions.append({**prediction_labels, 'box': (x, y, w, h)})
+            self._draw_predictions(annotated_image, (x, y, w, h), prediction_labels)
+        return annotated_image, predictions