from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
import gradio as gr
import os
from ONNX0630 import main as predict_smiles
from PIL import Image
import io

# Initialize FastAPI app
app = FastAPI(title="Chemical Structure to SMILES API")

# API endpoint to predict SMILES from an image
@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    try:
        # Read and save the uploaded image
        contents = await file.read()
        image = Image.open(io.BytesIO(contents))
        temp_path = f"temp_{file.filename}"
        image.save(temp_path)
        
        # Call the model function
        smiles = predict_smiles(temp_path)
        
        # Clean up temporary file
        os.remove(temp_path)
        
        return JSONResponse(content={"smiles": smiles})
    except Exception as e:
        return JSONResponse(content={"error": str(e)}, status_code=500)

# Gradio interface
def gradio_predict(image):
    try:
        # Save the uploaded image
        temp_path = "temp_image.png"
        image.save(temp_path)
        
        # Call the model function
        smiles = predict_smiles(temp_path)
        
        # Clean up
        os.remove(temp_path)
        
        return smiles
    except Exception as e:
        return f"Error: {str(e)}"

# Define Gradio interface
iface = gr.Interface(
    fn=gradio_predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Textbox(),
    title="Chemical Structure to SMILES Converter",
    description="Upload an image of a chemical structure to get its SMILES string."
)

# Launch Gradio with FastAPI
app = gr.mount_gradio_app(app, iface, path="/")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)