from flask import Flask, request, jsonify
import torch
from torchvision import transforms
from PIL import Image
import io
import torch.nn as nn
import timm

# Define the model class
class LeukemiaModel(nn.Module):
    def __init__(self, num_classes=4):
        super(LeukemiaModel, self).__init__()
        
        # Pretrained EfficientNetB7
        self.base_model = timm.create_model('efficientnet_b7', pretrained=False, num_classes=0)
        
        # Custom layers
        self.flatten = nn.Flatten()
        self.dense1 = nn.Linear(2560 * 7 * 7, 256)  # Adjust input size if needed
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(p=0.3)
        self.dense2 = nn.Linear(256, num_classes)  # Final output layer for classification
        self.softmax = nn.Softmax(dim=1)  # Softmax for multi-class classification
        
    def forward(self, x):
        # Forward pass through the base model (EfficientNetB7)
        x = self.base_model(x)
        
        # Flatten and pass through custom layers
        x = self.flatten(x)
        x = self.dense1(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.dense2(x)
        x = self.softmax(x)
        
        return x


# Initialize Flask app
app = Flask(__name__)

# Load the pre-trained model
model = LeukemiaModel(num_classes=4)
model.load_state_dict(torch.load('leukemia_model.pth', weights_only=True))
model.eval()  # Set the model to evaluation mode

# Define the image transformation pipeline
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # Resize image to match EfficientNet input size
    transforms.ToTensor(),  # Convert image to tensor
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # ImageNet normalization
])

@app.route('/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return jsonify({'error': 'No file part'}), 400
    
    file = request.files['file']
    
    # Read and process the image
    try:
        img_bytes = file.read()
        img = Image.open(io.BytesIO(img_bytes))
        img = transform(img).unsqueeze(0)  # Add batch dimension
    except Exception as e:
        return jsonify({'error': str(e)}), 400
    
    # Inference
    with torch.no_grad():
        outputs = model(img)
        _, predicted = torch.max(outputs, 1)  # Get the class with the highest probability
    
    return jsonify({'prediction': int(predicted.item())})

if __name__ == '__main__':
    app.run(debug=True)