Update app.py

app.py

@@ -1,6 +1,5 @@
 import io
 import os
-import sys
 import time
 import numpy as np
 import cv2
@@ -12,34 +11,14 @@ import uvicorn
 
 app = FastAPI()
 
-# 🟢 Clone FastDepth từ GitHub nếu chưa có
-if not os.path.exists("fastdepth"):
-    os.system("git clone https://github.com/dwofk/fast-depth.git fastdepth")
-
-# 🟢 Thêm `fastdepth` vào `sys.path`
-sys.path.append(os.path.abspath("fastdepth"))
-
-# 🟢 Tải đúng file trọng số nếu chưa có
-weights_path = "fastdepth/models/fastdepth_nyu.pt"
-if not os.path.exists(weights_path):
-    print("🔻 Trọng số chưa có, đang tải từ GitHub...")
-    os.system(f"wget -O {weights_path} https://github.com/dwofk/fast-depth/raw/master/models/fastdepth_nyu.pt")
-else:
-    print("✅ Trọng số đã có sẵn.")
-
-# 🟢 Import FastDepth
-from fastdepth.models import MobileNetSkipAdd  
-
-# 🟢 Load mô hình FastDepth đúng trọng số
+# 🟢 Tải mô hình MiDaS từ PyTorch Hub
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = MobileNetSkipAdd(output_size=(224, 224))
-try:
-    model.load_state_dict(torch.load(weights_path, map_location=device))
-    print("✅ Mô hình FastDepth đã được load thành công!")
-except FileNotFoundError:
-    print("❌ Không tìm thấy file trọng số! Kiểm tra lại đường dẫn.")
+model = torch.hub.load("intel-isl/MiDaS", "MiDaS_small").to(device)  # 🟢 Dùng phiên bản nhẹ MiDaS_small
+model.eval()
 
-model.eval().to(device)
+# 🟢 Load transform cho MiDaS
+midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
+transform = midas_transforms.small_transform  # 🟢 Dùng transform phù hợp với MiDaS_small
 
 @app.post("/analyze_path/")
 async def analyze_path(file: UploadFile = File(...)):
@@ -53,25 +32,24 @@ async def analyze_path(file: UploadFile = File(...)):
     # 🟢 Lật ảnh trước khi tính toán Depth Map
     flipped_image = cv2.flip(image_np, -1)
     
-    # 🟢 Chuyển đổi lại thành ảnh PIL để đưa vào FastDepth
+    # 🟢 Chuyển đổi lại thành ảnh PIL để đưa vào MiDaS
     flipped_image_pil = Image.fromarray(flipped_image)
 
-    # 🟢 Chuyển đổi ảnh thành tensor (chuẩn hóa cho FastDepth)
-    transform = torchvision.transforms.Compose([
-        torchvision.transforms.Resize((224, 224)),
-        torchvision.transforms.ToTensor(),
-    ])
-    img_tensor = transform(flipped_image_pil).unsqueeze(0).to(device)
+    # 🟢 Chuyển đổi ảnh thành tensor phù hợp với MiDaS
+    img_tensor = transform(flipped_image_pil).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 MiDaS
     with torch.no_grad():
-        depth_map = model(img_tensor).squeeze().cpu().numpy()
+        depth_map = model(img_tensor)
+        depth_map = torch.nn.functional.interpolate(
+            depth_map.unsqueeze(1), size=flipped_image_pil.size[::-1], mode="bicubic", align_corners=False
+        ).squeeze().cpu().numpy()
 
     end_time = time.time()
-    print(f"⏳ FastDepth xử lý trong {end_time - start_time:.4f} giây")
+    print(f"⏳ MiDaS xử lý trong {end_time - start_time:.4f} giây")
 
     # 🟢 Đo thời gian xử lý đường đi
     start_detect_time = time.time()