Upload model.py

model.py

@@ -0,0 +1,344 @@
+"""
+Complete Age and Gender Prediction Model with Pipeline Support
+Author: Abhilash Sahoo
+License: Apache 2.0
+
+This file provides both manual usage and Hugging Face pipeline support.
+"""
+
+import torch
+import torch.nn as nn
+from transformers import (
+    ViTModel, 
+    ViTImageProcessor,
+    PreTrainedModel,
+    PretrainedConfig,
+    ImageClassificationPipeline
+)
+from PIL import Image
+import numpy as np
+from typing import Union, Dict, Any, List
+import requests
+from io import BytesIO
+
+
+class AgeGenderConfig(PretrainedConfig):
+    """Configuration class for AgeGenderViTModel"""
+    model_type = "age-gender-vit"
+    
+    def __init__(
+        self,
+        vit_model_name="google/vit-base-patch16-224",
+        hidden_size=768,
+        intermediate_size=256,
+        final_size=64,
+        dropout_rate=0.1,
+        num_age_classes=100,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.vit_model_name = vit_model_name
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.final_size = final_size
+        self.dropout_rate = dropout_rate
+        self.num_age_classes = num_age_classes
+
+
+class AgeGenderViTModel(PreTrainedModel):
+    """
+    Vision Transformer model for simultaneous age estimation and gender classification.
+    
+    Architecture: ViT-Base + Dual-head (768 → 256 → 64 → 1)
+    - Age head: Regression output for age estimation (0-100 years)
+    - Gender head: Binary classification for gender prediction (Male/Female)
+    """
+    config_class = AgeGenderConfig
+    
+    def __init__(self, config=None):
+        if config is None:
+            config = AgeGenderConfig()
+        super().__init__(config)
+        
+        # Load pre-trained ViT model
+        self.vit = ViTModel.from_pretrained(config.vit_model_name)
+        
+        # Age regression head: 768 → 256 → 64 → 1
+        self.age_head = nn.Sequential(
+            nn.Linear(config.hidden_size, config.intermediate_size),  # 768 → 256
+            nn.ReLU(),
+            nn.Dropout(config.dropout_rate),
+            nn.Linear(config.intermediate_size, config.final_size),   # 256 → 64
+            nn.ReLU(),
+            nn.Dropout(config.dropout_rate),
+            nn.Linear(config.final_size, 1)                          # 64 → 1
+        )
+        
+        # Gender classification head: 768 → 256 → 64 → 1
+        self.gender_head = nn.Sequential(
+            nn.Linear(config.hidden_size, config.intermediate_size),  # 768 → 256
+            nn.ReLU(),
+            nn.Dropout(config.dropout_rate),
+            nn.Linear(config.intermediate_size, config.final_size),   # 256 → 64
+            nn.ReLU(),
+            nn.Dropout(config.dropout_rate),
+            nn.Linear(config.final_size, 1),                         # 64 → 1
+            nn.Sigmoid()
+        )
+        
+        # For pipeline compatibility, add a dummy classifier
+        self.classifier = nn.Linear(2, 2)  # Dummy layer for pipeline compatibility
+        
+    def forward(self, pixel_values, **kwargs):
+        """Forward pass through the model"""
+        # Extract features using ViT
+        vit_outputs = self.vit(pixel_values=pixel_values)
+        pooled_output = vit_outputs.pooler_output
+        
+        # Get predictions from both heads
+        age_output = self.age_head(pooled_output)
+        gender_output = self.gender_head(pooled_output)
+        
+        # For pipeline compatibility, create fake logits
+        fake_logits = torch.cat([age_output, gender_output], dim=1)
+        
+        # Return object with both pipeline format and raw outputs
+        return type('ModelOutput', (), {
+            'logits': fake_logits,
+            'age_logits': age_output,
+            'gender_logits': gender_output
+        })()
+
+
+class AgeGenderImageClassificationPipeline(ImageClassificationPipeline):
+    """Custom pipeline for age-gender classification"""
+    
+    def postprocess(self, model_outputs, top_k=1, **kwargs):
+        """Custom postprocessing for age-gender predictions"""
+        if isinstance(model_outputs, list):
+            outputs = model_outputs[0]  # Single image output
+        else:
+            outputs = model_outputs
+        
+        # Extract age and gender logits
+        age_logits = outputs.age_logits
+        gender_logits = outputs.gender_logits
+        
+        # Process predictions
+        age = int(torch.clamp(age_logits, 0, 100).item())
+        gender_prob = gender_logits.item()
+        gender = "Female" if gender_prob > 0.5 else "Male"
+        confidence = gender_prob if gender_prob > 0.5 else 1 - gender_prob
+        
+        # Return in pipeline format
+        return [{
+            "label": f"{age} years, {gender}",
+            "score": confidence,
+            "age": age,
+            "gender": gender,
+            "gender_confidence": round(confidence, 3),
+            "gender_probability_female": round(gender_prob, 3),
+            "gender_probability_male": round(1 - gender_prob, 3)
+        }]
+
+
+def create_model_and_processor():
+    """
+    Create model and processor for manual usage
+    
+    Returns:
+        tuple: (model, processor)
+    """
+    # Create model
+    model = AgeGenderViTModel()
+    
+    # Load pre-trained weights
+    model_url = "https://huggingface.co/abhilash88/age-gender-prediction/resolve/main/pytorch_model.bin"
+    weights = torch.hub.load_state_dict_from_url(model_url, map_location='cpu')
+    model.load_state_dict(weights)
+    model.eval()
+    
+    # Create processor
+    processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
+    
+    return model, processor
+
+
+def predict_age_gender(
+    image: Union[str, Image.Image, np.ndarray],
+    model=None,
+    processor=None,
+    device='auto'
+) -> Dict[str, Any]:
+    """
+    Predict age and gender from an image (Manual approach)
+    
+    Args:
+        image: Image path, URL, PIL Image, or numpy array
+        model: Pre-loaded model (optional)
+        processor: Pre-loaded processor (optional)
+        device: Device to use ('auto', 'cpu', 'cuda')
+    
+    Returns:
+        Dictionary with predictions
+    """
+    # Auto-detect device
+    if device == 'auto':
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    
+    # Load model and processor if not provided
+    if model is None or processor is None:
+        model, processor = create_model_and_processor()
+    
+    model = model.to(device)
+    
+    # Handle different input types
+    if isinstance(image, str):
+        if image.startswith(('http://', 'https://')):
+            # Download image from URL
+            response = requests.get(image)
+            image = Image.open(BytesIO(response.content)).convert('RGB')
+        else:
+            # Load local image
+            image = Image.open(image).convert('RGB')
+    elif isinstance(image, np.ndarray):
+        image = Image.fromarray(image).convert('RGB')
+    elif not isinstance(image, Image.Image):
+        raise ValueError(f"Unsupported image type: {type(image)}")
+    
+    # Process image
+    inputs = processor(images=image, return_tensors="pt")
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+    
+    # Make prediction
+    with torch.no_grad():
+        outputs = model(**inputs)
+        age_pred = outputs.age_logits
+        gender_pred = outputs.gender_logits
+    
+    # Post-process predictions
+    age = int(torch.clamp(age_pred, 0, 100).item())
+    gender_prob = gender_pred.item()
+    gender = "Female" if gender_prob > 0.5 else "Male"
+    confidence = gender_prob if gender_prob > 0.5 else 1 - gender_prob
+    
+    return {
+        "age": age,
+        "gender": gender,
+        "gender_confidence": round(confidence, 3),
+        "gender_probability_female": round(gender_prob, 3),
+        "gender_probability_male": round(1 - gender_prob, 3),
+        "summary": f"{age} years, {gender} ({confidence:.1%} confidence)"
+    }
+
+
+def predict_age_gender_pipeline(image_input: Union[str, Image.Image]) -> Dict[str, Any]:
+    """
+    Predict using Hugging Face pipeline (requires proper repo setup)
+    
+    Args:
+        image_input: Image path, URL, or PIL Image
+        
+    Returns:
+        Dictionary with predictions
+    """
+    from transformers import pipeline
+    
+    try:
+        # Create pipeline
+        classifier = pipeline(
+            "image-classification",
+            model="abhilash88/age-gender-prediction",
+            trust_remote_code=True
+        )
+        
+        # Make prediction
+        result = classifier(image_input)[0]  # Get first result
+        
+        return {
+            "age": result["age"],
+            "gender": result["gender"], 
+            "confidence": result["gender_confidence"],
+            "summary": result["label"]
+        }
+    
+    except Exception as e:
+        print(f"Pipeline failed: {e}")
+        print("Falling back to manual approach...")
+        return predict_age_gender(image_input)
+
+
+# Simple usage functions
+def simple_predict(image_path: str) -> str:
+    """
+    Simplest possible usage - just returns a string
+    
+    Args:
+        image_path: Path to image or URL
+        
+    Returns:
+        String with prediction
+    """
+    result = predict_age_gender(image_path)
+    return result["summary"]
+
+
+def batch_predict(image_list: List[str]) -> List[Dict[str, Any]]:
+    """
+    Predict on multiple images
+    
+    Args:
+        image_list: List of image paths or URLs
+        
+    Returns:
+        List of prediction dictionaries
+    """
+    # Load model once for efficiency
+    model, processor = create_model_and_processor()
+    
+    results = []
+    for image_path in image_list:
+        try:
+            result = predict_age_gender(image_path, model, processor)
+            result["image_path"] = image_path
+            results.append(result)
+        except Exception as e:
+            results.append({
+                "image_path": image_path,
+                "error": str(e),
+                "age": None,
+                "gender": None
+            })
+    
+    return results
+
+
+# Example usage and testing
+if __name__ == "__main__":
+    print("🚀 Testing Age-Gender Prediction Model...")
+    
+    try:
+        # Test simple prediction
+        print("📝 Testing simple prediction...")
+        test_url = "https://images.unsplash.com/photo-1507003211169-0a1dd7228f2d?w=300&h=300&fit=crop&crop=face"
+        
+        # Method 1: Simple string output
+        simple_result = simple_predict(test_url)
+        print(f"✅ Simple: {simple_result}")
+        
+        # Method 2: Detailed output
+        detailed_result = predict_age_gender(test_url)
+        print(f"✅ Detailed: {detailed_result}")
+        
+        # Method 3: Try pipeline (may fail if repo not updated)
+        try:
+            pipeline_result = predict_age_gender_pipeline(test_url)
+            print(f"✅ Pipeline: {pipeline_result}")
+        except:
+            print("❌ Pipeline not working yet (needs repo file updates)")
+        
+        print("🎉 Model is working perfectly!")
+        
+    except Exception as e:
+        print(f"❌ Error: {e}")
+        print("Note: This test requires internet connection for test image")