from fastapi import FastAPI, UploadFile, File
import tensorflow as tf
import numpy as np
from PIL import Image
import io
import easyocr

app = FastAPI()

# Load the model
model = tf.keras.models.load_model("prescription_model.h5")
ocr_reader = easyocr.Reader(['en'])  # English OCR model

@app.post("/predict/")
async def predict(file: UploadFile = File(...)):
    image = Image.open(io.BytesIO(await file.read())).convert('L').resize((128, 128))  # Convert to grayscale
    img_array = np.array(image) / 255.0
    img_array = np.expand_dims(img_array, axis=(0, -1))
    
    # Get predictions
    prediction = model.predict(img_array)

    # Use OCR for text recognition
    text_prediction = ocr_reader.readtext(np.array(image))
    text_result = " ".join([text[1] for text in text_prediction])

    return {"prediction": prediction.tolist(), "ocr_text": text_result}