scoring_Yolo_Model.py

#!/usr/bin/env python3
"""
FIXED ONNX Scoring Script - Based on diagnostic findings
"""

import json
import numpy as np
import onnxruntime as ort
from PIL import Image
import io
import base64
import os
import logging
from typing import List, Dict, Tuple, Any

# Configure logging for Azure ML
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Global variables for model and classes
model = None

# EXACT class names from working script - DO NOT CHANGE
class_names = ['Wall', 'Detail', 'Wall2']

# Model configuration - Updated for YOLO26
INPUT_SIZE = (640, 640)
CONFIDENCE_THRESHOLD = 0.7  # 70% minimum confidence for YOLO26


def init():
    """
    Initializes the ONNX model when the Azure ML container starts.
    """
    global model
    
    try:
        # Model path - read from root directory of repository
        model_path = 'best.onnx'
        
        # Check if model file exists
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"ONNX model not found at {model_path}")
        
        # Configure ONNX Runtime providers - GPU first, CPU fallback
        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
        
        # Check available providers
        available_providers = ort.get_available_providers()
        logger.info(f"🔍 Available ONNX providers: {available_providers}")
        
        # Use GPU if available, otherwise fallback to CPU
        if 'CUDAExecutionProvider' in available_providers:
            logger.info("🚀 Using GPU (CUDA) acceleration")
        else:
            logger.info("⚠️ GPU not available, using CPU")
            providers = ['CPUExecutionProvider']
        
        # Initialize ONNX Runtime session
        session_options = ort.SessionOptions()
        session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        
        model = ort.InferenceSession(
            model_path, 
            sess_options=session_options,
            providers=providers
        )
        
        # Log model information
        input_details = model.get_inputs()[0]
        output_details = model.get_outputs()
        
        logger.info(f"✅ ONNX model loaded successfully from {model_path}")
        logger.info(f"📊 Model input: {input_details.name} {input_details.shape}")
        logger.info(f"📊 Model outputs: {[output.name for output in output_details]}")
        logger.info(f"🏷️ Classes: {class_names}")
        
    except Exception as e:
        logger.error(f"❌ Error loading ONNX model: {e}")
        raise e

def preprocess_image_for_onnx(image_bytes: str, input_size=(640, 640)):
    """
    Preprocesses image - FIXED normalization based on diagnostic
    """
    try:
        # Decode base64 image 
        image_data = base64.b64decode(image_bytes)
        pil_image = Image.open(io.BytesIO(image_data)).convert("RGB")
        original_size = pil_image.size  # (width, height)
        resized_image = pil_image.resize(input_size)
        
        # FIXED: Use [0,1] normalization (confirmed working in diagnostic)
        img_np = np.array(resized_image).astype(np.float32) / 255.0
        
        img_np = img_np.transpose(2, 0, 1)
        img_np = np.expand_dims(img_np, axis=0)
        
        logger.info(f"🔍 Preprocessed image shape: {img_np.shape}, range: [{np.min(img_np):.3f}, {np.max(img_np):.3f}]")
        
        return img_np, original_size, input_size
        
    except Exception as e:
        logger.error(f"❌ Error preprocessing image: {e}")
        raise

