abhilash88
/

age-gender-prediction

@@ -1,242 +0,0 @@
-"""
-Complete Working Pipeline Implementation for Age-Gender Prediction
-This enables: pipeline("image-classification", model="abhilash88/age-gender-prediction", trust_remote_code=True)
-"""
-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 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):
-    """Age-Gender ViT Model with pipeline support"""
-    config_class = AgeGenderConfig
-    def __init__(self, config=None):
-        if config is None:
-            config = AgeGenderConfig()
-        super().__init__(config)
-        # ViT backbone
-        self.vit = ViTModel.from_pretrained(config.vit_model_name)
-        # Age head: 768 → 256 → 64 → 1
-        self.age_head = nn.Sequential(
-            nn.Linear(config.hidden_size, config.intermediate_size),
-            nn.ReLU(),
-            nn.Dropout(config.dropout_rate),
-            nn.Linear(config.intermediate_size, config.final_size),
-            nn.ReLU(),
-            nn.Dropout(config.dropout_rate),
-            nn.Linear(config.final_size, 1)
-        )
-        # Gender head: 768 → 256 → 64 → 1
-        self.gender_head = nn.Sequential(
-            nn.Linear(config.hidden_size, config.intermediate_size),
-            nn.ReLU(),
-            nn.Dropout(config.dropout_rate),
-            nn.Linear(config.intermediate_size, config.final_size),
-            nn.ReLU(),
-            nn.Dropout(config.dropout_rate),
-            nn.Linear(config.final_size, 1),
-            nn.Sigmoid()
-        )
-        # For pipeline compatibility, we need a classifier head that outputs logits
-        self.classifier = nn.Linear(2, 2)  # Dummy classifier for pipeline
-    def forward(self, pixel_values, **kwargs):
-        """Forward pass - returns format expected by pipeline"""
-        outputs = self.vit(pixel_values=pixel_values)
-        pooled_output = outputs.pooler_output
-        # Get age and gender predictions
-        age_output = self.age_head(pooled_output)
-        gender_output = self.gender_head(pooled_output)
-        # For pipeline compatibility, create fake logits
-        # We'll process these in the pipeline postprocessing
-        fake_logits = torch.cat([age_output, gender_output], dim=1)
-        return type('ModelOutput', (), {
-            'logits': fake_logits,
-            'age_logits': age_output,
-            'gender_logits': gender_output
-        })()
-class AgeGenderImageClassificationPipeline(ImageClassificationPipeline):
-    """Custom pipeline for age-gender classification"""
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-    def postprocess(self, model_outputs, top_k=1, **kwargs):
-        """Custom postprocessing for age-gender predictions"""
-        outputs = model_outputs[0]  # Single image output
-        # 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": round(gender_prob, 3)
-        }]
-# Simple wrapper function for easy usage
-def predict_age_gender_pipeline(image_path_or_url: str) -> Dict[str, Any]:
-    """
-    Simple function using the pipeline
-    Args:
-        image_path_or_url: Path to image file or URL
-    Returns:
-        Dictionary with age and gender predictions
-    """
-    from transformers import pipeline
-    # Create pipeline
-    classifier = pipeline(
-        "image-classification",
-        model="abhilash88/age-gender-prediction",
-        trust_remote_code=True
-    )
-    # Make prediction
-    result = classifier(image_path_or_url)[0]  # Get first result
-    return {
-        "age": result["age"],
-        "gender": result["gender"],
-        "confidence": result["gender_confidence"],
-        "summary": result["label"]
-    }
-# Manual implementation (guaranteed to work)
-def predict_age_gender_manual(image_input: Union[str, Image.Image, np.ndarray]) -> Dict[str, Any]:
-    """
-    Manual prediction function (always works)
-    Args:
-        image_input: Image path, URL, PIL Image, or numpy array
-    Returns:
-        Dictionary with predictions
-    """
-    # Create model
-    model = AgeGenderViTModel()
-    # Load 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")
-    # Handle different input types
-    if isinstance(image_input, str):
-        if image_input.startswith(('http://', 'https://')):
-            response = requests.get(image_input)
-            image = Image.open(BytesIO(response.content)).convert('RGB')
-        else:
-            image = Image.open(image_input).convert('RGB')
-    elif isinstance(image_input, np.ndarray):
-        image = Image.fromarray(image_input).convert('RGB')
-    else:
-        image = image_input.convert('RGB')
-    # Process and predict
-    inputs = processor(images=image, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model(inputs["pixel_values"])
-        age_pred = outputs.age_logits
-        gender_pred = outputs.gender_logits
-    # Process results
-    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,
-        "confidence": round(confidence, 3),
-        "gender_probability": round(gender_prob, 3)
-    }
-# Test both approaches
-if __name__ == "__main__":
-    print("🧪 Testing Age-Gender Prediction...")
-    # Test image URL
-    test_url = "https://images.unsplash.com/photo-1507003211169-0a1dd7228f2d?w=300&h=300&fit=crop&crop=face"
-    try:
-        print("🔧 Testing manual approach...")
-        result_manual = predict_age_gender_manual(test_url)
-        print(f"✅ Manual result: {result_manual}")
-    except Exception as e:
-        print(f"❌ Manual failed: {e}")
-    try:
-        print("🚀 Testing pipeline approach...")
-        result_pipeline = predict_age_gender_pipeline(test_url)
-        print(f"✅ Pipeline result: {result_pipeline}")
-    except Exception as e:
-        print(f"❌ Pipeline failed: {e}")
-        print("Note: Pipeline requires the files to be uploaded to your HF repo")