Backend + Frontend done

.streamlit/config.toml

@@ -0,0 +1,15 @@
+[theme]
+# Primary accent color (a professional blue, similar to your example)
+primaryColor="#4A90E2" 
+
+# Main background color
+backgroundColor="#F0F4F8"
+
+# Sidebar and card background color
+secondaryBackgroundColor="#FFFFFF"
+
+# Default text color
+textColor="#31333F"
+
+# Font
+font="sans serif"

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "git.ignoreLimitWarning": true
+}

app.py

@@ -2,195 +2,125 @@ import streamlit as st
 import cv2
 import numpy as np
 from PIL import Image
-from transformers import pipeline
-from mtcnn import MTCNN
-from collections import defaultdict
+import tensorflow as tf
+import sys
+import os
+import tempfile
+import time
+from streamlit_option_menu import option_menu
 
-st.set_page_config(layout="wide", page_title="Facial Age Detection")
+# --- Page Config (Set once at the top) ---
+st.set_page_config(page_title="Facial Analysis", page_icon="👤", layout="wide", initial_sidebar_state="expanded")
 
-st.title("Facial Age Detection")
-st.write("Detect age groups from images, videos, or a live webcam feed.")
-st.write("This application uses an EfficientFormer-L1 model fine-tuned on the Facial Age dataset.")
-
-# --- Helper Functions and Classes ---
-
-@st.cache_resource
-def load_model():
-    """Load the age detection model pipeline."""
-    model_path = "artifacts/model_trainer/facial_age_detector_model"
-    pipe = pipeline('image-classification', model=model_path, device=0) # Use 0 for GPU
-    return pipe
-
-@st.cache_resource
-def load_face_detector():
-    """Load the MTCNN face detector."""
-    return MTCNN()
-
-def iou(boxA, boxB):
-    """Calculate Intersection over Union."""
-    xA = max(boxA[0], boxB[0])
-    yA = max(boxA[1], boxB[1])
-    xB = min(boxA[2], boxB[2])
-    yB = min(boxA[3], boxB[3])
-    interArea = max(0, xB - xA) * max(0, yB - yA)
-    boxAArea = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1])
-    boxBArea = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1])
-    return interArea / float(boxAArea + boxBArea - interArea)
-
-class EMATracker:
-    """Exponential Moving Average Tracker for smoothing predictions."""
-    def __init__(self, alpha=0.3):
-        self.alpha = alpha
-        self.tracked_objects = {} # {track_id: {box: [], ema_preds: {}}}
-
-    def update(self, detections, id_counter):
-        # Detections are a list of face boxes
-        # Simple tracking by IOU
-        
-        # Match detections to existing tracks
-        matches = {} # {track_id: det_idx}
-        used_det_indices = set()
-        
-        # This is a simple greedy matching. For more robust tracking, consider Hungarian algorithm.
-        for track_id, data in self.tracked_objects.items():
-            best_iou = 0
-            best_det_idx = -1
-            for i, det_box in enumerate(detections):
-                if i in used_det_indices: continue
-                current_iou = iou(data['box'], det_box)
-                if current_iou > best_iou and current_iou > 0.3: # IOU threshold
-                    best_iou = current_iou
-                    best_det_idx = i
-            if best_det_idx != -1:
-                matches[track_id] = best_det_idx
-                used_det_indices.add(best_det_idx)
-
-        # Update matched tracks
-        for track_id, det_idx in matches.items():
-            self.tracked_objects[track_id]['box'] = detections[det_idx]
-            
-        # Add new tracks
-        for i, det_box in enumerate(detections):
-            if i not in used_det_indices:
-                self.tracked_objects[id_counter] = {'box': det_box, 'ema_preds': defaultdict(float)}
-                id_counter += 1
-        
-        # Remove old tracks (optional, for long videos)
-        
-        return id_counter
-
-    def apply_ema(self, track_id, new_preds):
-        """Applies EMA to the predictions for a given track."""
-        if track_id not in self.tracked_objects:
-            return {}
-        
-        current_ema = self.tracked_objects[track_id]['ema_preds']
-        
-        # Initialize if new
-        if not current_ema:
-            for pred in new_preds:
-                current_ema[pred['label']] = pred['score']
-        else:
-            # Update existing values
-            for pred in new_preds:
-                label = pred['label']
-                current_ema[label] = (self.alpha * pred['score']) + ((1 - self.alpha) * current_ema[label])
-        
-        self.tracked_objects[track_id]['ema_preds'] = current_ema
-        
-        # Return the top prediction from EMA
-        if not current_ema: return None
-        top_label = max(current_ema, key=current_ema.get)
-        return f"{top_label} ({current_ema[top_label]:.2f})"
-
-
-# --- Load Models ---
+# --- Backend Loading (Robust and Unchanged) ---
 try:
