Update app.py

app.py

@@ -10,10 +10,23 @@ import uvicorn
 
 app = FastAPI()
 
-# 🟢 Tải mô hình MobileNetDepth (MobileNet v3 Large)
+# 🟢 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
+import sys
+sys.path.append(fastdepth_path)
+
+# 🟢 Import FastDepth sau khi đã tải về
+from FastDepth.models import MobileNetSkipAdd  
+
+# 🟢 Load mô hình FastDepth
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = torchvision.models.mobilenet_v3_large(pretrained=True).to(device)
-model.eval()
+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)
 
 @app.post("/analyze_path/")
 async def analyze_path(file: UploadFile = File(...)):
@@ -21,9 +34,9 @@ async def analyze_path(file: UploadFile = File(...)):
     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 MobileNetDepth
+    # 🟢 Chuyển đổi ảnh thành tensor phù hợp với FastDepth
     transform = torchvision.transforms.Compose([
-        torchvision.transforms.Resize((224, 224)),  # MobileNetDepth yêu cầu ảnh 224x224
+        torchvision.transforms.Resize((224, 224)),  # FastDepth yêu cầu ảnh 224x224
         torchvision.transforms.ToTensor(),
     ])
     img_tensor = transform(image).unsqueeze(0).to(device)
@@ -31,13 +44,15 @@ async def analyze_path(file: UploadFile = File(...)):
     # 🟢 Bắt đầu đo thời gian dự đoán Depth Map
     start_time = time.time()
     
-    # 🟢 Dự đoán Depth Map với MobileNetDepth
+    # 🟢 Dự đoán Depth Map với FastDepth
     with torch.no_grad():
-        depth_map = model(img_tensor)
-        depth_map = depth_map.squeeze().cpu().numpy()
+        depth_map = model(img_tensor).squeeze().cpu().numpy()
 
     end_time = time.time()
-    print(f"⏳ MobileNetDepth xử lý trong {end_time - start_time:.4f} giây")
+    print(f"⏳ FastDepth xử lý trong {end_time - start_time:.4f} giây")
+
+    # 🟢 Kiểm tra kích thước Depth Map
+    print(f"📏 Depth Map Shape: {depth_map.shape}")
 
     # 🟢 Đo thời gian xử lý đường đi
     start_detect_time = time.time()
@@ -49,6 +64,9 @@ 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