ml_server.py

#!/usr/bin/env python3
"""
SIH Certificate Classifier API Server
====================================

Flask server that provides certificate authenticity verification
using the trained MobileNetV2 model from CNN_SIH.ipynb

Endpoints:
- POST /predict - Analyze certificate authenticity
- GET /health - Health check
- GET /model-info - Model information
"""

import os
import io
import base64
import numpy as np
from PIL import Image
import tensorflow as tf
from tensorflow.keras.models import load_model
from flask import Flask, request, jsonify
from flask_cors import CORS
import cv2
import json
from datetime import datetime

app = Flask(__name__)

# Configure CORS via env var (comma-separated origins). Defaults to '*'.
ALLOWED_ORIGINS_ENV = os.getenv('ALLOWED_ORIGINS', '*')
if ALLOWED_ORIGINS_ENV.strip() == '*':
    _cors = CORS(app, resources={r"/*": {"origins": "*"}}, supports_credentials=False)
    _allowed_origins_list = None  # wildcard
else:
    _allowed_origins_list = [o.strip() for o in ALLOWED_ORIGINS_ENV.split(',') if o.strip()]
    _cors = CORS(app, resources={r"/*": {"origins": _allowed_origins_list}}, supports_credentials=False)

@app.after_request
def add_cors_headers(response):
    try:
        origin = request.headers.get('Origin', '')
        if _allowed_origins_list is None:
            response.headers['Access-Control-Allow-Origin'] = '*'
        else:
            if origin in _allowed_origins_list:
                response.headers['Access-Control-Allow-Origin'] = origin
                response.headers['Vary'] = 'Origin'
        response.headers['Access-Control-Allow-Headers'] = 'Content-Type, Authorization, X-Requested-With'
        response.headers['Access-Control-Allow-Methods'] = 'GET, POST, OPTIONS'
    except Exception:
        pass
    return response