-    age_pipe = load_model()
-    face_detector = load_face_detector()
-except Exception as e:
-    st.error(f"Error loading models: {e}. Please ensure the model is trained and located at 'artifacts/model_trainer/facial_age_detector_model'.")
+    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()
-
-
-# --- UI Sidebar ---
-st.sidebar.header("Input Options")
-app_mode = st.sidebar.selectbox("Choose the app mode", ["Image", "Video", "Live Webcam"])
-
-# --- Main App Logic ---
-
-if app_mode == "Image":
-    uploaded_file = st.sidebar.file_uploader("Upload an image...", type=["jpg", "jpeg", "png"])
-    if uploaded_file is not None:
-        image = Image.open(uploaded_file).convert("RGB")
-        img_array = np.array(image)
-        
-        st.image(image, caption='Uploaded Image.', use_column_width=True)
-        st.write("")
-        st.write("Detecting faces and predicting age...")
-        
-        faces = face_detector.detect_faces(img_array)
-        
-        if not faces:
-            st.warning("No faces detected in the image.")
-        else:
-            for face in faces:
-                x, y, w, h = face['box']
-                face_img = img_array[y:y+h, x:x+w]
-                pil_face = Image.fromarray(face_img)
-                
-                # Predict age
-                age_preds = age_pipe(pil_face)
-                top_pred = age_preds[0]
-                
-                # Draw on image
-                cv2.rectangle(img_array, (x, y), (x+w, y+h), (0, 255, 0), 2)
-                label = f"Age: {top_pred['label']} ({top_pred['score']:.2f})"
-                cv2.putText(img_array, label, (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0,255,0), 2)
-            
-            st.image(img_array, caption='Processed Image.', use_column_width=True)
-
-elif app_mode == "Live Webcam":
-    st.sidebar.info("Webcam feed will start automatically. Press 'Stop' to end.")
-    run = st.sidebar.button('Start Webcam')
-    stop = st.sidebar.button('Stop Webcam')
-    FRAME_WINDOW = st.image([])
-    
-    cap = cv2.VideoCapture(0)
-    tracker = EMATracker()
-    track_id_counter = 0
-
-    while run and not stop:
-        ret, frame = cap.read()
-        if not ret:
-            st.error("Failed to capture image from webcam.")
-            break
-        
-        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        faces = face_detector.detect_faces(frame_rgb)
+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) ---
+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()
+
+if not pipeline:
+    st.error("AI Pipeline failed to load. Please check the terminal for errors.")
+else:
+    if app_mode == "Image":
+        uploaded_file = st.file_uploader("Upload an image for analysis", type=["jpg", "jpeg", "png"])
+        if uploaded_file:
+            image = Image.open(uploaded_file).convert("RGB")
+            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))
+                st.image(annotated_image, caption='Processed Image', use_column_width=True)
+                if predictions:
+                    with st.expander("View Details", expanded=True):
+                        for i, p in enumerate(predictions):
+                            st.write(f"**Face {i+1}:** Gender: `{p['gender']}`, Age Group: `{p['age']}`")
+                else: st.warning("No faces detected.")
+
+    elif app_mode == "Video":
+        uploaded_file = st.file_uploader("Upload a video for analysis", type=["mp4", "mov", "avi"])
+        if uploaded_file:
+            tfile = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
+            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):
+                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())):
+                    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()
+                st.success("Video processing complete!")
+                st.video(out_tfile.name)
+                with open(out_tfile.name, "rb") as f:
+                    st.download_button("Download Processed Video", f, "output.mp4", "video/mp4", use_container_width=True)
+
+    elif app_mode == "Live Feed":
+        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)
         
-        detection_boxes = [f['box'] for f in faces]
-        track_id_counter = tracker.update(detection_boxes, track_id_counter)
+        _, center_col, _ = st.columns([1, 2, 1])
+        with center_col:
+            FRAME_WINDOW = st.image([])
+            fps_display = st.empty()
         
-        for track_id, data in tracker.tracked_objects.items():
-            x, y, w, h = data['box']
-            if w > 20 and h > 20: # Filter small detections
-                face_img = frame_rgb[y:y+h, x:x+w]
-                pil_face = Image.fromarray(face_img)
-                
-                age_preds = age_pipe(pil_face)
-                smoothed_label = tracker.apply_ema(track_id, age_preds)
-                
-                cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 0), 2)
-                if smoothed_label:
-                    cv2.putText(frame, smoothed_label, (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0,255,0), 2)
-
-        FRAME_WINDOW.image(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-    
-    cap.release()
-    st.sidebar.info("Webcam stopped.")
-
-# Add a placeholder for Video processing, which would be similar to Webcam but with a file uploader.
-elif app_mode == "Video":
-    st.sidebar.warning("Video processing is similar to the webcam feed but processes a file. This feature is not fully implemented in this demo but follows the same logic.")
-    # You would use cv2.VideoCapture(video_path) and loop through frames.
+        if st.session_state.webcam_running:
+            cap = cv2.VideoCapture(0)
+            while st.session_state.webcam_running:
+                start_time = time.time()
+                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))
+                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)
+            cap.release()
+            cv2.destroyAllWindows()
+            st.session_state.webcam_running = False
+            st.rerun()

config/config.yaml

@@ -2,17 +2,22 @@ artifacts_root: artifacts
 
 data_ingestion:
   root_dir: artifacts/data_ingestion
-  dataset_name: frabbisw/facial-age
-  local_data_file: artifacts/data_ingestion/data.zip
-  unzip_dir: artifacts/data_ingestion
+  dataset_name: "HuggingFaceM4/FairFace"
+  dataset_config: "0.25"
+  local_data_dir: artifacts/data_ingestion/dataset
 
 data_preparation:
   root_dir: artifacts/data_preparation
-  data_path: artifacts/data_ingestion/face_age
-  dataset_name: facial_age_prepared_dataset
+  # Path to the raw dataset downloaded by the previous stage
+  raw_data_path: artifacts/data_ingestion/dataset
+  # Path where we will save the cleaned dataframe
+  cleaned_data_path: artifacts/data_preparation/fairface_cleaned.csv
 
-model_trainer:
-  root_dir: artifacts/model_trainer
-  trained_model_path: artifacts/model_trainer/facial_age_detector_model
-  # Using EfficientFormer-L1, a much lighter model than ViT
-  model_name: "snap-research/efficientformer-l1-300"
+multi_task_model_trainer:
+  root_dir: artifacts/multi_task_model_trainer
+  # The path to our cleaned CSV file from the previous stage
+  data_path: artifacts/data_preparation/fairface_cleaned.csv
+  # Where to save the final multi-task model
+  trained_model_path: artifacts/multi_task_model_trainer/facial_demographics_model
+  # The base model from Hugging Face
+  model_name: "google/efficientnet-b2"

