0-dev-0/Depth_Estimation

1#!/usr/bin/env python2# coding: utf-83 4import asyncio5import uuid6import shutil7import os8import time9 10import torch11import numpy as np12import cv213import matplotlib14matplotlib.use('Agg') # Set backend BEFORE importing pyplot15import matplotlib.pyplot as plt16 17from fastapi import FastAPI, UploadFile, File, Request, HTTPException, Query18from fastapi.responses import FileResponse, HTMLResponse19from fastapi.staticfiles import StaticFiles20 21# --- Configuration ---22# Use the standard /tmp directory which is typically writable in containers23UPLOAD_DIR = "/tmp/temp_uploads"24OUTPUT_DIR = "/tmp/temp_outputs"25# --- END CHANGE ---26 27os.makedirs(UPLOAD_DIR, exist_ok=True)28os.makedirs(OUTPUT_DIR, exist_ok=True)29 30# Choose MiDaS model type:31# Options: "MiDaS_small", "DPT_Hybrid", "DPT_Large" (Highest quality)32MODEL_TYPE = "DPT_Hybrid"33 34# --- Model Loading ---35print("Loading MiDaS model...")36start_load_time = time.time()37try:38    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")39    print(f"Using device: {device}")40 41    midas_model = torch.hub.load("intel-isl/MiDaS", MODEL_TYPE)42    midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")43 44    if MODEL_TYPE == "MiDaS_small":45        transform = midas_transforms.small_transform46    else:47        transform = midas_transforms.dpt_transform48 49    midas_model.to(device)50    midas_model.eval()51    load_time = time.time() - start_load_time52    print(f"MiDaS model '{MODEL_TYPE}' loaded successfully in {load_time:.2f} seconds.")53 54except Exception as e:55    print(f"FATAL ERROR: Could not load MiDaS model '{MODEL_TYPE}' from torch.hub.")56    print(f"Error details: {e}")57    print("Please ensure internet connectivity and correct MODEL_TYPE.")58    import sys59    sys.exit(1)60 61# --- Helper Function: Depth Prediction ---62def predict_depth(image_path: str, output_dir: str, colormap: str = 'plasma') -> str:63    """64    Reads image, predicts depth, saves visualization with decorations, returns output path.65    """66    try:67        print(f"Processing image: {image_path}")68        img_bgr = cv2.imread(image_path)69        if img_bgr is None:70            raise ValueError(f"Could not read image file: {image_path}")71 72        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)73        input_batch = transform(img_rgb).to(device)74 75        with torch.no_grad():76            start_pred_time = time.time()77            prediction = midas_model(input_batch)78            pred_time = time.time() - start_pred_time79            print(f"Inference time: {pred_time:.3f} seconds")80 81            prediction = torch.nn.functional.interpolate(82                prediction.unsqueeze(1),83                size=img_rgb.shape[:2],84                mode="bicubic",85                align_corners=False,86            ).squeeze()87 88        depth_map = prediction.cpu().numpy()89 90        depth_min = depth_map.min()91        depth_max = depth_map.max()92        if depth_max - depth_min > 1e-6:93            normalized_map = (depth_map - depth_min) / (depth_max - depth_min)94        else:95            normalized_map = np.zeros_like(depth_map)96 97        # --- Save Visualization ---98        output_filename = f"depth_{uuid.uuid4().hex[:8]}.png"99        output_path = os.path.join(output_dir, output_filename)100 101        plt.figure(figsize=(10, 6))102        plt.imshow(normalized_map, cmap=colormap)103        plt.title(f"Estimated Depth Map ({MODEL_TYPE})")104        cbar = plt.colorbar()105        cbar.set_label("Relative Inverse Depth (Higher Value = Closer)")106        plt.axis('off')107        plt.tight_layout()108        plt.savefig(output_path)109        plt.close()110 111        print(f"Saved depth map visualization to: {output_path}")112        return output_path113 114    except Exception as e:115        print(f"Error during depth prediction for {image_path}: {e}")116        import traceback117        traceback.print_exc()118        raise119 120# --- FastAPI Application Setup ---121app = FastAPI(title="MiDaS Depth Estimation API")122 123# --- API Endpoints ---124@app.get("/", response_class=HTMLResponse)125async def serve_frontend():126    """Serves the main HTML frontend."""127    html_file_path = "index.html"128    if not os.path.exists(html_file_path):129        return HTMLResponse(content="<html><body><h1>Error</h1><p>Frontend file 'index.html' not found.</p></body></html>", status_code=404)130    try:131        with open(html_file_path, "r", encoding="utf-8") as f:132            return HTMLResponse(content=f.read(), status_code=200)133    except Exception as e:134        print(f"Error reading index.html: {e}")135        raise HTTPException(status_code=500, detail="Internal Server Error reading frontend file.")136 137@app.post("/predict/", response_class=FileResponse)138async def create_prediction(139    file: UploadFile = File(...),140    colormap: str = Query("plasma", enum=["viridis", "plasma", "magma", "inferno", "cividis", "gray"])141):142    """143    Accepts an uploaded image, predicts depth using the specified colormap,144    and returns the depth map image.145    """146    temp_filepath = None147    output_path = None148 149    allowed_extensions = {'png', 'jpg', 'jpeg', 'bmp', 'webp'}150    file_ext = file.filename.split('.')[-1].lower() if '.' in file.filename else ''151    if file_ext not in allowed_extensions:152        raise HTTPException(status_code=400, detail=f"Invalid file type '{file_ext}'. Allowed types: {allowed_extensions}")153 154    try:155        # --- CORRECT LOGIC BLOCK ---156        temp_filename = f"upload_{uuid.uuid4().hex[:8]}.{file_ext}"157        temp_filepath = os.path.join(UPLOAD_DIR, temp_filename)158 159        start_save_time = time.time()160        with open(temp_filepath, "wb") as buffer:161            shutil.copyfileobj(file.file, buffer)162        save_time = time.time() - start_save_time163        print(f"Uploaded file saved to: {temp_filepath} in {save_time:.3f} seconds")164 165        print(f"Using colormap: {colormap}")166        output_path = predict_depth(temp_filepath, OUTPUT_DIR, colormap=colormap)167 168        return FileResponse(169            output_path,170            media_type='image/png',171            filename=f"depth_{colormap}_{file.filename}"172        )173        # --- END CORRECT LOGIC BLOCK ---174 175    # REMOVED THE DUPLICATE TRY BLOCK THAT WAS HERE176 177    except ValueError as ve:178         print(f"Value Error processing file: {ve}")179         raise HTTPException(status_code=400, detail=str(ve))180    except Exception as e:181        print(f"An unexpected error occurred during prediction for {file.filename}: {e}")182        raise HTTPException(status_code=500, detail=f"Internal server error during prediction: {e}")183 184    finally:185        # --- Cleanup ---186        if temp_filepath and os.path.exists(temp_filepath):187            try:188                os.remove(temp_filepath)189                print(f"Removed temporary upload file: {temp_filepath}")190            except Exception as e_rem_in:191                 print(f"Warning: Error removing temporary input file {temp_filepath}: {e_rem_in}")192 193# --- Main execution block ---194if __name__ == "__main__":195    import uvicorn196    import os # Import os to potentially read PORT197 198    print("Starting FastAPI server for LOCAL DEVELOPMENT...")199 200    # Use environment variable for port if available (good practice), otherwise default201    port = int(os.environ.get("PORT", 8000))202    # For local testing, 127.0.0.1 is usually fine.203    # For testing container networking, you might use "0.0.0.0" here too.204    host = "127.0.0.1"205 206    uvicorn.run(207        "app:app",208        host=host,209        port=port,210        reload=True # Keep reload=True for local development ease211        # log_level="info" # Optional: Set log level212    )213 214    # NOTE FOR DEPLOYMENT:215    # Production servers (like Gunicorn with Uvicorn workers) are typically used216    # in deployment and configured separately (e.g., via Procfile or platform settings).217    # They usually bind to host="0.0.0.0" and use a $PORT environment variable.218    # This __main__ block is NOT executed by those servers.