import torch import numpy as np from PIL import Image from transformers import AutoImageProcessor import cv2 from huggingface_hub import hf_hub_download 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__), '..')) 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: from src.cnnClassifier.components.multi_task_model_trainer import MultiTaskEfficientNet from src.cnnClassifier.utils.common import read_yaml 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 ---") cache_dir = os.getenv("HF_HOME", "/app/hf_cache") 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.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.model = self._load_model() 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, 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}") 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_hq(self, image_array: np.ndarray) -> (np.ndarray, list): 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'] 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 predict_lq(self, image_array: np.ndarray) -> (np.ndarray, list): 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