dvc.lock

@@ -3,24 +3,24 @@ stages:
   data_ingestion:
     cmd: python src/cnnClassifier/pipeline/stage_01_data_ingestion.py
     deps:
-    - path: config/config.yaml
-      hash: md5
-      md5: 3cea2dfb36f0a5e40dd599dad9458ca4
-      size: 609
     - path: src/cnnClassifier/components/data_ingestion.py
       hash: md5
-      md5: 80b591ef3eedaf256ef85f4d196a0d43
-      size: 1591
+      md5: 11a88b6a3a45651504f59cab654dd889
+      size: 1737
+    - path: src/cnnClassifier/config/configuration.py
+      hash: md5
+      md5: f2ee2d8b6bf946faa613a7d498bf789a
+      size: 2458
     - path: src/cnnClassifier/pipeline/stage_01_data_ingestion.py
       hash: md5
-      md5: 2e1c2ad52ddc9763ff2a241576a7477c
-      size: 904
+      md5: ba77afc62b94ad61c5990601bc2a6f15
+      size: 945
     outs:
     - path: artifacts/data_ingestion
       hash: md5
-      md5: 35941f86a72fc72e64cb3195753ae21d.dir
-      size: 1758455894
-      nfiles: 19557
+      md5: 010e49306f9a2f8bd5f6235a9bb4c40a.dir
+      size: 1158773512
+      nfiles: 12
   model_training:
     cmd: python src/cnnClassifier/pipeline/stage_02_model_training.py
     deps:
@@ -51,3 +51,59 @@ stages:
       md5: 621f61ba7beea89c3bef7a921afdcc9d.dir
       size: 183039001
       nfiles: 12
+  data_preparation:
+    cmd: python src/cnnClassifier/pipeline/stage_02_data_preparation.py
+    deps:
+    - path: artifacts/data_ingestion/dataset
+      hash: md5
+      md5: 86434e33cd2b0a60b09c0624d29f1fda.dir
+      size: 579423180
+      nfiles: 8
+    - path: config/config.yaml
+      hash: md5
+      md5: bc8e095ab04797e455847fc34f3db546
+      size: 908
+    - path: src/cnnClassifier/components/data_preparation.py
+      hash: md5
+      md5: 39d7a55e908ab9b099f55c21a3019014
+      size: 2175
+    - path: src/cnnClassifier/pipeline/stage_02_data_preparation.py
+      hash: md5
+      md5: d04a9dd31ed636a27a79235a0dff46a6
+      size: 953
+    outs:
+    - path: artifacts/data_preparation
+      hash: md5
+      md5: 412fbf5339e9b82998c7617d19fce476.dir
+      size: 580361278
+      nfiles: 97699
+  multi_task_model_training:
+    cmd: python src/cnnClassifier/pipeline/stage_03_multi_task_model_training.py
+    deps:
+    - path: artifacts/data_preparation
+      hash: md5
+      md5: 412fbf5339e9b82998c7617d19fce476.dir
+      size: 580361278
+      nfiles: 97699
+    - path: config/config.yaml
+      hash: md5
+      md5: bc8e095ab04797e455847fc34f3db546
+      size: 908
+    - path: params.yaml
+      hash: md5
+      md5: ce8c137aa11f22d0901fb41485e9bfde
+      size: 239
+    - path: src/cnnClassifier/components/multi_task_model_trainer.py
+      hash: md5
+      md5: 5429e22ede43731b1806a9218c41c6d7
+      size: 7510
+    - path: src/cnnClassifier/pipeline/stage_03_multi_task_model_training.py
+      hash: md5
+      md5: 317cd31673bc55d1d73dea54cb75a8e0
+      size: 974
+    outs:
+    - path: artifacts/multi_task_model_trainer
+      hash: md5
+      md5: b641e844039bc54d5c29145dfb0bab77.dir
+      size: 1898189766
+      nfiles: 123

dvc.yaml

@@ -4,17 +4,28 @@ stages:
     deps:
       - src/cnnClassifier/pipeline/stage_01_data_ingestion.py
       - src/cnnClassifier/components/data_ingestion.py
+      - src/cnnClassifier/config/configuration.py
+    outs:
+      - artifacts/data_ingestion # This output now includes the saved dataset
+
+  data_preparation: # <<< NEW STAGE
+    cmd: python src/cnnClassifier/pipeline/stage_02_data_preparation.py
+    deps:
+      - src/cnnClassifier/pipeline/stage_02_data_preparation.py
+      - src/cnnClassifier/components/data_preparation.py
       - config/config.yaml
+      - artifacts/data_ingestion/dataset # Depends on the output of the first stage
     outs:
-      - artifacts/data_ingestion
+      - artifacts/data_preparation
 
-  model_training:
-    cmd: python src/cnnClassifier/pipeline/stage_02_model_training.py
+  multi_task_model_training: # <<< NEW STAGE
+    cmd: python src/cnnClassifier/pipeline/stage_03_multi_task_model_training.py
     deps:
-      - src/cnnClassifier/pipeline/stage_02_model_training.py
-      - src/cnnClassifier/components/model_trainer.py
+      - src/cnnClassifier/pipeline/stage_03_multi_task_model_training.py
+      - src/cnnClassifier/components/multi_task_model_trainer.py
       - config/config.yaml
       - params.yaml
-      - artifacts/data_ingestion  # Depends on the output of the previous stage
+      - artifacts/data_preparation
     outs:
-      - artifacts/model_trainer
+      - artifacts/multi_task_model_trainer
+      

requirements.txt

