Update app.py

app.py

@@ -1,22 +1,30 @@
 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()
 
-# 🟢 Chọn thiết bị xử lý (GPU nếu có)
+# 🟢 Kiểm tra GPU
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# 🟢 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()
+# 🟢 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])
+])
 
 @app.post("/analyze_path/")
 async def analyze_path(file: UploadFile = File(...)):
@@ -24,27 +32,22 @@ async def analyze_path(file: UploadFile = File(...)):
     image_bytes = await file.read()
     image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
 
-    # 🔵 Resize ảnh để xử lý nhanh hơn
-    image = image.resize((384, 384))  # Giảm kích thước giúp tăng tốc độ xử lý
-
-    # 🟢 Chuẩn bị ảnh cho mô hình
-    inputs = feature_extractor(images=image, return_tensors="pt").to(device)
+    # 🔵 Resize và chuẩn hóa ảnh
+    input_tensor = transform(image).unsqueeze(0).to(device)
 
-    # 🟢 Bắt đầu đo thời gian dự đoán Depth Map
+    # 🟢 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():
-        outputs = model(**inputs)
-
-    # 🟢 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")
-
+        depth_map = midas(input_tensor)
     end_time = time.time()
-    print(f"⏳ DPT xử lý trong {end_time - start_time:.4f} giây")
+    print(f"⏳ MiDaS xử lý trong {end_time - start_time:.4f} giây")
+
+    # 🟢 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")
 
-    # 🟢 Đo thời gian xử lý đường đi
+    # 🟢 Xử lý phát hiện đường đi
     start_detect_time = time.time()
     command = detect_path(depth_map)
     end_detect_time = time.time()