Update app.py

app.py

@@ -1,59 +1,48 @@
 import io
-import os  # 🟢 Thêm dòng này để tránh lỗi NameError
-import sys
 import time
 import numpy as np
 import cv2
 import torch
-import torchvision
+from transformers import AutoImageProcessor, ZoeDepthForDepthEstimation
 from fastapi import FastAPI, File, UploadFile
 from PIL import Image
 import uvicorn
 
 app = FastAPI()
 
-# 🟢 Clone FastDepth nếu chưa có
-fastdepth_path = "FastDepth"
-if not os.path.exists(fastdepth_path):
-    os.system("git clone https://github.com/dwofk/fast-depth.git FastDepth")
-
-# 🟢 Thêm FastDepth vào sys.path để import được
-sys.path.append(fastdepth_path)
-
-# 🟢 Import FastDepth sau khi đã tải về
-from FastDepth.models import MobileNetSkipAdd  
-
-# 🟢 Load mô hình FastDepth
+# 🟢 Tải mô hình ZoeDepth từ Hugging Face
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = MobileNetSkipAdd(output_size=(224, 224))  # 🟢 FastDepth hỗ trợ đầu ra 224x224
-model.load_state_dict(torch.load(f"{fastdepth_path}/models/fastdepth_nyu.pt", map_location=device))
-model.eval().to(device)
+image_processor = AutoImageProcessor.from_pretrained("Intel/zoedepth-nyu-kitti")
+model = ZoeDepthForDepthEstimation.from_pretrained("Intel/zoedepth-nyu-kitti").to(device)
+model.eval()
 
 @app.post("/analyze_path/")
 async def analyze_path(file: UploadFile = File(...)):
     # 🟢 Đọc file ảnh từ ESP32
     image_bytes = await file.read()
     image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-    
-    # 🟢 Chuyển đổi ảnh thành tensor phù hợp với FastDepth
-    transform = torchvision.transforms.Compose([
-        torchvision.transforms.Resize((224, 224)),  # FastDepth yêu cầu ảnh 224x224
-        torchvision.transforms.ToTensor(),
-    ])
-    img_tensor = transform(image).unsqueeze(0).to(device)
+
+    # 🟢 Chuẩn bị ảnh cho mô hình ZoeDepth
+    inputs = image_processor(images=image, return_tensors="pt").to(device)
 
     # 🟢 Bắt đầu đo thời gian dự đoán Depth Map
     start_time = time.time()
-    
-    # 🟢 Dự đoán Depth Map với FastDepth
+
+    # 🟢 Dự đoán Depth Map với ZoeDepth
     with torch.no_grad():
-        depth_map = model(img_tensor).squeeze().cpu().numpy()
+        outputs = model(**inputs)
 
-    end_time = time.time()
-    print(f"⏳ FastDepth xử lý trong {end_time - start_time:.4f} giây")
+    # 🟢 Xử lý ảnh sau khi dự đoán
+    post_processed_output = image_processor.post_process_depth_estimation(
+        outputs,
+        source_sizes=[(image.height, image.width)],
+    )
+    predicted_depth = post_processed_output[0]["predicted_depth"]
+    depth_map = predicted_depth * 255 / predicted_depth.max()
+    depth_map = depth_map.detach().cpu().numpy().astype("uint8")
 
-    # 🟢 Kiểm tra kích thước Depth Map
-    print(f"📏 Depth Map Shape: {depth_map.shape}")
+    end_time = time.time()
+    print(f"⏳ ZoeDepth xử lý trong {end_time - start_time:.4f} giây")
 
     # 🟢 Đo thời gian xử lý đường đi
     start_detect_time = time.time()
@@ -65,9 +54,6 @@ async def analyze_path(file: UploadFile = File(...)):
 
 def detect_path(depth_map):
     """Phân tích đường đi từ ảnh Depth Map"""
-    if len(depth_map.shape) != 2:  # 🟢 Kiểm tra nếu depth_map không phải 2D
-        raise ValueError("Depth map không phải ảnh 2D hợp lệ!")
-
     h, w = depth_map.shape
     center_x = w // 2
     scan_y = h - 20  # Quét dòng gần đáy ảnh