@@ -25,12 +25,12 @@ 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

research/evaluation.py

@@ -0,0 +1,163 @@
+# research/evaluation.py
+
+import torch
+import pandas as pd
+import json
+from pathlib import Path
+import sys
+import os
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import accuracy_score, confusion_matrix, classification_report
+import seaborn as sns
+import matplotlib.pyplot as plt
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+# Add the project's src directory to the Python path
+# This allows us to import our custom modules
+src_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'src'))
+sys.path.append(src_path)
+
+# Now we can import our custom classes
+from cnnClassifier.components.multi_task_model_trainer import MultiTaskEfficientNet, FairFaceDataset
+from cnnClassifier.utils.common import read_yaml
+from torchvision.transforms import Compose, Resize, ToTensor, Normalize
+from transformers import AutoImageProcessor
+
+# ==============================================================================
+# CONFIGURATION
+# ==============================================================================
+# Define paths directly. We are not using the config manager.
+MODEL_PATH = Path("artifacts/multi_task_model_trainer/facial_demographics_model")
+DATA_PATH = Path("artifacts/data_preparation/fairface_cleaned.csv")
+PARAMS_PATH = Path("params.yaml")
+EVALUATION_OUTPUT_DIR = Path("artifacts/manual_evaluation")
+EVALUATION_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+
+# Load parameters
+params = read_yaml(PARAMS_PATH)
+IMAGE_SIZE = params.IMAGE_SIZE
+BATCH_SIZE = params.BATCH_SIZE
+TEST_SPLIT_SIZE = params.TEST_SPLIT_SIZE
+RANDOM_STATE = params.RANDOM_STATE
+
+# ==============================================================================
+# MAIN EVALUATION LOGIC
+# ==============================================================================
+def evaluate_model():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"--- Running evaluation on device: {device} ---")
+
+    # 1. Load data and prepare the test split
+    print("Loading and preparing test data...")
+    df = pd.read_csv(DATA_PATH)
+    
+    label_maps = {}
+    for task in ['age', 'gender', 'race']:
+        labels = sorted(df[task].unique())
+        label_maps[f'{task}_label2id'] = {label: i for i, label in enumerate(labels)}
+        label_maps[f'{task}_id2label'] = {i: label for i, label in enumerate(labels)}
+        df[f'{task}_id'] = df[task].map(label_maps[f'{task}_label2id'])
+    
+    # Use the same random_state to ensure we get the identical test split as in training
+    _, test_df = train_test_split(
+        df, 
+        test_size=TEST_SPLIT_SIZE, 
+        random_state=RANDOM_STATE, 
+        stratify=df['age']
+    )
+    
+    # 2. Create the PyTorch DataLoader
+    model_config = read_yaml(Path("config/config.yaml"))
+    base_model_name = model_config.multi_task_model_trainer.model_name
+    processor = AutoImageProcessor.from_pretrained(base_model_name)
+    _transforms = Compose([
+        Resize((IMAGE_SIZE, IMAGE_SIZE)),
+        ToTensor(),
+        Normalize(mean=processor.image_mean, std=processor.image_std)
+    ])
+    
+    test_dataset = FairFaceDataset(dataframe=test_df, processor=processor, transforms=_transforms)
+    test_dataloader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)
+
+    # 3. Load the trained model
+    print(f"Loading model from: {MODEL_PATH}")
+    model = MultiTaskEfficientNet(
+        model_name=str(MODEL_PATH), # Pass the path as the model name
+        num_labels_age=len(label_maps['age_id2label']),
+        num_labels_gender=len(label_maps['gender_id2label']),
+        num_labels_race=len(label_maps['race_id2label']),
+    ).to(device)
+    
+    # Load the trained weights
+    model.load_state_dict(torch.load(MODEL_PATH / 'pytorch_model.bin', map_location=device))
+    model.eval()
+
+    # 4. Run predictions on the test set
+    print("Running predictions on the test set...")
+    all_preds = {'age': [], 'gender': [], 'race': []}
+    all_labels = {'age': [], 'gender': [], 'race': []}
+
+    for batch in tqdm(test_dataloader, desc="Evaluating"):
+        pixel_values = batch['pixel_values'].to(device)
+        labels = batch['labels']
+        
+        with torch.no_grad():
+            outputs = model(pixel_values=pixel_values)
+        
+        all_preds['age'].extend(outputs['age_logits'].argmax(1).cpu().numpy())
+        all_preds['gender'].extend(outputs['gender_logits'].argmax(1).cpu().numpy())
+        all_preds['race'].extend(outputs['race_logits'].argmax(1).cpu().numpy())
+
+        all_labels['age'].extend(labels['age'].cpu().numpy())
+        all_labels['gender'].extend(labels['gender'].cpu().numpy())
+        all_labels['race'].extend(labels['race'].cpu().numpy())
+
+    # 5. Calculate metrics, generate reports, and save artifacts
+    print("--- Evaluation Results ---")
+    metrics = {}
+    for task in ['age', 'gender', 'race']:
+        accuracy = accuracy_score(all_labels[task], all_preds[task])
+        print(f"\n--- {task.capitalize()} ---")
+        print(f"Accuracy: {accuracy:.4f}")
+        
+        report_str = classification_report(
+            all_labels[task], 
+            all_preds[task], 
+            target_names=list(label_maps[f'{task}_id2label'].values())
+        )
+        print("Classification Report:")
+        print(report_str)
+        
+        metrics[f'{task}_accuracy'] = accuracy
+        
+        # Confusion Matrix
+        cm = confusion_matrix(all_labels[task], all_preds[task])
+        plt.figure(figsize=(12, 10))
+        sns.heatmap(cm, annot=True, fmt='d', xticklabels=list(label_maps[f'{task}_id2label'].values()), yticklabels=list(label_maps[f'{task}_id2label'].values()), cmap='Blues')
+        plt.title(f'Confusion Matrix for {task.capitalize()}', fontsize=16)
+        plt.ylabel('Actual', fontsize=12)
+        plt.xlabel('Predicted', fontsize=12)
+        plt.xticks(rotation=45)
+        plt.yticks(rotation=0)
+        
+        cm_path = EVALUATION_OUTPUT_DIR / f'{task}_confusion_matrix.png'
+        plt.savefig(cm_path, bbox_inches='tight')
+        plt.close() # Close the plot to avoid displaying it in the console
+        print(f"Saved {task} confusion matrix to {cm_path}")
+
+    # Save metrics to a JSON file
+    metrics_path = EVALUATION_OUTPUT_DIR / "metrics.json"
+    with open(metrics_path, 'w') as f:
+        json.dump(metrics, f, indent=4)
+    print(f"\nSaved final metrics to {metrics_path}")
+    
+    # Save the label maps used for this evaluation run
+    label_maps_path = EVALUATION_OUTPUT_DIR / "label_maps.json"
+    with open(label_maps_path, 'w') as f:
+        json.dump(label_maps, f, indent=4)
+    print(f"Saved label maps to {label_maps_path}")
+
+
+if __name__ == '__main__':
+    evaluate_model()

