app.py

import os
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import cv2
import numpy as np
from flask import Flask, request, jsonify, render_template, send_from_directory
from flask_cors import CORS

# Initialize Flask
app = Flask(__name__, static_folder='static', static_url_path='/')
CORS(app) # Enable CORS for all routes

# --- AI Model Initialization (LlamaCpp) ---
try:
    from llama_cpp import Llama
    LLM_PATH = os.path.join(os.path.dirname(__file__), "models", "qwen2.5-1.5b-instruct-q4_k_m.gguf")
    llm = Llama(
        model_path=LLM_PATH,
        n_ctx=2048,
        n_threads=4,
        verbose=False,
        chat_format="chatml"
    )
    print(f">>> Local LLM loaded from {LLM_PATH}")
except Exception as e:
    print(f">>> Failed to load local LLM: {e}")
    llm = None

# --- Configuration ---
MODEL_PATH = os.path.join(os.path.dirname(__file__), "models", "dr_model_complete.pth")
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
CLASSES = ['No Diabetic Retinopathy', 'Mild DR', 'Moderate DR', 'Severe DR', 'Proliferative DR']

# --- Model Definition (Same as Pipeline) ---
class DRViTModel(nn.Module):
    def __init__(self, num_classes=5):
        super().__init__()
        self.vit = models.vit_b_16(weights=None)
        self.vit.heads.head = nn.Linear(self.vit.heads.head.in_features, num_classes)

    def forward(self, x):
        return self.vit(x)

# Loading the model
model = DRViTModel().to(DEVICE)
try:
    model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
    model.eval()
    print(">>> AI Model loaded successfully!")
except Exception as e:
    print(f">>> Error loading model: {e}")

# --- Image Preprocessing ---
def circle_crop_img(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    _, thresh = cv2.threshold(gray, 10, 255, cv2.THRESH_BINARY)
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours: return img
    cnt = max(contours, key=cv2.contourArea)
    x, y, w, h = cv2.boundingRect(cnt)
    return img[y:y+h, x:x+w]

def apply_clahe(img):
    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    cl = clahe.apply(l)
    return cv2.cvtColor(cv2.merge((cl, a, b)), cv2.COLOR_LAB2BGR)

def preprocess_image(image_path):
    img = cv2.imread(image_path)
    img = circle_crop_img(img)
    img = apply_clahe(img)
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    pil_img = Image.fromarray(img_rgb)
    
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    return transform(pil_img).unsqueeze(0).to(DEVICE)

def get_ai_explanation(diagnosis):
    if not llm:
        return "AI Explanation Service Unavailable (Local LLM not loaded)."
        
    system_prompt = (
        "You are VisiCore AI, a specialized medical assistant for Diabetic Retinopathy (DR). "
        "The patient has been diagnosed with: "
        f"***{diagnosis}***. "
        "Explain what this diagnosis means in simple terms, its potential risks, and next steps for the user. "
        "Keep it professional, empathetic, and concise (2-3 sentences). "
        "Always remind the user that this is an AI screening and they MUST consult an ophthalmologist."
    )
    
    try:
        response = llm.create_chat_completion(
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"Explain my diagnosis: {diagnosis}"}
            ]
        )
        return response['choices'][0]['message']['content']
    except Exception as e:
        print(f"Local AI Explanation Error: {e}")
    return None

# --- Routes ---
@app.route('/')
def home():
    return send_from_directory(app.static_folder, 'index.html')

@app.route('/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return jsonify({'error': 'No file uploaded'})
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'No file selected'})

    # Save temp image
    temp_path = "temp_prediction.jpg"
    file.save(temp_path)
    
    try:
        input_tensor = preprocess_image(temp_path)
        with torch.no_grad():
            output = model(input_tensor)
            probabilities = torch.nn.functional.softmax(output[0], dim=0)
            confidence, predicted_idx = torch.max(probabilities, 0)
        
        diagnosis = CLASSES[predicted_idx.item()]
        
        # Get AI explanation
        explanation = get_ai_explanation(diagnosis)
        
        result = {
            'class': diagnosis,
            'confidence': f"{confidence.item()*100:.2f}%",
            'explanation': explanation,
            'status': 'Analysis Complete'
        }
        os.remove(temp_path)
        return jsonify(result)
    except Exception as e:
        if os.path.exists(temp_path):
            os.remove(temp_path)
        return jsonify({'error': str(e)})

@app.route('/chat', methods=['POST'])
def chat():
    data = request.json
    user_message = data.get('message', '')
    history = data.get('history', [])  # Get chat history from frontend
    diagnosis = data.get('diagnosis', 'No diagnosis yet')
    
    if not user_message:
        return jsonify({'error': 'No message provided'})

    if not llm:
        return jsonify({'error': 'Local AI engine is not running.'})

    messages = [
        {"role": "system", "content": f"You are VisiCore AI, a specialized medical assistant for Diabetic Retinopathy (DR). The patient has been diagnosed with: ***{diagnosis}***. Always remind the user that this is an AI screening and they MUST consult an ophthalmologist."}
    ]
    
    # Add history to messages
    for msg in history:
        messages.append({"role": msg['role'], "content": msg['text']})
        
    # Add latest user message
    messages.append({"role": "user", "content": user_message})

    try:
        response = llm.create_chat_completion(messages=messages)
        ai_message = response['choices'][0]['message']['content']
        return jsonify({'reply': ai_message})
    except Exception as e:
        return jsonify({'error': f"Local AI Error: {str(e)}"})

if __name__ == '__main__':

    app.run(debug=True, port=5000)