Spaces:
Sleeping
Sleeping
File size: 7,391 Bytes
89ee5b3 65c0c50 89ee5b3 65c0c50 89ee5b3 65c0c50 89ee5b3 65c0c50 89ee5b3 65c0c50 89ee5b3 65c0c50 89ee5b3 65c0c50 89ee5b3 65c0c50 89ee5b3 65c0c50 89ee5b3 65c0c50 89ee5b3 65c0c50 89ee5b3 |
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 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 |
import os
import base64
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.middleware.cors import CORSMiddleware
import shutil
import uuid
import logging
from typing import Dict, List, Any
import json
# Import scene graph service
from app.scene_graph_service import process_image
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Use /tmp directory which should be writable
UPLOAD_DIR = "/tmp/uploads"
OUTPUT_DIR = "/tmp/outputs"
# Create necessary directories
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)
os.makedirs("app/models", exist_ok=True)
# Initialize FastAPI app
app = FastAPI(title="Scene Graph Generation API")
# Add CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
@app.get("/")
def read_root():
return {
"message": "Scene Graph Generation API is running",
"usage": "POST /generate with an image file to generate a scene graph",
"docs": "Visit /docs for API documentation"
}
@app.post("/generate")
async def generate_scene_graph(
image: UploadFile = File(...),
confidence_threshold: float = Form(0.5),
use_fixed_boxes: bool = Form(False),
) -> Dict[str, Any]:
try:
# Debug information
logger.info(f"Received file: {image.filename}, content_type: {image.content_type}")
# Input validation with improved error handling
if image is None:
raise HTTPException(status_code=400, detail="No image file provided")
if not image.content_type:
# Set a default content type if none provided
logger.warning("No content type provided, assuming image/jpeg")
image.content_type = "image/jpeg"
if not image.content_type.startswith("image/"):
raise HTTPException(
status_code=400, detail=f"Uploaded file must be an image, got {image.content_type}"
)
if not (0 <= confidence_threshold <= 1):
raise HTTPException(
status_code=400, detail="Confidence threshold must be between 0 and 1"
)
# Generate unique ID for this job
job_id = str(uuid.uuid4())
short_id = job_id.split("-")[0]
# Create directories for this job in /tmp which should be writable
upload_job_dir = os.path.join(UPLOAD_DIR, job_id)
output_job_dir = os.path.join(OUTPUT_DIR, job_id)
# Create directories with explicit permission setting
os.makedirs(upload_job_dir, exist_ok=True, mode=0o777)
os.makedirs(output_job_dir, exist_ok=True, mode=0o777)
logger.info(f"Created upload directory: {upload_job_dir}")
logger.info(f"Created output directory: {output_job_dir}")
# Determine file extension
file_ext = os.path.splitext(image.filename)[1] if image.filename else ".jpg"
if not file_ext:
file_ext = ".jpg"
# Save the uploaded image to /tmp
image_filename = f"{short_id}{file_ext}"
image_path = os.path.join(upload_job_dir, image_filename)
# Save the file with error handling
try:
# Explicitly open with write permissions
with open(image_path, "wb") as buffer:
contents = await image.read()
buffer.write(contents)
# Check if file was created and has size
if not os.path.exists(image_path):
raise HTTPException(status_code=400, detail=f"Failed to save uploaded file to {image_path}")
if os.path.getsize(image_path) == 0:
raise HTTPException(status_code=400, detail="Uploaded file is empty")
logger.info(f"Image saved to {image_path} ({os.path.getsize(image_path)} bytes)")
except Exception as e:
logger.error(f"Error saving file: {str(e)}")
raise HTTPException(status_code=500, detail=f"Error saving uploaded file: {str(e)}")
# Define model paths
model_path = "app/models/model.pth"
vocabulary_path = "app/models/vocabulary.json"
# Check if model files exist
if not os.path.exists(model_path):
logger.error(f"Model file not found: {model_path}")
raise HTTPException(status_code=500, detail=f"Model file not found: {model_path}")
if not os.path.exists(vocabulary_path):
logger.error(f"Vocabulary file not found: {vocabulary_path}")
raise HTTPException(status_code=500, detail=f"Vocabulary file not found: {vocabulary_path}")
# Process the image
objects, relationships, annotated_image_path, graph_path = process_image(
image_path=image_path,
model_path=model_path,
vocabulary_path=vocabulary_path,
confidence_threshold=confidence_threshold,
use_fixed_boxes=use_fixed_boxes,
output_dir=output_job_dir,
base_filename=short_id,
)
logger.info(f"Processing complete. Annotated image: {annotated_image_path}, Graph: {graph_path}")
# Verify output files exist
if not os.path.exists(annotated_image_path):
logger.error(f"Annotated image not generated: {annotated_image_path}")
raise HTTPException(status_code=500, detail="Failed to generate annotated image")
if not os.path.exists(graph_path):
logger.error(f"Graph image not generated: {graph_path}")
raise HTTPException(status_code=500, detail="Failed to generate graph image")
# Read the generated images as base64
try:
with open(annotated_image_path, "rb") as img_file:
annotated_image_base64 = base64.b64encode(img_file.read()).decode("utf-8")
with open(graph_path, "rb") as img_file:
graph_image_base64 = base64.b64encode(img_file.read()).decode("utf-8")
logger.info("Successfully encoded images as base64")
except Exception as e:
logger.error(f"Error reading output images: {str(e)}")
raise HTTPException(status_code=500, detail=f"Error reading output images: {str(e)}")
# Prepare response with base64 encoded images
response = {
"objects": objects,
"relationships": relationships,
"annotated_image": annotated_image_base64,
"graph_image": graph_image_base64
}
# Clean up
try:
shutil.rmtree(upload_job_dir)
shutil.rmtree(output_job_dir)
logger.info("Cleaned up temporary directories")
except Exception as e:
logger.warning(f"Error cleaning up temporary files: {str(e)}")
return response
except Exception as e:
logger.error(f"Error processing image: {str(e)}", exc_info=True)
raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
@app.get("/health")
def health_check():
return {"status": "healthy"}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860) |