src/cnnClassifier/components/data_ingestion.py

@@ -1,34 +1,40 @@
-import os
-import zipfile
+# src/cnnClassifier/components/data_ingestion.py
+
+from datasets import load_dataset
 from cnnClassifier import logger
 from cnnClassifier.entity.config_entity import DataIngestionConfig
+from pathlib import Path
 
 class DataIngestion:
     def __init__(self, config: DataIngestionConfig):
         self.config = config
 
-    def download_file(self):
+    def download_dataset(self):
         """
-        Downloads the dataset from Kaggle.
-        Make sure to have your kaggle.json file in ~/.kaggle/ or set KAGGLE_USERNAME and KAGGLE_KEY env variables.
+        Downloads and saves the FairFace dataset from the Hugging Face Hub.
         """
         try:
-            logger.info(f"Downloading dataset from kaggle: {self.config.dataset_name}")
-            os.system(f"kaggle datasets download {self.config.dataset_name} -p {os.path.dirname(self.config.local_data_file)}")
-            # The downloaded file will be named 'facial-age.zip'. We need to rename it to 'data.zip' as per our config.
-            downloaded_zip_path = os.path.join(os.path.dirname(self.config.local_data_file), 'facial-age.zip')
-            os.rename(downloaded_zip_path, self.config.local_data_file)
-            logger.info(f"Dataset downloaded and saved at {self.config.local_data_file}")
-        except Exception as e:
-            logger.error(f"Failed to download dataset. Error: {e}")
-            raise e
+            logger.info(f"Downloading dataset '{self.config.dataset_name}' from Hugging Face Hub...")
+            
+            # load_dataset handles everything: download, verification, and caching
+            # It returns a DatasetDict, typically with 'train' and 'validation' splits
+            fairface_dataset = load_dataset(
+                self.config.dataset_name,
+                name=self.config.dataset_config,
+                cache_dir=self.config.root_dir # Use our root_dir for caching
+            )
+            
+            # Save the downloaded dataset to our specified artifacts directory
+            # This makes it a persistent part of our DVC pipeline
+            save_path = Path(self.config.local_data_dir)
+            fairface_dataset.save_to_disk(save_path)
+            
+            logger.info(f"Dataset successfully downloaded and saved to {save_path}")
 
-    def extract_zip_file(self):
-        """
-        Extracts the zip file into the data directory
-        """
-        unzip_path = self.config.unzip_dir
-        os.makedirs(unzip_path, exist_ok=True)
-        with zipfile.ZipFile(self.config.local_data_file, 'r') as zip_ref:
-            zip_ref.extractall(unzip_path)
-        logger.info(f"Dataset extracted to {unzip_path}")
+            # Optional: Log the structure of the downloaded dataset
+            logger.info(f"Dataset splits: {list(fairface_dataset.keys())}")
+            logger.info(f"Training set features: {fairface_dataset['train'].features}")
+
+        except Exception as e:
+            logger.error(f"Failed to download or save dataset. Error: {e}")
+            raise e

src/cnnClassifier/components/data_preparation.py

@@ -0,0 +1,50 @@
+from datasets import load_from_disk
+import pandas as pd
+from cnnClassifier import logger
+from cnnClassifier.entity.config_entity import DataPreparationConfig
+from pathlib import Path
+from PIL import Image # <<< ADD THIS IMPORT
+import io              # <<< ADD THIS IMPORT
+
+class DataPreparation:
+    def __init__(self, config: DataPreparationConfig):
+        self.config = config
+
+    def create_cleaned_dataframe(self):
+        try:
+            logger.info("Loading raw dataset to create cleaned CSV...")
+            raw_dataset = load_from_disk(self.config.raw_data_path)
+            
+            df_train = raw_dataset['train'].to_pandas()
+            df_val = raw_dataset['validation'].to_pandas()
+            combined_df = pd.concat([df_train, df_val], ignore_index=True)
+            
+            image_dir = Path("artifacts/data_preparation/images")
+            image_dir.mkdir(parents=True, exist_ok=True)
+            
+            combined_df['image_file_path'] = [
+                str(image_dir / f"{i}.jpg") for i in range(len(combined_df))
+            ]
+            
+            # --- IMPORTANT ---
+            # We only need the file path for the CSV, so we drop the bulky 'image' column
+            final_df_for_csv = combined_df.drop(columns=['image'])
+            
+            logger.info(f"Saving cleaned metadata to {self.config.cleaned_data_path}")
+            final_df_for_csv.to_csv(self.config.cleaned_data_path, index=False)
+
+            # --- CORRECTED IMAGE SAVING LOOP ---
+            logger.info(f"Deterministically saving images to {image_dir}...")
+            for i, row in combined_df.iterrows():
+                image_path = Path(row['image_file_path'])
+                image_dict = row['image']
+                
+                # Recreate the PIL Image from the dictionary's bytes data
+                pil_image = Image.open(io.BytesIO(image_dict['bytes']))
+                
+                # Now save the reconstructed PIL Image
+                pil_image.save(image_path)
+
+        except Exception as e:
+            logger.error(f"Failed during data preparation. Error: {e}")
+            raise e

