src/api/api.operator.py

from flask import Flask, request, jsonify
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
import io
import base64
import numpy as np
import cv2

app = Flask(__name__)

# [Previous model architecture code remains the same...]
class AttentionModule(nn.Module):
    def __init__(self, channels: int):
        super().__init__()
        self.conv1 = nn.Conv2d(channels, channels // 8, kernel_size=1)
        self.conv2 = nn.Conv2d(channels // 8, channels, kernel_size=1)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        attention = F.avg_pool2d(x, 1)
        attention = F.relu(self.conv1(attention))
        attention = torch.sigmoid(self.conv2(attention))
        return x * attention

class ResidualBlock(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, stride: int = 1):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, 
                              stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3,
                              stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)
        
        self.shortcut = nn.Sequential()
        if stride != 1 or in_channels != out_channels:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=1,
                         stride=stride, bias=False),
                nn.BatchNorm2d(out_channels)
            )
        
        self.attention = AttentionModule(out_channels)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out = self.attention(out)
        out += self.shortcut(x)
        out = F.relu(out)
        return out

class OptimizedOCR(nn.Module):
    def __init__(self, num_classes: int = 36):
        super().__init__()
        self.in_channels = 64
        
        self.conv1 = nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        
        self.layer1 = self.make_layer(64, stride=1, num_blocks=2)
        self.layer2 = self.make_layer(128, stride=2, num_blocks=2)
        self.layer3 = self.make_layer(256, stride=2, num_blocks=2)
        
        self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Sequential(
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, num_classes)
        )

    def make_layer(self, out_channels: int, stride: int, num_blocks: int) -> nn.Sequential:
        strides = [stride] + [1]*(num_blocks-1)
        layers = []
        for stride in strides:
            layers.append(ResidualBlock(self.in_channels, out_channels, stride))
            self.in_channels = out_channels
        return nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.avg_pool(out)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        return out

# Global variables
device = torch.device("mps")  # Using MPS for M2 Ultra
model = None
idx_to_label = None

def create_label_map():
    """Creates and returns the label mapping dictionary"""
    label_map = {}
    # Add digits (0-9)
    for i in range(10):
        label_map[i] = str(i)
    # Add uppercase letters (A-Z)
    for i in range(26):
        label_map[i + 10] = chr(65 + i)
    return label_map

def initialize_model():
    """Initializes the model with pre-trained weights"""
    global model, idx_to_label
    
    # Create label mapping
    idx_to_label = create_label_map()
    
    # Initialize model
    model = OptimizedOCR(num_classes=36).to(device)
    model.load_state_dict(torch.load('best_model.pth', map_location=device))
    model.eval()
    
    print("Model initialized successfully")

def preprocess_image(image_data: bytes) -> torch.Tensor:
    """Enhanced preprocessing function to convert any input image to binary format"""
    # Convert bytes to numpy array
    nparr = np.frombuffer(image_data, np.uint8)
    image = cv2.imdecode(nparr, cv2.IMREAD_GRAYSCALE)
    
    # Ensure image is loaded
    if image is None:
        raise ValueError("Failed to load image")
    
    # Resize to 27x27
    image = cv2.resize(image, (27, 27), interpolation=cv2.INTER_AREA)
    
    # Apply adaptive thresholding to handle varying lighting conditions
    binary = cv2.adaptiveThreshold(
        image,
        255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY,
        11,  # Block size
        2    # C constant
    )
    
    # Determine if we need to invert the image
    # Count white pixels
    white_pixels = np.sum(binary == 255)
    total_pixels = binary.size
    
    # If majority is white, invert the image
    if white_pixels > total_pixels / 2:
        binary = cv2.bitwise_not(binary)
    
    # Convert to tensor and normalize
    image_tensor = torch.FloatTensor(binary).unsqueeze(0).unsqueeze(0) / 255.0
    
    return image_tensor

@app.route('/ping', methods=['GET'])
def ping():
    """Health check endpoint"""
    return jsonify({"status": "healthy", "message": "OCR service is running"})

@app.route('/predict', methods=['POST'])
def predict():
    """Prediction endpoint"""
    if 'image' not in request.files:
        return jsonify({"error": "No image file provided"}), 400
    
    try:
        image_file = request.files['image']
        image_data = image_file.read()
        
        # Preprocess image
        image_tensor = preprocess_image(image_data)
        image_tensor = image_tensor.to(device)
        
        # Make prediction
        with torch.no_grad():
            output = model(image_tensor)
            probabilities = torch.softmax(output, dim=1)
            predicted_idx = output.argmax(1).item()
            confidence = float(probabilities[0][predicted_idx])
            predicted_char = idx_to_label[predicted_idx]
        
        # For debugging, you can also return the preprocessed image
        # Convert tensor back to base64 encoded image
        preprocessed_img = (image_tensor.squeeze().cpu().numpy() * 255).astype(np.uint8)
        _, buffer = cv2.imencode('.png', preprocessed_img)
        preprocessed_b64 = base64.b64encode(buffer).decode('utf-8')
        
        return jsonify({
            "character": predicted_char,
            "confidence": confidence,
            "preprocessed_image": preprocessed_b64,  # This helps in debugging
            "status": "success"
        })
    
    except Exception as e:
        return jsonify({
            "error": str(e),
            "status": "error"
        }), 500

if __name__ == '__main__':
    initialize_model()
    app.run(host='0.0.0.0', port=5001)