midas / app.py
adpro's picture
Update app.py
0848216 verified
raw
history blame
2.68 kB
import io
import time
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
app = FastAPI()
# 🟢 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 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(...)):
# 🟢 Đọ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((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)
# 🟢 Bắt đầu đo thời gian dự đoán Depth Map
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")
end_time = time.time()
print(f"⏳ DPT xử lý trong {end_time - start_time:.4f} giây")
# 🟢 Đo thời gian xử lý đường đi
start_detect_time = time.time()
command = detect_path(depth_map)
end_detect_time = time.time()
print(f"⏳ detect_path() xử lý trong {end_detect_time - start_detect_time:.4f} giây")
return {"command": command}
def detect_path(depth_map):
"""Phân tích đường đi từ ảnh Depth Map"""
h, w = depth_map.shape
center_x = w // 2
scan_y = int(h * 0.8) # Quét dòng 80% từ trên xuống
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 - 40:center_x + 40])
# 🟢 Cải thiện logic xử lý
threshold = 100 # Ngưỡng phân biệt vật cản
if center_region > threshold:
return "forward"
elif left_region > right_region:
return "left"
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)