src/cnnClassifier/components/multi_task_model_trainer.py

@@ -0,0 +1,161 @@
+import torch
+import torch.nn as nn
+import pandas as pd
+import numpy as np
+from torch.utils.data import Dataset
+from transformers import (
+    AutoImageProcessor,
+    AutoModelForImageClassification,
+    TrainingArguments, 
+    Trainer
+)
+from torchvision.transforms import (
+    Compose, 
+    Normalize, 
+    RandomRotation, 
+    RandomHorizontalFlip,
+    Resize,
+    ToTensor
+)
+from cnnClassifier.entity.config_entity import MultiTaskModelTrainerConfig
+from cnnClassifier import logger
+from PIL import Image
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import accuracy_score
+
+class MultiTaskEfficientNet(nn.Module):
+    def __init__(self, model_name, num_labels_age, num_labels_gender, num_labels_race):
+        super().__init__()
+        self.efficientnet_base = AutoModelForImageClassification.from_pretrained(model_name, ignore_mismatched_sizes=True)
+        original_classifier = self.efficientnet_base.classifier
+        feature_dim = original_classifier.in_features
+        self.efficientnet_base.classifier = nn.Identity()
+
+        self.age_classifier = nn.Linear(feature_dim, num_labels_age)
+        self.gender_classifier = nn.Linear(feature_dim, num_labels_gender)
+        self.race_classifier = nn.Linear(feature_dim, num_labels_race)
+
+    def forward(self, pixel_values, labels=None):
+        features = self.efficientnet_base.efficientnet(pixel_values)
+        pooled_features = features.last_hidden_state.mean(dim=[2, 3])
+        age_logits = self.age_classifier(pooled_features)
+        gender_logits = self.gender_classifier(pooled_features)
+        race_logits = self.race_classifier(pooled_features)
+        
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            age_loss = loss_fct(age_logits, labels[:, 0])
+            gender_loss = loss_fct(gender_logits, labels[:, 1])
+            race_loss = loss_fct(race_logits, labels[:, 2])
+            loss = (2.0 * age_loss) + gender_loss + race_loss
+
+        return {"loss": loss, "age_logits": age_logits, "gender_logits": gender_logits, "race_logits": race_logits}
+
+class FairFaceDataset(Dataset):
+    def __init__(self, dataframe, processor, transforms):
+        self.dataframe = dataframe
+        self.processor = processor
+        self.transforms = transforms
+        self.normalize = Normalize(mean=processor.image_mean, std=processor.image_std)
+
+    def __len__(self):
+        return len(self.dataframe)
+
+    def __getitem__(self, idx):
+        row = self.dataframe.iloc[idx]
+        image_path = row['image_file_path']
+        image = Image.open(image_path).convert("RGB")
+        pixel_values = self.transforms(image)
+        pixel_values = self.normalize(pixel_values)
+        
+        labels = torch.tensor([row['age_id'], row['gender_id'], row['race_id']], dtype=torch.long)
+        return {"pixel_values": pixel_values, "labels": labels}
+
+def compute_multitask_metrics(eval_pred):
+    predictions, labels = eval_pred
+    age_logits, gender_logits, race_logits = predictions['age_logits'], predictions['gender_logits'], predictions['race_logits']
+    age_preds = np.argmax(age_logits, axis=1)
+    gender_preds = np.argmax(gender_logits, axis=1)
+    race_preds = np.argmax(race_logits, axis=1)
+    age_labels, gender_labels, race_labels = labels[:, 0], labels[:, 1], labels[:, 2]
+    age_acc = accuracy_score(age_labels, age_preds)
+    gender_acc = accuracy_score(gender_labels, gender_preds)
+    race_acc = accuracy_score(race_labels, race_preds)
+    overall_acc = (age_acc + gender_acc + race_acc) / 3.0
+    return {"age_accuracy": age_acc, "gender_accuracy": gender_acc, "race_accuracy": race_acc, "overall_accuracy": overall_acc}
+
+class MultiTaskModelTrainer:
+    def __init__(self, config: MultiTaskModelTrainerConfig):
+        self.config = config
+        self.processor = AutoImageProcessor.from_pretrained(config.model_name)
+
+    def train(self):
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        logger.info(f"Using device: {device}")
+
+        logger.info("Loading and preparing dataset from cleaned CSV...")
+        df = pd.read_csv(self.config.data_path)
+        
+        label_maps = {}
+        for task in ['age', 'gender', 'race']:
+            labels = sorted(df[task].unique())
+            label_maps[f'{task}_label2id'] = {label: i for i, label in enumerate(labels)}
+            df[f'{task}_id'] = df[task].map(label_maps[f'{task}_label2id'])
+        
+        num_classes_age = len(label_maps['age_label2id'])
+        num_classes_gender = len(label_maps['gender_label2id'])
+        num_classes_race = len(label_maps['race_label2id'])
+        train_df, test_df = train_test_split(df, test_size=self.config.test_split_size, random_state=self.config.random_state, stratify=df['age'])
+        
+        train_transforms = Compose([
+            Resize((self.config.image_size, self.config.image_size)),
+            RandomHorizontalFlip(),
+            RandomRotation(10),
+            ToTensor(), # Normalization is now in the Dataset
+        ])
+        
+        val_transforms = Compose([
+            Resize((self.config.image_size, self.config.image_size)),
+            ToTensor(),
+        ])
+        
+        train_dataset = FairFaceDataset(dataframe=train_df, processor=self.processor, transforms=train_transforms)
+        test_dataset = FairFaceDataset(dataframe=test_df, processor=self.processor, transforms=val_transforms)
+
+        model = MultiTaskEfficientNet(model_name=self.config.model_name, num_labels_age=num_classes_age, num_labels_gender=num_classes_gender, num_labels_race=num_classes_race).to(device)
+        
+        args = TrainingArguments(
+            output_dir=self.config.root_dir,
+            logging_dir=f'{self.config.root_dir}/logs',
+            evaluation_strategy="epoch",
+            learning_rate=self.config.learning_rate,
+            per_device_train_batch_size=self.config.batch_size,
+            per_device_eval_batch_size=self.config.batch_size,
+            num_train_epochs=self.config.num_train_epochs,
+            weight_decay=self.config.weight_decay,
+            save_strategy='epoch',
+            load_best_model_at_end=True,
+            metric_for_best_model="eval_overall_accuracy",
+            dataloader_num_workers=4,
+            lr_scheduler_type='cosine',
+            report_to="none"
+        )
+        
+        class EvalTrainer(Trainer):
+            def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None):
+                has_labels = "labels" in inputs
+                inputs = self._prepare_inputs(inputs)
+                with torch.no_grad():
+                    outputs = model(**inputs)
+                    loss = outputs.get("loss")
+                predictions = {"age_logits": outputs["age_logits"], "gender_logits": outputs["gender_logits"], "race_logits": outputs["race_logits"]}
+                return (loss, predictions, inputs["labels"] if has_labels else None)
+
+        trainer = EvalTrainer(model=model, args=args, train_dataset=train_dataset, eval_dataset=test_dataset, compute_metrics=compute_multitask_metrics)
+
+        trainer.train()
+        
+        logger.info(f"Saving final model and processor to {self.config.trained_model_path}")
+        trainer.save_model(self.config.trained_model_path)
+        self.processor.save_pretrained(self.config.trained_model_path)

