File size: 4,905 Bytes
f733284 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 |
# ==============================================================================
# Aura Mind Glow - Main Server (FastAPI + Gradio)
# ==============================================================================
"""
This script is the main entry point for the application. It launches a FastAPI
server that provides the diagnosis API and also serves the entire Gradio UI.
To run this server for development:
1. Make sure you have installed all packages from requirements.txt.
2. Run the command: uvicorn api_server:app --host 127.0.0.1 --port 7860
When deployed to Hugging Face Spaces, the Procfile will handle this command.
"""
# --- Essential Imports ---
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
from PIL import Image
import os
import warnings
import tempfile
import re
import io
import gradio as gr
# --- Import Core Components from Modules ---
# This setup is now shared between the API and the Gradio App
from vision_model import load_vision_model
from knowledge_base import KnowledgeBase
from agent_setup import initialize_adk
from bigquery_search import search_bigquery_for_remedy
from vector_store import embed_and_store_documents
# --- Import the Gradio UI from app.py ---
# We import the 'demo' object directly. The app.py script should not call demo.launch()
try:
from app import demo as gradio_app
print("β
Gradio UI imported successfully from app.py.")
except ImportError as e:
gradio_app = None
print(f"β CRITICAL: Could not import Gradio UI from app.py: {e}")
print("Ensure app.py defines a Gradio Blocks object named 'demo' and does not call .launch().")
print("β
All server libraries imported successfully.")
# --- Global Initialization ---
warnings.filterwarnings("ignore")
os.environ["TORCH_COMPILE_DISABLE"] = "1"
print("Performing initial setup for server (this may take a moment)...")
VISION_MODEL, PROCESSOR = load_vision_model()
KB = KnowledgeBase()
RETRIEVER = KB
embed_and_store_documents()
adk_components = initialize_adk(VISION_MODEL, PROCESSOR, RETRIEVER)
DIAGNOSIS_TOOL = adk_components["diagnosis_tool"] if adk_components else None
if not DIAGNOSIS_TOOL:
print("β CRITICAL: Diagnosis tool could not be initialized. The API will not work.")
print("β
Server setup complete.")
# --- FastAPI App and Endpoint Logic ---
app = FastAPI(
title="Aura Mind Glow API",
description="Provides access to the plant diagnosis model and serves the Gradio UI.",
version="1.0.0",
)
def run_diagnosis_logic(image: Image.Image):
"""
Core logic for running diagnosis and getting remedies.
"""
temp_file_path = None
try:
with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file:
image.save(temp_file.name)
temp_file_path = temp_file.name
diagnosis = DIAGNOSIS_TOOL(temp_file_path)
if "Could not parse" in diagnosis:
return {"error": f"Could not identify condition: {diagnosis}"}
cleaned_diagnosis = re.sub(r'[^\w\s.\\-,\"]', '', diagnosis)
cleaned_diagnosis = re.sub(r'\s+', ' ', cleaned_diagnosis).strip()
local_remedy_list = search_documents(cleaned_diagnosis)
local_remedy = local_remedy_list[0] if local_remedy_list else "No remedy found in local knowledge base."
search_query = "healthy maize" if "healthy" in cleaned_diagnosis.lower() else "phosphorus" if "phosphorus" in cleaned_diagnosis.lower() else "general"
cloud_remedy = search_bigquery_for_remedy(search_query)
return {
"diagnosis": diagnosis,
"remedy_local": local_remedy,
"remedy_cloud": cloud_remedy
}
finally:
if temp_file_path:
os.remove(temp_file_path)
@app.post("/diagnose/", tags=["Diagnosis"])
async def diagnose_endpoint(file: UploadFile = File(...)):
"""
Receives an image file, performs diagnosis, and returns the result as JSON.
"""
if not file.content_type.startswith('image/'):
raise HTTPException(status_code=400, detail="File provided is not an image.")
try:
image_bytes = await file.read()
image = Image.open(io.BytesIO(image_bytes))
result = run_diagnosis_logic(image)
if "error" in result:
raise HTTPException(status_code=500, detail=result["error"])
return JSONResponse(content=result)
except Exception as e:
print(f"β API Error: {e}")
raise HTTPException(status_code=500, detail=f"An internal server error occurred: {e}")
# --- Mount the Gradio App ---
if gradio_app:
app = gr.mount_gradio_app(app, gradio_app, path="/")
print("β
Gradio UI has been mounted on the FastAPI server at the root path '/'.")
# Note: The 'if __name__ == "__main__":' block with uvicorn.run() is removed.
# The Procfile will be used by Hugging Face to run the server. |