Update app.py
Browse files
app.py
CHANGED
@@ -1,25 +1,78 @@
|
|
|
|
|
|
|
|
|
|
1 |
import torch
|
2 |
from transformers import DPTFeatureExtractor, DPTForDepthEstimation
|
|
|
3 |
from PIL import Image
|
4 |
-
import
|
5 |
|
|
|
|
|
|
|
6 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
7 |
|
8 |
-
#
|
9 |
-
feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-
|
10 |
-
model = DPTForDepthEstimation.from_pretrained("Intel/dpt-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
11 |
|
12 |
-
|
13 |
-
|
14 |
-
|
|
|
|
|
15 |
|
16 |
-
|
17 |
-
|
|
|
18 |
|
19 |
-
#
|
20 |
-
|
21 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
22 |
|
23 |
-
#
|
24 |
-
|
25 |
-
|
|
|
1 |
+
import io
|
2 |
+
import time
|
3 |
+
import numpy as np
|
4 |
+
import cv2
|
5 |
import torch
|
6 |
from transformers import DPTFeatureExtractor, DPTForDepthEstimation
|
7 |
+
from fastapi import FastAPI, File, UploadFile
|
8 |
from PIL import Image
|
9 |
+
import uvicorn
|
10 |
|
11 |
+
app = FastAPI()
|
12 |
+
|
13 |
+
# 🟢 Chọn thiết bị xử lý (GPU nếu có)
|
14 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
15 |
|
16 |
+
# 🟢 Tải model DPT-Hybrid thay cho ZoeDepth để tăng tốc
|
17 |
+
feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-hybrid")
|
18 |
+
model = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid").to(device)
|
19 |
+
model.eval()
|
20 |
+
|
21 |
+
@app.post("/analyze_path/")
|
22 |
+
async def analyze_path(file: UploadFile = File(...)):
|
23 |
+
# 🟢 Đọc file ảnh từ ESP32
|
24 |
+
image_bytes = await file.read()
|
25 |
+
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
|
26 |
+
|
27 |
+
# 🔵 Resize ảnh để xử lý nhanh hơn
|
28 |
+
image = image.resize((384, 384)) # Giảm kích thước giúp tăng tốc độ xử lý
|
29 |
+
|
30 |
+
# 🟢 Chuẩn bị ảnh cho mô hình
|
31 |
+
inputs = feature_extractor(images=image, return_tensors="pt").to(device)
|
32 |
+
|
33 |
+
# 🟢 Bắt đầu đo thời gian dự đoán Depth Map
|
34 |
+
start_time = time.time()
|
35 |
+
|
36 |
+
# 🟢 Dự đoán Depth Map với DPT-Hybrid
|
37 |
+
with torch.no_grad():
|
38 |
+
outputs = model(**inputs)
|
39 |
+
|
40 |
+
# 🟢 Xử lý ảnh sau khi dự đoán
|
41 |
+
predicted_depth = outputs.predicted_depth.squeeze().cpu().numpy()
|
42 |
+
depth_map = (predicted_depth * 255 / predicted_depth.max()).astype("uint8")
|
43 |
+
|
44 |
+
end_time = time.time()
|
45 |
+
print(f"⏳ DPT xử lý trong {end_time - start_time:.4f} giây")
|
46 |
+
|
47 |
+
# 🟢 Đo thời gian xử lý đường đi
|
48 |
+
start_detect_time = time.time()
|
49 |
+
command = detect_path(depth_map)
|
50 |
+
end_detect_time = time.time()
|
51 |
+
print(f"⏳ detect_path() xử lý trong {end_detect_time - start_detect_time:.4f} giây")
|
52 |
+
|
53 |
+
return {"command": command}
|
54 |
|
55 |
+
def detect_path(depth_map):
|
56 |
+
"""Phân tích đường đi từ ảnh Depth Map"""
|
57 |
+
h, w = depth_map.shape
|
58 |
+
center_x = w // 2
|
59 |
+
scan_y = int(h * 0.8) # Quét dòng 80% từ trên xuống
|
60 |
|
61 |
+
left_region = np.mean(depth_map[scan_y, :center_x])
|
62 |
+
right_region = np.mean(depth_map[scan_y, center_x:])
|
63 |
+
center_region = np.mean(depth_map[scan_y, center_x - 40:center_x + 40])
|
64 |
|
65 |
+
# 🟢 Cải thiện logic xử lý
|
66 |
+
threshold = 100 # Ngưỡng phân biệt vật cản
|
67 |
+
if center_region > threshold:
|
68 |
+
return "forward"
|
69 |
+
elif left_region > right_region:
|
70 |
+
return "left"
|
71 |
+
elif right_region > left_region:
|
72 |
+
return "right"
|
73 |
+
else:
|
74 |
+
return "backward"
|
75 |
|
76 |
+
# 🟢 Chạy server FastAPI
|
77 |
+
if __name__ == "__main__":
|
78 |
+
uvicorn.run(app, host="0.0.0.0", port=7860)
|