def postprocess_onnx_output(predictions, original_size, input_size, conf_threshold=None):
    """
    FIXED post-processing based on diagnostic findings
    Output shape: (1, 7, 8400) -> (8400, 7)
    """
    try:
        if conf_threshold is None:
            conf_threshold = CONFIDENCE_THRESHOLD
        
        # Remove batch dimension: (1, 7, 8400) -> (7, 8400)
        predictions = predictions[0]
        logger.info(f"🔍 After removing batch dim: {predictions.shape}")
        
        # FIXED: Transpose to get (8400, 7) format
        predictions = predictions.transpose(1, 0)  # (7, 8400) -> (8400, 7)
        logger.info(f"🔍 After transpose: {predictions.shape}")
        
        boxes = predictions[:, :4]  # Center_x, Center_y, Width, Height
        scores = predictions[:, 4:] # Class scores (one score per class)
        
        logger.info(f"🔍 Boxes shape: {boxes.shape}, range: [{np.min(boxes):.3f}, {np.max(boxes):.3f}]")
        logger.info(f"🔍 Scores shape: {scores.shape}, range: [{np.min(scores):.6f}, {np.max(scores):.6f}]")
        
        # Convert boxes from cx, cy, w, h to x1, y1, x2, y2
        boxes_xyxy = np.zeros_like(boxes)
        boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2  # x1
        boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2  # y1
        boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2  # x2
        boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2  # y2
        
        detections = []
        confidence_stats = []
        
        for i in range(len(scores)):
            class_scores = scores[i]
            max_score_idx = np.argmax(class_scores)
            confidence = class_scores[max_score_idx]
            confidence_stats.append(confidence)
            
            if confidence > conf_threshold:
                x1, y1, x2, y2 = boxes_xyxy[i]
                
                # Scale bounding box coordinates back to original image size
                scale_x = original_size[0] / input_size[0]
                scale_y = original_size[1] / input_size[1]
                
                x1_orig = x1 * scale_x
                y1_orig = y1 * scale_y
                x2_orig = x2 * scale_x
                y2_orig = y2 * scale_y
                
                # Ensure valid bounding box
                if x2_orig > x1_orig and y2_orig > y1_orig:
                    detections.append({
                        'class_name': class_names[int(max_score_idx)],
                        'confidence': float(confidence),
                        'bbox_orig': [x1_orig, y1_orig, x2_orig, y2_orig],
                        'class_id': int(max_score_idx)
                    })
        
        # Log statistics
        confidence_stats = np.array(confidence_stats)
        logger.info(f"🔍 Raw predictions: {len(scores)} detections")
        logger.info(f"🔍 Confidence threshold: {conf_threshold}")
        logger.info(f"🔍 Confidence stats - Min: {np.min(confidence_stats):.6f}, Max: {np.max(confidence_stats):.6f}, Mean: {np.mean(confidence_stats):.6f}")
        logger.info(f"🔍 Detections above threshold ({conf_threshold}): {len(detections)}")
        
        if len(detections) > 0:
            logger.info("🔍 Found detections:")
            for i, det in enumerate(detections[:5]):  # Show top 5
                logger.info(f"   {i+1}. {det['class_name']} (conf: {det['confidence']:.6f})")
        else:
            # Show top confidences for debugging
            top_indices = np.argsort(confidence_stats)[-10:][::-1]
            logger.info(f"🔍 Top 10 confidences found:")
            for idx in top_indices:
                conf = confidence_stats[idx]
                class_idx = np.argmax(scores[idx])
                logger.info(f"   Detection {idx}: {class_names[class_idx]} = {conf:.6f} {'✅' if conf > conf_threshold else '❌'}")
        
        return detections
        
    except Exception as e:
        logger.error(f"❌ Error in postprocessing: {e}")
        return []

def run(raw_data: str) -> Dict[str, Any]:
    """
    Main inference function - FIXED based on diagnostic
    """
    try:
        # Parse input data
        data = json.loads(raw_data)
        
        # Handle batch format (list of tiles with offsets)
        if 'tiles' in data:
            tiles_data = data['tiles']
        else:
            # Handle single image format (fallback)
            tiles_data = [{
                'image': data['image'],
                'x_offset': 0,
                'y_offset': 0
            }]
        
        logger.info(f"Processing {len(tiles_data)} tiles with Azure ML YOLO endpoint")
        logger.info(f"📊 Using confidence threshold: {CONFIDENCE_THRESHOLD}")
        
        all_detections = []
        
        for tile_data in tiles_data:
            image_base64 = tile_data['image']
            x_offset = tile_data['x_offset']
            y_offset = tile_data['y_offset']
            
            try:
                # Preprocess image
                input_image, original_size, model_input_size = preprocess_image_for_onnx(image_base64)
                
                # Get input and output names from the ONNX model session
                input_name = model.get_inputs()[0].name
                output_names = [output.name for output in model.get_outputs()]
                
                # Run inference using the ONNX Runtime session
                outputs = model.run(output_names, {input_name: input_image})
                
                logger.info(f"🔍 Model output shape: {outputs[0].shape}")
                
                # Post-process the ONNX output
                tile_detections = postprocess_onnx_output(
                    outputs[0], 
                    original_size, 
                    model_input_size, 
                    conf_threshold=CONFIDENCE_THRESHOLD
                )
                
                # Add tile offset to detection coordinates
                for det in tile_detections:
                    x1_tile, y1_tile, x2_tile, y2_tile = det['bbox_orig']
                    
                    x1_orig = x1_tile + x_offset
                    y1_orig = y1_tile + y_offset
                    x2_orig = x2_tile + x_offset
                    y2_orig = y2_tile + y_offset
                    
                    all_detections.append({
                        'class_name': det['class_name'],
                        'confidence': det['confidence'],
                        'bbox_orig': [x1_orig, y1_orig, x2_orig, y2_orig]
                    })
                
                logger.info(f"Processed tile at ({x_offset},{y_offset}): {len(tile_detections)} detections")
                        
            except Exception as e:
                logger.error(f"Error during YOLO inference for tile ({x_offset},{y_offset}): {e}")
        
        # Return in format function_app.py expects
        response = {
            'detections': all_detections,
            'num_detections': len(all_detections),
            'processed_tiles': len(tiles_data),
            'status': 'success'
        }
        
        logger.info(f"✅ Inference complete: {len(all_detections)} detections from {len(tiles_data)} tiles")
        return response
        
    except Exception as e:
        error_msg = str(e)
        logger.error(f"❌ Error during inference: {error_msg}")
        
        return {
            'detections': [],
            'num_detections': 0,
            'error': error_msg,
            'status': 'error'
        }

# Entry point for local testing
if __name__ == "__main__":
    print("Initializing model...")
    init()
    print("Model initialized successfully!")