src/cnnClassifier/config/configuration.py

@@ -1,6 +1,10 @@
 from cnnClassifier.constants import *
 from cnnClassifier.utils.common import read_yaml, create_directories
-from cnnClassifier.entity.config_entity import DataIngestionConfig, DataPreparationConfig, ModelTrainerConfig
+from cnnClassifier.entity.config_entity import (
+    DataIngestionConfig,
+    DataPreparationConfig,
+    MultiTaskModelTrainerConfig # <-- Import the new one
+)
 
 class ConfigurationManager:
     def __init__(
@@ -15,46 +19,45 @@ class ConfigurationManager:
 
     def get_data_ingestion_config(self) -> DataIngestionConfig:
         config = self.config.data_ingestion
-
         create_directories([config.root_dir])
 
         data_ingestion_config = DataIngestionConfig(
             root_dir=config.root_dir,
             dataset_name=config.dataset_name,
-            local_data_file=config.local_data_file,
-            unzip_dir=config.unzip_dir
+            dataset_config=config.dataset_config, 
+            local_data_dir=config.local_data_dir
         )
         return data_ingestion_config
-
-    def get_data_preparation_config(self) -> DataPreparationConfig:
+    
+    def get_data_preparation_config(self) -> DataPreparationConfig: # <<< NEW METHOD
         config = self.config.data_preparation
+
         create_directories([config.root_dir])
-        
+
         data_preparation_config = DataPreparationConfig(
             root_dir=config.root_dir,
-            data_path=config.data_path,
-            dataset_name=config.dataset_name
+            raw_data_path=config.raw_data_path,
+            cleaned_data_path=config.cleaned_data_path
         )
         return data_preparation_config
-    
-    def get_model_trainer_config(self) -> ModelTrainerConfig:
-        config = self.config.model_trainer
-        data_prep_config = self.config.data_preparation
+
+    def get_multi_task_model_trainer_config(self) -> MultiTaskModelTrainerConfig:
+        config = self.config.multi_task_model_trainer
         params = self.params
         create_directories([config.root_dir])
 
-        model_trainer_config = ModelTrainerConfig(
-        root_dir=Path(config.root_dir),
-        data_path=Path(data_prep_config.data_path),
-        trained_model_path=Path(config.trained_model_path),
-        model_name=config.model_name,
-        image_size=int(params.IMAGE_SIZE),
-        learning_rate=float(params.LEARNING_RATE), # <<< CORRECTED
-        batch_size=int(params.BATCH_SIZE),
-        num_train_epochs=int(params.NUM_TRAIN_EPOCHS),
-        weight_decay=float(params.WEIGHT_DECAY), # <<< CORRECTED
-        warmup_steps=int(params.WARMUP_STEPS),
-        test_split_size=float(params.TEST_SPLIT_SIZE), # <<< CORRECTED
-        random_state=int(params.RANDOM_STATE)
+        multi_task_model_trainer_config = MultiTaskModelTrainerConfig(
+            root_dir=Path(config.root_dir),
+            data_path=config.data_path,
+            trained_model_path=Path(config.trained_model_path),
+            model_name=config.model_name,
+            image_size=int(params.IMAGE_SIZE),
+            learning_rate=float(params.LEARNING_RATE),
+            batch_size=int(params.BATCH_SIZE),
+            num_train_epochs=int(params.NUM_TRAIN_EPOCHS),
+            weight_decay=float(params.WEIGHT_DECAY),
+            warmup_steps=int(params.WARMUP_STEPS),
+            test_split_size=float(params.TEST_SPLIT_SIZE),
+            random_state=int(params.RANDOM_STATE)
         )
-        return model_trainer_config
+        return multi_task_model_trainer_config

src/cnnClassifier/entity/config_entity.py

@@ -5,19 +5,21 @@ from pathlib import Path
 class DataIngestionConfig:
     root_dir: Path
     dataset_name: str
-    local_data_file: Path
-    unzip_dir: Path
+    local_data_dir: Path
+    dataset_config: str
 
 @dataclass(frozen=True)
-class DataPreparationConfig:
+class DataPreparationConfig: # <<< NEW DATACLASS
     root_dir: Path
-    data_path: Path
-    dataset_name: str
+    raw_data_path: Path
+    cleaned_data_path: Path
+
+# ... (other configs are above)
 
 @dataclass(frozen=True)
-class ModelTrainerConfig:
+class MultiTaskModelTrainerConfig:
     root_dir: Path
-    data_path: Path
+    data_path: str
     trained_model_path: Path
     model_name: str
     image_size: int

src/cnnClassifier/pipeline/prediction.py

@@ -0,0 +1,191 @@
+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 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
+
+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"
+        self.model_path = Path(model_path)
+        self.base_model_name = "google/efficientnet-b2"
+        params = read_yaml(Path("params.yaml"))
+        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.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
+    
+    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'
+        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)
+        model.eval()
+        return model
+
+    def _draw_predictions(self, image, box, labels):
+        x, y, w, h = [int(c) for c in box]
+        font, font_scale, font_thickness = cv2.FONT_HERSHEY_DUPLEX, 0.6, 1
+        text_color, bg_color = (255, 255, 255), (255, 75, 75)
+        text_lines = [f"Gender: {labels['gender']}", f"Age: {labels['age']}"]
+        max_width, line_height = 0, 25
+        for line in text_lines:
+            (w_text, _), _ = cv2.getTextSize(line, font, font_scale, font_thickness)
+            if w_text > max_width: max_width = w_text
+        total_height = len(text_lines) * line_height - 5
+        cv2.rectangle(image, (x, y), (x + w, y + h), bg_color, 2)
+        cv2.rectangle(image, (x-1, y - total_height), (x + max_width + 10, y), bg_color, -1)
+        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):
+        annotated_image, predictions = image_array.copy(), []
+        face_results = self.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)
+        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

