40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
from fastapi import FastAPI, HTTPException
|
|
from pydantic import BaseModel
|
|
from typing import List, Optional
|
|
|
|
class Distance(BaseModel):
|
|
distance_mm: float
|
|
confidence: float
|
|
detection_image_path: Optional[str] = None
|
|
|
|
class DistanceAPI:
|
|
def __init__(self, camera, detector, estimator):
|
|
self.camera = camera
|
|
self.detector = detector
|
|
self.estimator = estimator
|
|
|
|
async def get_distances(self) -> List[Distance]:
|
|
frame = self.camera.get_frame()
|
|
if frame is None:
|
|
raise HTTPException(status_code=500, detail="No frame available")
|
|
|
|
# 检测人物并保存可视化结果
|
|
persons, vis_path = self.detector.detect(
|
|
frame,
|
|
save_visualization=True,
|
|
save_path="output/detections"
|
|
)
|
|
|
|
distances = []
|
|
for person in persons:
|
|
bbox = person['bbox']
|
|
person_height = bbox[3] - bbox[1] # y2 - y1
|
|
distance = self.estimator.estimate_distance(person_height, frame.shape[0])
|
|
|
|
distances.append(Distance(
|
|
distance_mm=distance,
|
|
confidence=person['confidence'],
|
|
detection_image_path=vis_path
|
|
))
|
|
|
|
return distances |