Update app.py

app.py

@@ -1,60 +1,72 @@
-from fastapi import FastAPI, File, UploadFile
 import io
+import os
 import numpy as np
+import cv2
+import torch
+import torchvision
+from fastapi import FastAPI, File, UploadFile
 from PIL import Image
 import uvicorn
-import cv2
-from fastdepth import FastDepth
-model = FastDepth(pretrained=True)
-model.eval()
+
 app = FastAPI()
 
-import os
+# 🟢 Clone FastDepth từ GitHub (chỉ cần làm 1 lần)
 if not os.path.exists("fastdepth"):
     os.system("git clone https://github.com/dwofk/fast-depth.git fastdepth")
 
-from fastdepth import FastDepth  # Import sau khi clone
-
-app = FastAPI()
+# 🟢 Import FastDepth sau khi clone
+from fastdepth.models import MobileNetSkipAdd  # Model chính của FastDepth
 
 # 🟢 Load mô hình FastDepth
-model = FastDepth(pretrained=True)
-model.eval()
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = MobileNetSkipAdd()
+model.load_state_dict(torch.load("fastdepth/models/fastdepth_nyu.pt", map_location=device))
+model.eval().to(device)
 
-def analyzepath(image):
-    depth_map = model(image).squeeze().cpu().numpy()
-    return detect_path(depth_map)  # Xử lý đường đi nhanh hơn
 @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("L")
-    depth_map = np.array(image)
+    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
 
-    # 🟢 Lật ảnh (nếu cần)
-    flipped_depth_map = cv2.flip(depth_map, -1)
+    # 🟢 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(image).unsqueeze(0).to(device)
+
+    # 🟢 Dự đoán Depth Map với FastDepth
     with torch.no_grad():
         depth_map = model(img_tensor).squeeze().cpu().numpy()
+
+    # 🟢 Lật ngược ảnh (nếu cần)
+    flipped_depth_map = cv2.flip(depth_map, -1)
+
     # 🟢 Phân tích đường đi
     command = detect_path(flipped_depth_map)
 
     return {"command": command}
 
 def detect_path(depth_map):
-    _, thresh = cv2.threshold(depth_map, 200, 255, cv2.THRESH_BINARY)
-    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    """Phân tích đường đi từ ảnh Depth Map"""
+    h, w = depth_map.shape
+    center_x = w // 2
+    scan_y = h - 20  # Quét dòng gần đáy ảnh
 
-    if len(contours) == 0:
+    left_region = np.mean(depth_map[scan_y, :center_x])
+    right_region = np.mean(depth_map[scan_y, center_x:])
+    center_region = np.mean(depth_map[scan_y, center_x - 20:center_x + 20])
+
+    if center_region > 200:
         return "forward"
-    
-    left_region = np.mean(depth_map[:, :depth_map.shape[1]//3])
-    right_region = np.mean(depth_map[:, 2*depth_map.shape[1]//3:])
-    
-    if left_region > right_region:
+    elif left_region > right_region:
         return "left"
-    else:
+    elif right_region > left_region:
         return "right"
+    else:
+        return "backward"
+
+# 🟢 Chạy server FastAPI
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)