midas / app.py
adpro's picture
Update app.py
903cef3 verified
raw
history blame
2.86 kB
import io
import os
import sys
import numpy as np
import cv2
import torch
import torchvision
from fastapi import FastAPI, File, UploadFile
from PIL import Image
import uvicorn
app = FastAPI()
# 🟢 Clone FastDepth từ GitHub nếu chưa có
if not os.path.exists("fastdepth"):
os.system("git clone https://github.com/dwofk/fast-depth.git fastdepth")
# 🟢 Thêm `fastdepth` vào `sys.path`
sys.path.append(os.path.abspath("fastdepth"))
# 🟢 Đảm bảo file trọng số tồn tại
weights_path = "fastdepth/models/fastdepth_nyu.pt"
if not os.path.exists(weights_path):
os.system(f"wget -O {weights_path} https://github.com/dwofk/fast-depth/raw/master/models/fastdepth_nyu.pt")
# 🟢 Import FastDepth
from fastdepth.models import MobileNetSkipAdd
# 🟢 Load mô hình FastDepth
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MobileNetSkipAdd(output_size=(224, 224))
model.load_state_dict(torch.load(weights_path, map_location=device))
model.eval().to(device)
@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")
# 🟢 Chuyển ảnh sang NumPy để lật đúng chiều
image_np = np.array(image)
# 🟢 Lật ảnh trước khi tính toán Depth Map
flipped_image = cv2.flip(image_np, -1) # Lật trái ↔ phải, trên ↔ xuống
# 🟢 Chuyển đổi lại thành ảnh PIL để đưa vào FastDepth
flipped_image_pil = Image.fromarray(flipped_image)
# 🟢 Chuyển đổi ảnh thành tensor (chuẩn hóa cho FastDepth)
transform = torchvision.transforms.Compose([
torchvision.transforms.Resize((224, 224)),
torchvision.transforms.ToTensor(),
])
img_tensor = transform(flipped_image_pil).unsqueeze(0).to(device)
# 🟢 Dự đoán Depth Map với FastDepth
with torch.no_grad():
depth_map = model(img_tensor).squeeze().cpu().numpy()
# 🟢 Phân tích đường đi
command = detect_path(depth_map)
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 = h - 20 # Quét dòng gần đáy ảnh
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 - 20:center_x + 20])
if center_region > 200:
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)