Update app.py

app.py

@@ -1,54 +1,52 @@
 import io
 import time
-import torch
 import numpy as np
 import cv2
+import torch
+from transformers import DPTFeatureExtractor, DPTForDepthEstimation
 from fastapi import FastAPI, File, UploadFile
 from PIL import Image
 import uvicorn
-from torchvision import transforms
 
-# 🟢 Tạo FastAPI
 app = FastAPI()
 
-# 🟢 Kiểm tra GPU
+# 🟢 Chọn thiết bị xử lý (GPU nếu có)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# 🟢 Tải model MiDaS
-midas = torch.hub.load("isl-org/MiDaS", "DPT_Swin2_L_384")
-midas.to(device)
-midas.eval()
-
-# 🟢 Chuẩn bị bộ tiền xử lý ảnh
-transform = transforms.Compose([
-    transforms.Resize((384, 384)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-])
+# 🟢 Tải model DPT-Hybrid thay cho ZoeDepth để tăng tốc
+feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-hybrid-midas")
+model = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas").to(device)
+model.eval()
 
 @app.post("/analyze_path/")
 async def analyze_path(file: UploadFile = File(...)):
+        # 🟢 Bắt đầu đo thời gian dự đoán Depth Map
+    start_time = time.time()
     # 🟢 Đọc file ảnh từ ESP32
     image_bytes = await file.read()
     image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+
+    # 🔵 Resize ảnh để xử lý nhanh hơn
+    image = image.resize((256, 256))  # Giảm kích thước giúp tăng tốc độ xử lý
     image_np = np.array(image)
     flipped_image = cv2.flip(image_np, -1)
-    # 🔵 Resize và chuẩn hóa ảnh
-    input_tensor = transform(flipped_image).unsqueeze(0).to(device)
+    # 🟢 Chuẩn bị ảnh cho mô hình
+    inputs = feature_extractor(images=flipped_image, return_tensors="pt").to(device)
 
-    # 🟢 Dự đoán Depth Map với MiDaS
-    start_time = time.time()
+
+
+    # 🟢 Dự đoán Depth Map với DPT-Hybrid
     with torch.no_grad():
-        depth_map = midas(input_tensor)
-    end_time = time.time()
-    print(f"⏳ MiDaS xử lý trong {end_time - start_time:.4f} giây")
+        outputs = model(**inputs)
 
-    # 🟢 Chuẩn hóa ảnh Depth Map
-    depth_map = depth_map.squeeze().cpu().numpy()
-    depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min()) * 255
-    depth_map = depth_map.astype("uint8")
+    # 🟢 Xử lý ảnh sau khi dự đoán
+    predicted_depth = outputs.predicted_depth.squeeze().cpu().numpy()
+    depth_map = (predicted_depth * 255 / predicted_depth.max()).astype("uint8")
+
+    end_time = time.time()
+    print(f"⏳ DPT xử lý trong {end_time - start_time:.4f} giây")
 
-    # 🟢 Xử lý phát hiện đường đi
+    # 🟢 Đo thời gian xử lý đường đi
     start_detect_time = time.time()
     command = detect_path(depth_map)
     end_detect_time = time.time()