Update src/cnnClassifier/pipeline/prediction.py

src/cnnClassifier/pipeline/prediction.py

@@ -4,12 +4,13 @@ from PIL import Image
 from transformers import AutoImageProcessor
 import cv2
 from huggingface_hub import hf_hub_download
-from mtcnn import MTCNN  # For high-quality
+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
+import pandas as pd
 
 try:
     src_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
@@ -17,7 +18,6 @@ try:
     from components.multi_task_model_trainer import MultiTaskEfficientNet
     from utils.common import read_yaml
 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
 
@@ -25,65 +25,45 @@ class PredictionPipeline:
     def __init__(self, repo_id: str = "ALYYAN/Facial-Age-Det"):
         self.device = "cpu"
         self.repo_id = repo_id
-
         print("--- Initializing Prediction Pipeline by downloading artifacts from Hub ---")
-
-        # Define cache dir (matches Dockerfile ENV)
         cache_dir = os.getenv("HF_HOME", "/app/hf_cache")
-
-        # Download individual files → return full path
-        self.model_file = hf_hub_download(
-            repo_id=self.repo_id,
-            filename="checkpoint-26873/model.safetensors",
-            cache_dir=cache_dir
-        )
-        self.params_path = hf_hub_download(
-            repo_id=self.repo_id,
-            filename="params.yaml",
-            cache_dir=cache_dir
-        )
-        self.data_csv_path = hf_hub_download(
-            repo_id=self.repo_id,
-            filename="fairface_cleaned.csv",
-            cache_dir=cache_dir
-        )
-
+        self.model_file = hf_hub_download(repo_id=self.repo_id, filename="checkpoint-26873/model.safetensors", cache_dir=cache_dir)
+        self.params_path = hf_hub_download(repo_id=self.repo_id, filename="params.yaml", cache_dir=cache_dir)
+        self.data_csv_path = hf_hub_download(repo_id=self.repo_id, filename="fairface_cleaned.csv", cache_dir=cache_dir)
         self.base_model_name = "google/efficientnet-b2"
         self.params = read_yaml(Path(self.params_path))
+        
+        # --- THE FIX IS HERE: CALL THE METHOD THAT EXISTS ---
+        self.label_maps = self._load_label_maps_from_csv()
+        # --- END FIX ---
 
-        self.label_maps = self._load_label_maps()
         self.processor = AutoImageProcessor.from_pretrained(self.base_model_name)
-        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.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()
-
-        # Face detector
         haar_cascade_path = cv2.data.haarcascades + 'haarcascade_frontalface_default.xml'
         self.lq_face_detector = cv2.CascadeClassifier(haar_cascade_path)
         self.hq_face_detector = MTCNN()
-
         print(f"--- Pipeline Initialized Successfully on device: {self.device} ---")
 
-    
+    # --- THE MISSING METHOD ---
+    def _load_label_maps_from_csv(self):
+        print(f"Generating label maps from downloaded CSV: {self.data_csv_path}")
+        df = pd.read_csv(self.data_csv_path)
+        label_maps = {}
+        tasks = {'age': lambda x: int(str(x).split('-')[0]), 'gender': None}
+        for task, sort_key in tasks.items():
+            labels_str = [str(label) for label in df[task].unique()]
+            sorted_labels = sorted(labels_str, key=sort_key)
+            label_maps[f'{task}_id2label'] = {str(i): label for i, label in enumerate(sorted_labels)}
+        return label_maps
+    # --- END MISSING METHOD ---
+
     def _load_model(self):
-        num_age = len(self.label_maps['age_id2label'])
-        num_gender = len(self.label_maps['gender_id2label'])
-        num_race = 7
-    
+        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 = Path(self.model_file)
-        if not weight_file.exists():
-            raise FileNotFoundError(f"Weights not found: {weight_file}")
-    
-        if weight_file.suffix == ".safetensors":
-            state_dict = load_safetensors(weight_file, device="cpu")
-        else:
-            state_dict = torch.load(weight_file, map_location="cpu")
-    
+        if not weight_file.exists(): raise FileNotFoundError(f"Weights not found: {weight_file}")
+        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()
@@ -106,11 +86,9 @@ class PredictionPipeline:
             cv2.putText(image, line, (x + 5, y_text), font, font_scale, text_color, font_thickness, cv2.LINE_AA)
 
     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.hq_face_detector.detect_faces(image_array)
         if not face_results: return annotated_image, predictions
-
         for face in face_results:
             if face['confidence'] < 0.95: continue
             x, y, w, h = face['box']
@@ -129,12 +107,10 @@ class PredictionPipeline:
         return annotated_image, predictions
 
     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