Upload 2 files

app.py

@@ -0,0 +1,450 @@
+"""
+Project Phoenix - Cervical Cancer Cell Classification API
+Flask application for running inference on ConvNeXt V2 model from Hugging Face
+with explainability features (GRAD-CAM).
+"""
+
+import os
+import io
+import base64
+import numpy as np
+import cv2
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+# Flask
+from flask import Flask, request, jsonify
+from flask_cors import CORS
+from werkzeug.utils import secure_filename
+
+# Deep Learning
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+import torchvision.transforms as T
+
+# Transformers
+from transformers import (
+    ConvNextV2ForImageClassification,
+    AutoImageProcessor
+)
+
+# GRAD-CAM
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.image import show_cam_on_image
+from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
+
+# ========== CONFIGURATION ==========
+
+# Update this with your Hugging Face model ID
+# Example: "Meet2304/convnextv2-cervical-cell-classification"
+HF_MODEL_ID = os.getenv("HF_MODEL_ID", "Meet2304/convnextv2-cervical-cell-classification")
+
+# Class names
+CLASS_NAMES = [
+    'im_Dyskeratotic',
+    'im_Koilocytotic',
+    'im_Metaplastic',
+    'im_Parabasal',
+    'im_Superficial-Intermediate'
+]
+
+# Display names (cleaner for UI)
+DISPLAY_NAMES = [
+    'Dyskeratotic',
+    'Koilocytotic',
+    'Metaplastic',
+    'Parabasal',
+    'Superficial-Intermediate'
+]
+
+# Image preprocessing
+IMG_SIZE = 224
+ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'bmp'}
+MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB
+
+# Device
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+# ========== FLASK APP SETUP ==========
+
+app = Flask(__name__)
+CORS(app)  # Enable CORS for Next.js frontend
+
+app.config['MAX_CONTENT_LENGTH'] = MAX_FILE_SIZE
+
+# ========== MODEL LOADING ==========
+
+print("Loading model from Hugging Face...")
+print(f"Model ID: {HF_MODEL_ID}")
+print(f"Device: {DEVICE}")
+
+# Load image processor
+processor = AutoImageProcessor.from_pretrained(HF_MODEL_ID)
+print("✓ Processor loaded")
+
+# Load model
+model = ConvNextV2ForImageClassification.from_pretrained(HF_MODEL_ID)
+model = model.to(DEVICE)
+model.eval()
+print("✓ Model loaded and set to evaluation mode")
+
+print(f"Model configuration:")
+print(f"  - Number of classes: {model.config.num_labels}")
+print(f"  - Image size: {model.config.image_size}")
+print(f"  - Total parameters: {sum(p.numel() for p in model.parameters()):,}")
+
+# ========== HELPER FUNCTIONS ==========
+
+def allowed_file(filename: str) -> bool:
+    """Check if file extension is allowed."""
+    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
+
+
+def preprocess_image(image: Image.Image) -> Tuple[torch.Tensor, np.ndarray]:
+    """
+    Preprocess image for model input.
+
+    Returns:
+        Tuple of (preprocessed_tensor, original_image_array)
+    """
+    # Store original for visualization
+    original_image = np.array(image.convert('RGB'))
+
+    # Preprocess using the model's processor
+    inputs = processor(images=image, return_tensors="pt")
+    pixel_values = inputs['pixel_values'].to(DEVICE)
+
+    return pixel_values, original_image
+
+
+def predict_image(pixel_values: torch.Tensor, top_k: int = 5) -> Dict:
+    """
+    Make prediction on preprocessed image.
+
+    Args:
+        pixel_values: Preprocessed image tensor
+        top_k: Number of top predictions to return
+
+    Returns:
+        Dictionary with prediction results
+    """
+    model.eval()
+    with torch.no_grad():
+        outputs = model(pixel_values)
+        logits = outputs.logits
+
+    # Get probabilities
+    probabilities = F.softmax(logits, dim=-1)[0]
+
+    # Get top-k predictions
+    top_probs, top_indices = torch.topk(probabilities, k=min(top_k, len(CLASS_NAMES)))
+
+    # Get predicted class
+    predicted_class_idx = logits.argmax(-1).item()
+    predicted_class_name = DISPLAY_NAMES[predicted_class_idx]
+    predicted_confidence = probabilities[predicted_class_idx].item()
+
+    # Prepare results
+    results = {
+        'predicted_class': predicted_class_name,
+        'predicted_class_raw': CLASS_NAMES[predicted_class_idx],
+        'predicted_idx': predicted_class_idx,
+        'confidence': float(predicted_confidence),
+        'top_k_predictions': [
+            {
+                'class': DISPLAY_NAMES[idx],
+                'class_raw': CLASS_NAMES[idx],
+                'probability': float(prob)
+            }
+            for idx, prob in zip(top_indices, top_probs)
+        ],
+        'all_probabilities': {
+            DISPLAY_NAMES[i]: float(prob)
+            for i, prob in enumerate(probabilities)
+        }
+    }
+
+    return results
+
+
+class ConvNeXtGradCAMWrapper(nn.Module):
+    """Wrapper for ConvNeXtV2ForImageClassification to make it compatible with GRAD-CAM."""
+
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+
+    def forward(self, x):
+        outputs = self.model(pixel_values=x)
+        return outputs.logits
+
+
+def get_target_layers(model):
+    """Get the target layers for GRAD-CAM from ConvNeXt model."""
+    return [model.convnextv2.encoder.stages[-1].layers[-1]]
+
+
+def apply_gradcam(
+    pixel_values: torch.Tensor,
+    original_image: np.ndarray,
+    target_class: Optional[int] = None
+) -> Dict:
+    """
+    Apply GRAD-CAM to visualize model attention.
+
+    Args:
+        pixel_values: Preprocessed image tensor
+        original_image: Original image as numpy array
+        target_class: Target class index (None for predicted class)
+
+    Returns:
+        Dictionary with GRAD-CAM visualization and metadata
+    """
+    # Wrap the model
+    wrapped_model = ConvNeXtGradCAMWrapper(model)
+
+    # Get target layers
+    target_layers = get_target_layers(model)
+
+    # Initialize GRAD-CAM
+    cam = GradCAM(model=wrapped_model, target_layers=target_layers)
+
+    # Get prediction
+    model.eval()
+    with torch.no_grad():
+        outputs = model(pixel_values)
+        logits = outputs.logits
+        predicted_class = logits.argmax(-1).item()
+        probabilities = F.softmax(logits, dim=-1)[0]
+
+    # Use predicted class if target not specified
+    if target_class is None:
+        target_class = predicted_class
+
+    # Create target for GRAD-CAM
+    targets = [ClassifierOutputTarget(target_class)]
+
+    # Generate GRAD-CAM
+    grayscale_cam = cam(input_tensor=pixel_values, targets=targets)
+    grayscale_cam = grayscale_cam[0, :]
+
+    # Resize original image to match CAM dimensions
+    cam_h, cam_w = grayscale_cam.shape
+    rgb_image_for_overlay = cv2.resize(original_image, (cam_w, cam_h)).astype(np.float32) / 255.0
+
+    # Create visualization
+    visualization = show_cam_on_image(
+        rgb_image_for_overlay,
+        grayscale_cam,
+        use_rgb=True,
+        colormap=cv2.COLORMAP_JET
+    )
+
+    return {
+        'grayscale_cam': grayscale_cam,
+        'visualization': visualization,
+        'predicted_class': predicted_class,
+        'target_class': target_class,
+        'confidence': float(probabilities[predicted_class].item())
+    }
+
+
+def encode_image_to_base64(image_array: np.ndarray) -> str:
+    """Convert numpy array to base64 encoded PNG."""
+    # Convert to PIL Image
+    if image_array.dtype != np.uint8:
+        image_array = (image_array * 255).astype(np.uint8)
+
+    img = Image.fromarray(image_array)
+
+    # Save to bytes buffer
+    buffer = io.BytesIO()
+    img.save(buffer, format='PNG')
+    buffer.seek(0)
+
+    # Encode to base64
+    img_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
+    return f"data:image/png;base64,{img_base64}"
+
+
+# ========== API ENDPOINTS ==========
+
+@app.route('/health', methods=['GET'])
+def health_check():
+    """Health check endpoint."""
+    return jsonify({
+        'status': 'healthy',
+        'model_loaded': model is not None,
+        'device': str(DEVICE),
+        'model_id': HF_MODEL_ID
+    })
+
+
+@app.route('/predict', methods=['POST'])
+def predict():
+    """
+    Predict cervical cell classification.
+
+    Expects:
+        - image file in multipart/form-data
+        - Optional: top_k parameter for number of predictions
+
+    Returns:
+        JSON with prediction results
+    """
+    # Check if image file is present
+    if 'image' not in request.files:
+        return jsonify({'error': 'No image file provided'}), 400
+
+    file = request.files['image']
+
+    # Check if file is selected
+    if file.filename == '':
+        return jsonify({'error': 'No file selected'}), 400
+
+    # Check file extension
+    if not allowed_file(file.filename):
+        return jsonify({
+            'error': f'File type not allowed. Allowed types: {", ".join(ALLOWED_EXTENSIONS)}'
+        }), 400
+
+    try:
+        # Get top_k parameter (default: 5)
+        top_k = int(request.form.get('top_k', 5))
+
+        # Load and preprocess image
+        image = Image.open(file.stream)
+        pixel_values, original_image = preprocess_image(image)
+
+        # Make prediction
+        results = predict_image(pixel_values, top_k=top_k)
+
+        return jsonify({
+            'success': True,
+            'prediction': results
+        })
+
+    except Exception as e:
+        return jsonify({
+            'success': False,
+            'error': str(e)
+        }), 500
+
+
+@app.route('/predict_with_explainability', methods=['POST'])
+def predict_with_explainability():
+    """
+    Predict cervical cell classification with GRAD-CAM visualization.
+
+    Expects:
+        - image file in multipart/form-data
+        - Optional: top_k parameter for number of predictions
+        - Optional: target_class parameter for GRAD-CAM visualization
+
+    Returns:
+        JSON with prediction results and GRAD-CAM visualization
+    """
+    # Check if image file is present
+    if 'image' not in request.files:
+        return jsonify({'error': 'No image file provided'}), 400
+
+    file = request.files['image']
+
+    # Check if file is selected
+    if file.filename == '':
+        return jsonify({'error': 'No file selected'}), 400
+
+    # Check file extension
+    if not allowed_file(file.filename):
+        return jsonify({
+            'error': f'File type not allowed. Allowed types: {", ".join(ALLOWED_EXTENSIONS)}'
+        }), 400
+
+    try:
+        # Get parameters
+        top_k = int(request.form.get('top_k', 5))
+        target_class = request.form.get('target_class')
+        if target_class is not None:
+            target_class = int(target_class)
+
+        # Load and preprocess image
+        image = Image.open(file.stream)
+        pixel_values, original_image = preprocess_image(image)
+
+        # Make prediction
+        prediction_results = predict_image(pixel_values, top_k=top_k)
+
+        # Apply GRAD-CAM
+        gradcam_results = apply_gradcam(pixel_values, original_image, target_class)
+
+        # Encode visualization as base64
+        visualization_base64 = encode_image_to_base64(gradcam_results['visualization'])
+        original_image_base64 = encode_image_to_base64(original_image)
+
+        return jsonify({
+            'success': True,
+            'prediction': prediction_results,
+            'explainability': {
+                'method': 'GRAD-CAM',
+                'target_class': DISPLAY_NAMES[gradcam_results['target_class']],
+                'target_class_idx': gradcam_results['target_class'],
+                'visualization': visualization_base64,
+                'original_image': original_image_base64
+            }
+        })
+
+    except Exception as e:
+        return jsonify({
+            'success': False,
+            'error': str(e)
+        }), 500
+
+
+@app.route('/classes', methods=['GET'])
+def get_classes():
+    """Get list of available classes."""
+    return jsonify({
+        'classes': [
+            {
+                'idx': i,
+                'name': display_name,
+                'raw_name': raw_name
+            }
+            for i, (display_name, raw_name) in enumerate(zip(DISPLAY_NAMES, CLASS_NAMES))
+        ]
+    })
+
+
+@app.route('/', methods=['GET'])
+def index():
+    """Root endpoint with API information."""
+    return jsonify({
+        'name': 'Project Phoenix - Cervical Cancer Cell Classification API',
+        'version': '1.0.0',
+        'model': HF_MODEL_ID,
+        'device': str(DEVICE),
+        'endpoints': {
+            '/health': 'GET - Health check',
+            '/predict': 'POST - Predict cell classification',
+            '/predict_with_explainability': 'POST - Predict with GRAD-CAM visualization',
+            '/classes': 'GET - Get available classes'
+        },
+        'supported_formats': list(ALLOWED_EXTENSIONS),
+        'max_file_size': f'{MAX_FILE_SIZE / (1024 * 1024)}MB'
+    })
+
+
+# ========== MAIN ==========
+
+if __name__ == '__main__':
+    # Get port from environment variable or default to 5000
+    port = int(os.getenv('PORT', 5000))
+
+    # Run the app
+    app.run(
+        host='0.0.0.0',
+        port=port,
+        debug=os.getenv('FLASK_ENV') == 'development'
+    )

requirements.txt

@@ -0,0 +1,6 @@
+torch
+safetensors
+scikit-learn
+transformers
+numpy
+pillow