main.py

from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi.responses import FileResponse
from fastapi import FastAPI, UploadFile, File, HTTPException
from gradio_client import Client
import os
import tempfile
import shutil
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi.responses import FileResponse

app = FastAPI()

# Initialize Gradio Client
hf_token = os.environ.get('HF_TOKEN')
client = Client("Ashrafb/moondream1", hf_token=hf_token)
# Function to generate captions for uploaded images
def get_caption(image, additional_context):
    if additional_context.strip():  # Check if additional_context is not empty
        context = additional_context
    else:
        context = "What is this image? Describe this image to someone who is visually impaired."
    result = client.predict(image, context, api_name="/predict")
    return result  

@app.post("/uploadfile/")
async def upload_file(image: UploadFile = File(...), additional_context: str = Form(...)):
    try:
        # Create a temporary directory to store the uploaded image
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_image_path = os.path.join(temp_dir, image.filename)
            # Write the uploaded image data to a temporary file
            with open(temp_image_path, "wb") as temp_image:
                shutil.copyfileobj(image.file, temp_image)
            # Read the image data
            with open(temp_image_path, "rb") as image_file:
                image_data = image_file.read()
            # Generate caption using Gradio client
            result = get_caption(image_data, additional_context)
            return {"result": result}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


app.mount("/", StaticFiles(directory="static", html=True), name="static")

@app.get("/")
def index() -> FileResponse:
    return FileResponse(path="/app/static/index.html", media_type="text/html")