import os import tempfile import numpy as np import cv2 from pathlib import Path import logging from transformers import DepthProImageProcessorFast, DepthProForDepthEstimation import torch from PIL import Image from fastapi import FastAPI, File, UploadFile, Form, HTTPException from fastapi.responses import JSONResponse, HTMLResponse from typing import Any, Dict, List, Tuple, Union import pillow_heif import json from depth_pro.utils import load_rgb, extract_exif # Initialize FastAPI app app = FastAPI( title="Depth Pro Distance Estimation", description="Estimate distance and depth using Apple's Depth Pro model", version="1.0.0", docs_url="/docs", redoc_url="/redoc" ) # Force CPU usage device = 'cpu' def initialize_depth_pipeline(): """Initialize the Depth Pro pipeline""" try: print("Initializing Depth Pro pipeline...") image_processor = DepthProImageProcessorFast.from_pretrained("apple/DepthPro-hf") model = DepthProForDepthEstimation.from_pretrained("apple/DepthPro-hf").to(device) return model, image_processor except Exception as e: print(f"Error initializing pipeline: {e}") print("Falling back to dummy pipeline...") return None class DepthEstimator: def __init__(self, model=None, image_processor=None): self.device = torch.device('cpu') # Force CPU print("Initializing Depth Pro estimator...") self.model = model self.image_processor = image_processor print("Depth Pro estimator initialized successfully!") def estimate_depth(self, image_path): try: # Load image image = Image.open(image_path) # Resize image for processing resized_image, new_size = self.resize_image(image_path) rgb_image = load_rgb(resized_image.name) f_px = rgb_image[-1] eval_image = rgb_image[0] # Perform inference using model inputs = self.image_processor(eval_image, return_tensors="pt").to(self.device) with torch.no_grad(): outputs = self.model(**inputs) post_processed_output = self.image_processor.post_process_depth_estimation( outputs, target_sizes=[(new_size[1], new_size[0])], ) result = post_processed_output[0] field_of_view = result["field_of_view"] focal_length = result["focal_length"] depth = result["predicted_depth"] # Convert to numpy if needed if isinstance(depth, torch.Tensor): depth = depth.detach().cpu().numpy() elif not isinstance(depth, np.ndarray): depth = np.array(depth) # Estimate focal length (rough estimation) print(f_px,focal_length) return depth, new_size, focal_length except Exception as e: print(f"Error in depth estimation: {e}") return None, None, None def resize_image(self, image_path, max_size=1536): with Image.open(image_path) as img: ratio = max_size / max(img.size) new_size = (int(img.size[0] * ratio), int(img.size[1] * ratio)) img = img.resize(new_size, Image.Resampling.LANCZOS) with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file: img.save(temp_file, format="PNG") return temp_file, new_size def find_topmost_pixel(mask): '''Top Pixel from footpath mask''' footpath_pixels = np.where(mask > 0) if len(footpath_pixels[0]) == 0: return None min_y = np.min(footpath_pixels[0]) top_pixels_mask = footpath_pixels[0] == min_y top_x_coords = footpath_pixels[1][top_pixels_mask] center_idx = len(top_x_coords) // 2 return (min_y, top_x_coords[center_idx]) def find_bottommost_footpath_pixel(mask, topmost_pixel): """Find the bottommost pixel perpendicular to the topmost pixel within the mask""" if topmost_pixel is None: return None top_y, top_x = topmost_pixel # Find all mask pixels in the same x-column as the topmost pixel mask_y_coords, mask_x_coords = np.where(mask > 0) column_mask = mask_x_coords == top_x column_y_coords = mask_y_coords[column_mask] if len(column_y_coords) == 0: # If no pixels in the same column, find the bottommost pixel in the entire mask footpath_pixels = np.where(mask > 0) if len(footpath_pixels[0]) == 0: return None max_y = np.max(footpath_pixels[0]) bottom_pixels_mask = footpath_pixels[0] == max_y bottom_x_coords = footpath_pixels[1][bottom_pixels_mask] center_idx = len(bottom_x_coords) // 2 return (max_y, bottom_x_coords[center_idx]) # Find the bottommost pixel in the same x-column max_y_in_column = np.max(column_y_coords) return (max_y_in_column, top_x) def estimate_real_world_distance(depth_map, topmost_pixel, mask): """Estimate real-world distance between two pixels using depth information""" if topmost_pixel is None or depth_map is None: return None # Find the bottommost pixel perpendicular to the topmost pixel bottommost_pixel = find_bottommost_footpath_pixel(mask, topmost_pixel) if bottommost_pixel is None: return None top_y, top_x = topmost_pixel bottom_y, bottom_x = bottommost_pixel # Ensure coordinates are within bounds if (top_y >= depth_map.shape[0] or top_x >= depth_map.shape[1] or bottom_y >= depth_map.shape[0] or bottom_x >= depth_map.shape[1]): return None topmost_depth = depth_map[top_y, top_x] bottommost_depth = depth_map[bottom_y, bottom_x] # Check if depth values are valid if np.isnan(topmost_depth) or np.isnan(bottommost_depth): print("Invalid depth values (NaN) found") return None distance_meters = float(topmost_depth - bottommost_depth) print(f"Distance calculation:") print(f" Topmost pixel: ({top_y}, {top_x}) = {topmost_depth:.3f}m") print(f" Bottommost pixel: ({bottom_y}, {bottom_x}) = {bottommost_depth:.3f}m") print(f" Distance: {distance_meters:.3f}m") return distance_meters # Initialize depth estimator globally print("Initializing Depth Pro pipeline...") depth_model, image_processor = initialize_depth_pipeline() depth_estimator = DepthEstimator(depth_model, image_processor) @app.get("/health") async def health_check(): """Health check endpoint for Docker""" return {"status": "healthy", "service": "Depth Pro Distance Estimation"} @app.get("/api") async def api_info(): """API information endpoint""" return { "message": "Depth Pro Distance Estimation API", "docs": "/docs", "health": "/health", "estimate_endpoint": "/estimate-depth" } @app.post("/estimate-depth") async def estimate_depth_endpoint(file: UploadFile = File(...), mask: UploadFile = File(...)): """FastAPI endpoint for depth estimation and distance calculation""" try: # Save uploaded file temporarily with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file: content = await file.read() temp_file.write(content) temp_file_path = temp_file.name # Save uploaded mask temporarily with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as mtemp_file: content = await mask.read() mtemp_file.write(content) temp_file_path_mask = mtemp_file.name # Load image for pixel detection image = cv2.imread(temp_file_path) mask = cv2.imread(temp_file_path_mask) if image is None or mask is None: return JSONResponse( status_code=400, content={"error": "Could not load image or mask"} ) # Estimate depth depth_map, new_size, focal_length_px = depth_estimator.estimate_depth(temp_file_path) if depth_map is None: return JSONResponse( status_code=500, content={"error": "Depth estimation failed"} ) # Resize image and mask to match depth map size resized_image = cv2.resize(image, new_size) resized_mask = cv2.resize(mask, new_size) # Convert mask to grayscale if it's not already if len(resized_mask.shape) == 3: resized_mask = cv2.cvtColor(resized_mask, cv2.COLOR_BGR2GRAY) # Find key pixels from the mask topmost_pixel = find_topmost_pixel(resized_mask) # Calculate distance distance_meters = estimate_real_world_distance(depth_map, topmost_pixel, resized_mask) # Clean up os.unlink(temp_file_path) os.unlink(temp_file_path_mask) result = { "depth_map_shape": depth_map.shape, "focal_length_px": float(focal_length_px) if focal_length_px is not None else None, "topmost_pixel": [ int(topmost_pixel[0]), int(topmost_pixel[1])] if topmost_pixel else None, "distance_meters": distance_meters, "depth_stats": { "min_depth": float(np.min(depth_map)), "max_depth": float(np.max(depth_map)), "mean_depth": float(np.mean(depth_map)) } } return JSONResponse(content=result) except Exception as e: # Clean up on error if 'temp_file_path' in locals(): try: os.unlink(temp_file_path) except: pass if 'temp_file_path_mask' in locals(): try: os.unlink(temp_file_path_mask) except: pass return JSONResponse( status_code=500, content={"error": str(e)} ) @app.get("/", response_class=HTMLResponse) async def root(): """Root endpoint with simple HTML interface""" html_content = """ Depth Pro Distance Estimation

🔍 Depth Pro Distance Estimation

Upload an image and a footpath mask to estimate depth and calculate distances using Apple's Depth Pro model

Upload Image and Mask

Analysis Results:

🔗 API Endpoints

POST /estimate-depth - Upload image and footpath mask for depth estimation

GET /docs - API documentation

GET /health - Health check

✨ Features

""" return HTMLResponse(content=html_content) # FastAPI app is ready to run if __name__ == "__main__": import uvicorn uvicorn.run( app, host="0.0.0.0", port=7860, log_level="info", access_log=True )