src/cnnClassifier/pipeline/stage_01_data_ingestion.py

@@ -1,3 +1,5 @@
+# src/cnnClassifier/pipeline/stage_01_data_ingestion.py
+
 from cnnClassifier.config.configuration import ConfigurationManager
 from cnnClassifier.components.data_ingestion import DataIngestion
 from cnnClassifier import logger
@@ -12,8 +14,7 @@ class DataIngestionTrainingPipeline:
         config = ConfigurationManager()
         data_ingestion_config = config.get_data_ingestion_config()
         data_ingestion = DataIngestion(config=data_ingestion_config)
-        data_ingestion.download_file()
-        data_ingestion.extract_zip_file()
+        data_ingestion.download_dataset() # Call the new method
 
 
 if __name__ == '__main__':

src/cnnClassifier/pipeline/stage_02_data_preparation.py

@@ -0,0 +1,27 @@
+# src/cnnClassifier/pipeline/stage_02_data_preparation.py
+
+from cnnClassifier.config.configuration import ConfigurationManager
+from cnnClassifier.components.data_preparation import DataPreparation
+from cnnClassifier import logger
+
+STAGE_NAME = "Data Preparation stage"
+
+class DataPreparationTrainingPipeline:
+    def __init__(self):
+        pass
+
+    def main(self):
+        config = ConfigurationManager()
+        data_preparation_config = config.get_data_preparation_config()
+        data_preparation = DataPreparation(config=data_preparation_config)
+        data_preparation.create_cleaned_dataframe()
+
+if __name__ == '__main__':
+    try:
+        logger.info(f">>>>>> stage {STAGE_NAME} started <<<<<<")
+        obj = DataPreparationTrainingPipeline()
+        obj.main()
+        logger.info(f">>>>>> stage {STAGE_NAME} completed <<<<<<\n\nx==========x")
+    except Exception as e:
+        logger.exception(e)
+        raise e

src/cnnClassifier/pipeline/stage_03_multi_task_model_training.py

@@ -0,0 +1,26 @@
+# src/cnnClassifier/pipeline/stage_03_multi_task_model_training.py
+from cnnClassifier.config.configuration import ConfigurationManager
+from cnnClassifier.components.multi_task_model_trainer import MultiTaskModelTrainer
+from cnnClassifier import logger
+
+STAGE_NAME = "Multi-Task Model Training stage"
+
+class MultiTaskModelTrainingPipeline:
+    def __init__(self):
+        pass
+
+    def main(self):
+        config = ConfigurationManager()
+        multi_task_model_trainer_config = config.get_multi_task_model_trainer_config()
+        trainer = MultiTaskModelTrainer(config=multi_task_model_trainer_config)
+        trainer.train()
+
+if __name__ == '__main__':
+    try:
+        logger.info(f">>>>>> stage {STAGE_NAME} started <<<<<<")
+        obj = MultiTaskModelTrainingPipeline()
+        obj.main()
+        logger.info(f">>>>>> stage {STAGE_NAME} completed <<<<<<\n\nx==========x")
+    except Exception as e:
+        logger.exception(e)
+        raise e