adpro commited on
Commit
8d36ed5
·
verified ·
1 Parent(s): a310278

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +68 -15
app.py CHANGED
@@ -1,25 +1,78 @@
 
 
 
 
1
  import torch
2
  from transformers import DPTFeatureExtractor, DPTForDepthEstimation
 
3
  from PIL import Image
4
- import requests
5
 
 
 
 
6
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
7
 
8
- # Load model chuyển sang GPU
9
- feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
10
- model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large").to(device)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
11
 
12
- # Tải ảnh
13
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
14
- image = Image.open(requests.get(url, stream=True).raw)
 
 
15
 
16
- # Chuẩn bị input
17
- inputs = feature_extractor(images=image, return_tensors="pt").to(device)
 
18
 
19
- # Dự đoán depth map
20
- with torch.no_grad():
21
- outputs = model(**inputs)
 
 
 
 
 
 
 
22
 
23
- # Lấy output chuyển về dạng ảnh
24
- depth_map = outputs.predicted_depth
25
- depth_map = depth_map.squeeze().cpu().numpy()
 
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)