Shridhartd's picture
Update app.py
22c395c verified
import requests
from fastapi import FastAPI, UploadFile, File
import shutil
import os
# Initialize FastAPI
app = FastAPI()
# Create an upload directory
UPLOAD_DIR = "/app/uploads"
if not os.path.exists(UPLOAD_DIR):
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.chmod(UPLOAD_DIR, 0o777)
# Hugging Face API Key (Replace with your actual API key)
HUGGINGFACE_API_KEY = "HF_API_KEY"
# Hugging Face Model URL
MODEL_URL = "https://api-inference.huggingface.co/models/Salesforce/blip-image-captioning-base"
def get_caption(image_path):
"""Sends image to Hugging Face API for captioning"""
headers = {
"Authorization": f"Bearer {HUGGINGFACE_API_KEY}",
"Content-Type": "application/octet-stream"
}
with open(image_path, "rb") as image_file:
image_data = image_file.read()
response = requests.post(MODEL_URL, headers=headers, data=image_data)
if response.status_code == 200:
try:
return response.json()[0]['generated_text']
except (IndexError, KeyError):
return f"Error: Unexpected response format: {response.text}"
else:
return f"Error: {response.text}"
@app.post("/upload-image/")
async def upload_image(file: UploadFile = File(...)):
"""Uploads image & gets caption"""
file_path = f"{UPLOAD_DIR}/{file.filename}"
# Save the uploaded file
with open(file_path, "wb") as buffer:
shutil.copyfileobj(file.file, buffer)
# Get caption
caption = get_caption(file_path)
return {"filename": file.filename, "caption": caption}
@app.get("/")
async def root():
return {"message": "Image-to-Text API using Hugging Face Inference API is running."}