class SIHCertificateClassifier:
    def __init__(self):
        self.model = None
        self.class_indices = None
        self.model_loaded = False
        self.input_size = (224, 224)
        self.channels = 3
        self.uses_mobilenet_preprocessing = False
        self.real_index = 1  # default to 1=REAL (per cnn_final)
        self.fake_index = 0  # default to 0=FAKE (per cnn_final)
        self.load_model()
    
    def load_model(self):
        """Load the SIH MobileNetV2 certificate classifier"""
        try:
            # Resolve paths (support running from anywhere)
            base_dir = os.path.dirname(os.path.abspath(__file__))
            candidate_paths = [
                os.path.join(base_dir, 'certificate_classifier.h5'),
                os.path.join(base_dir, 'build', 'certificate_classifier.h5'),
            ]
            model_path = next((p for p in candidate_paths if os.path.exists(p)), None)
            indices_candidates = [
                os.path.join(base_dir, 'class_indices.npy'),
                os.path.join(base_dir, 'class_indices.json'),
                os.path.join(base_dir, 'build', 'class_indices.npy'),
                os.path.join(base_dir, 'build', 'class_indices.json'),
            ]
            indices_path = next((p for p in indices_candidates if os.path.exists(p)), None)
            
            if not model_path:
                print(f"❌ Model file not found in candidates: {candidate_paths}")
                print("📝 Please ensure certificate_classifier.h5 is in the same directory")
                return False
            
            print(f"🔄 Loading model from: {model_path}")
            self.model = load_model(model_path)
            print("✅ Model loaded successfully")
            
            # Load class indices if available
            if indices_path and indices_path.endswith('.npy'):
                self.class_indices = np.load(indices_path, allow_pickle=True).item()
                print("✅ Class indices loaded (npy):", self.class_indices)
            elif indices_path and indices_path.endswith('.json'):
                with open(indices_path, 'r', encoding='utf-8') as f:
                    self.class_indices = json.load(f)
                print("✅ Class indices loaded (json):", self.class_indices)
            else:
                # Default mapping aligned with cnn_final: 0=FAKE, 1=REAL
                self.class_indices = {'0_FAKE_CERT': 0, '1_REAL_CERT': 1}
                print("📝 Using default class indices (cnn_final):", self.class_indices)

            # Resolve real/fake index from mapping or env
            self._resolve_class_order()
            
            self.model_loaded = True
            
            # Print model info
            print("📊 Model Architecture:")
            print(f"   - Input shape: {self.model.input_shape}")
            print(f"   - Output shape: {self.model.output_shape}")
            print(f"   - Total parameters: {self.model.count_params():,}")
            print(f"   - Class order: REAL={self.real_index}, FAKE={self.fake_index}")
            
            # Infer input size and channels
            try:
                ishape = self.model.input_shape
                # Handle list of inputs vs single input
                if isinstance(ishape, list):
                    ishape = ishape[0]
                # ishape example: (None, H, W, C) or (None, C, H, W)
                if len(ishape) == 4:
                    if ishape[-1] in (1, 3):
                        self.input_size = (int(ishape[1] or 224), int(ishape[2] or 224))
                        self.channels = int(ishape[3])
                    elif ishape[1] in (1, 3):
                        # channels_first
                        self.input_size = (int(ishape[2] or 224), int(ishape[3] or 224))
                        self.channels = int(ishape[1])
                print(f"   - Inferred input size: {self.input_size}, channels: {self.channels}")
            except Exception as ie:
                print(f"⚠️ Failed to infer input size, using default (224,224,3): {ie}")
                self.input_size = (224, 224)
                self.channels = 3

            # Heuristic: use MobileNetV2 preprocessing if model name hints
            model_name = getattr(self.model, 'name', '')
            if 'mobilenet' in model_name.lower() or 'v2' in model_name.lower():
                self.uses_mobilenet_preprocessing = True
            print(f"   - Preprocessing: {'MobileNetV2 preprocess_input' if self.uses_mobilenet_preprocessing else 'normalize [0,1]'}")
            
            return True
            
        except Exception as e:
            print(f"❌ Failed to load SIH model: {e}")
            return False

    def _resolve_class_order(self):
        """Resolve class indices for REAL and FAKE from available mapping or env."""
        # Try to detect from class_indices keys
        real_idx = None
        fake_idx = None
        try:
            if isinstance(self.class_indices, dict):
                for k, v in self.class_indices.items():
                    ku = str(k).upper()
                    if 'REAL' in ku and real_idx is None:
                        real_idx = int(v)
                    if 'FAKE' in ku and fake_idx is None:
                        fake_idx = int(v)
        except Exception:
            pass

        # If still missing, check env CLASS_ORDER
        if real_idx is None or fake_idx is None:
            order = os.getenv('CLASS_ORDER', 'FAKE_REAL').upper()
            if order == 'REAL_FAKE':
                real_idx, fake_idx = 0, 1
            else:
                # default FAKE_REAL matches cnn_final
                real_idx, fake_idx = 1, 0

        self.real_index = real_idx
        self.fake_index = fake_idx
        print(f"🔎 Resolved class order => REAL: {self.real_index}, FAKE: {self.fake_index}")
    
    def preprocess_image(self, image):
        """
        Preprocess image for SIH model (matches CNN_SIH.ipynb preprocessing)
        """
        try:
            # Convert PIL to numpy array
            img_array = np.array(image)
            
            # Convert RGB to BGR if needed (for cv2 compatibility)
            if len(img_array.shape) == 3 and img_array.shape[2] == 3:
                img_array = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
                img_array = cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB)
            
            # Resize to model's expected input size
            img_resized = cv2.resize(img_array, self.input_size)
            img_resized = img_resized.astype('float32')
            
            # If model expects grayscale, convert
            if self.channels == 1 and len(img_resized.shape) == 3 and img_resized.shape[2] == 3:
                img_resized = cv2.cvtColor(img_resized, cv2.COLOR_RGB2GRAY)
                # add channel back
                img_resized = np.expand_dims(img_resized, axis=-1)

            # Preprocess
            if self.uses_mobilenet_preprocessing:
                try:
                    from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
                    img_normalized = preprocess_input(img_resized)
                except Exception as _:
                    print("⚠️ MobileNetV2 preprocess_input unavailable; falling back to [0,1] normalization")
                    img_normalized = img_resized / 255.0
            else:
                img_normalized = img_resized / 255.0
            
            # Add batch dimension
            img_batch = np.expand_dims(img_normalized, axis=0)
            
            print(f"🖼️ Image preprocessed: {img_batch.shape}")
            return img_batch
            
        except Exception as e:
            print(f"❌ Image preprocessing failed: {e}")
            return None
    
    def predict_authenticity(self, image):
        """
        Predict certificate authenticity using SIH model
        """
        if not self.model_loaded:
            return self.generate_mock_prediction()
        
        try:
            # Preprocess image
            processed_img = self.preprocess_image(image)
            if processed_img is None:
                raise ValueError("Image preprocessing failed")
            
            # Make prediction
            print("🔄 Running SIH model prediction...")
            prediction = self.model.predict(processed_img, verbose=0)
            
            print(f"📊 Raw prediction: {prediction}")
            
            # Support both binary-sigmoid (1 unit) and softmax (2 units)
            out = prediction[0]
            if out.shape[-1] == 1:
                # Binary sigmoid: treat value as P(REAL)
                real_prob = float(out[0])
                fake_prob = float(1.0 - real_prob)
                is_authentic = real_prob >= 0.5
                predicted_index = 0 if is_authentic else 1
                confidence = real_prob if is_authentic else fake_prob
            else:
                # Softmax with 2 units: [REAL, FAKE]
                # Use resolved indices
                real_prob = float(out[self.real_index])
                fake_prob = float(out[self.fake_index])
                predicted_index = int(np.argmax(out))
                is_authentic = predicted_index == self.real_index
                confidence = float(np.max(out))
            
            # Map index to label (auto-fix swapped folder names like in SIH notebook)
            idx_to_label = {int(v): str(k) for k, v in self.class_indices.items()}
            # Normalize to simple labels
            predicted_label = 'REAL_CERT' if is_authentic else 'FAKE_CERT'
            
            # Determine confidence level
            confidence_level = 'HIGH' if confidence > 0.85 else 'MODERATE' if confidence > 0.65 else 'LOW'
            
            result = {
                'authentic': bool(is_authentic),
                'prediction': predicted_label,
                'confidence': float(confidence),
                'confidence_level': confidence_level,
                'real_probability': float(real_prob),
                'fake_probability': float(fake_prob),
                'predicted_index': int(predicted_index),
                'model_type': getattr(self.model, 'name', 'SIH_Model'),
                'timestamp': datetime.now().isoformat(),
                'mock_result': False
            }
            
            print(f"🎯 SIH Prediction: {predicted_label} ({confidence:.3f} confidence)")
            return result
            
        except Exception as e:
            print(f"❌ SIH prediction failed: {e}")
            return self.generate_mock_prediction(error=str(e))
    
    def generate_mock_prediction(self, error=None):
        """Generate mock prediction when model unavailable"""
        confidence = 0.75 + np.random.random() * 0.2  # 0.75-0.95 range
        is_authentic = np.random.random() > 0.3  # 70% authentic rate
        
        return {
            'authentic': bool(is_authentic),
            'prediction': 'REAL_CERT' if is_authentic else 'FAKE_CERT',
            'confidence': float(confidence),
            'confidence_level': 'MODERATE',
            'real_probability': float(confidence if is_authentic else 1 - confidence),
            'fake_probability': float(1 - confidence if is_authentic else confidence),
            'predicted_index': 0 if is_authentic else 1,
            'model_type': 'Mock_SIH_Model',
            'timestamp': datetime.now().isoformat(),
            'mock_result': True,
            'error': error
        }
    
    def get_model_info(self):
        """Get model information"""
        if not self.model_loaded:
            return {
                'loaded': False,
                'error': 'SIH model not loaded',
                'expected_files': ['certificate_classifier.h5', 'class_indices.npy|class_indices.json']
            }
        
        return {
            'loaded': True,
            'model_type': getattr(self.model, 'name', 'SIH_Model'),
            'input_shape': str(self.model.input_shape),
            'output_shape': str(self.model.output_shape),
            'parameters': self.model.count_params(),
            'classes': self.class_indices,
            'architecture': 'CNN per cnn_final; auto-detected',
            'training_details': {
                'optimizer': 'Adam(0.0001)',
                'loss': 'categorical_crossentropy or binary_crossentropy',
                'input_size': f'{self.input_size[0]}x{self.input_size[1]}x{self.channels}',
                'preprocessing': 'mobilenet_v2.preprocess_input' if self.uses_mobilenet_preprocessing else 'normalize [0,1]'
            }
        }

