import os, glob
from flask import Flask, request, render_template
from transformers import pipeline
from ultralytics import YOLO

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'static/uploads'
app.config['MARKED_FOLDER'] = 'static/marked'
PREDICT_FOLDER = os.path.join(app.config['MARKED_FOLDER'], 'predict')

# สร้างโฟลเดอร์หากยังไม่มี
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
os.makedirs(PREDICT_FOLDER, exist_ok=True)

# โหลดโมเดล
pipe = pipeline("image-classification", model="JorgeGIT/Finetuned-Leukemia-cell")
yolo_model = YOLO("runs/detect/leukemia-yolo-model/weights/best.pt")

# 🎯 ฟังก์ชันรวมความเสี่ยง
def calculate_combined_risk(blast_risk, huggingface_score, age, sex):
    score = 0.6 * blast_risk + 0.4 * huggingface_score
    if age and int(age) > 60:
        score += 5
    if sex == 'male':
        score += 2
    return min(round(score, 1), 100)

@app.route('/', methods=['GET'])
def home():
    return render_template('index.html')

@app.route('/upload', methods=['GET', 'POST'])
def upload_file():
    if request.method == 'POST':
        file = request.files.get('file')
        age = request.form.get('age')
        sex = request.form.get('sex')

        if not file or file.filename == '' or not age or not sex:
            return "กรุณากรอกข้อมูลให้ครบ"

        filename = file.filename
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)

        results = yolo_model.predict(
            source=filepath,
            save=True,
            project=app.config['MARKED_FOLDER'],
            name='predict',
            exist_ok=True
        )

        list_of_files = glob.glob(os.path.join(PREDICT_FOLDER, '*.jpg'))
        latest_file = max(list_of_files, key=os.path.getctime)
        marked_filename = os.path.relpath(latest_file, 'static').replace('\\', '/')

        total_cells = len(results[0].boxes)
        abnormal_cells = 0
        for box in results[0].boxes:
            cls_id = int(box.cls[0].item())
            class_name = results[0].names[cls_id]
            if 'blast' in class_name.lower():
                abnormal_cells += 1
        blast_risk = round((abnormal_cells / total_cells) * 100, 1) if total_cells else 0.0

        result = pipe(filepath)
        lla_result = [
            {'label': 'โอกาสการเกิดโรค Leukemia = ', 'score': round(res['score'] * 1000, 1)}
            for res in result if res['label'].lower() == 'lla'
        ]
        huggingface_score = lla_result[0]['score'] if lla_result else 0.0

        combined_risk = calculate_combined_risk(blast_risk, huggingface_score, age, sex)

        return render_template('result.html',
            filename=filename,
            marked_filename=marked_filename,
            total_cells=total_cells,
            abnormal_cells=abnormal_cells,
            blast_risk=blast_risk,
            huggingface_score=huggingface_score,
            combined_risk=combined_risk,
            age=age,
            sex=sex
        )

    return render_template('upload.html')
if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)