# Initialize classifier
classifier = SIHCertificateClassifier()

@app.route('/health', methods=['GET'])
def health_check():
    """Health check endpoint"""
    return jsonify({
        'status': 'healthy',
        'service': 'SIH Certificate Classifier API',
        'model_loaded': classifier.model_loaded,
        'timestamp': datetime.now().isoformat()
    })

@app.route('/model-info', methods=['GET'])
def model_info():
    """Get model information"""
    return jsonify(classifier.get_model_info())

@app.route('/predict', methods=['POST', 'OPTIONS'])
def predict_certificate():
    """
    Predict certificate authenticity
    
    Expected input:
    - multipart/form-data with 'image' file
    - or JSON with base64 encoded image
    """
    try:
        # Handle CORS preflight
        if request.method == 'OPTIONS':
            return ('', 204)
        image = None
        
        # Handle multipart file upload
        if 'image' in request.files:
            file = request.files['image']
            if file.filename == '':
                return jsonify({'error': 'No file selected'}), 400
            
            # Validate file type
            if not file.filename.lower().endswith(('.png', '.jpg', '.jpeg')):
                return jsonify({'error': 'Invalid file type. Please upload PNG, JPG, or JPEG'}), 400
            
            try:
                image = Image.open(file.stream)
                # Convert to RGB if needed
                if image.mode != 'RGB':
                    image = image.convert('RGB')
            except Exception as e:
                return jsonify({'error': f'Invalid image file: {str(e)}'}), 400
        
        # Handle JSON base64 upload
        elif request.is_json:
            data = request.get_json()
            if 'image' not in data:
                return jsonify({'error': 'No image data provided'}), 400
            
            try:
                # Decode base64 image
                image_data = base64.b64decode(data['image'].split(',')[1])  # Remove data:image/jpeg;base64, prefix
                image = Image.open(io.BytesIO(image_data))
                if image.mode != 'RGB':
                    image = image.convert('RGB')
            except Exception as e:
                return jsonify({'error': f'Invalid base64 image: {str(e)}'}), 400
        
        else:
            return jsonify({'error': 'No image provided. Send as multipart/form-data or JSON base64'}), 400
        
        # Run prediction
        print(f"🔍 Analyzing certificate image: {image.size}")
        result = classifier.predict_authenticity(image)
        
        # Add API metadata
        result.update({
            'api_version': '1.0.0',
            'processing_time': datetime.now().isoformat(),
            'image_size': f"{image.size[0]}x{image.size[1]}"
        })
        
        return jsonify(result)
        
    except Exception as e:
        print(f"❌ API error: {e}")
        return jsonify({
            'error': f'Internal server error: {str(e)}',
            'timestamp': datetime.now().isoformat()
        }), 500

@app.errorhandler(404)
def not_found(error):
    return jsonify({
        'error': 'Endpoint not found',
        'available_endpoints': {
            'GET /health': 'Health check',
            'GET /model-info': 'Model information',
            'POST /predict': 'Certificate authenticity prediction'
        }
    }), 404

@app.errorhandler(500)
def internal_error(error):
    return jsonify({
        'error': 'Internal server error',
        'timestamp': datetime.now().isoformat()
    }), 500

if __name__ == '__main__':
    print("🚀 Starting SIH Certificate Classifier API Server")
    print("=" * 50)
    print("📋 Available endpoints:")
    print("   GET  /health     - Health check")
    print("   GET  /model-info - Model information") 
    print("   POST /predict    - Certificate authentication")
    print("=" * 50)
    
    # Configuration
    host = os.getenv('HOST', '0.0.0.0')
    port = int(os.getenv('PORT', 5000))
    debug = os.getenv('DEBUG', 'False').lower() == 'true'
    
    print(f"🌐 Server starting on http://{host}:{port}")
    if classifier.model_loaded:
        print("✅ SIH model loaded and ready!")
    else:
        print("⚠️  Running with mock predictions (model not found)")
    print("=" * 50)
    
    app.run(host=host, port=port, debug=debug)