最近新的项目有一个需求,需要进行目标检测,返回物品的坐标。
部署时候采用的是对并发支持比较好的fastapi框架,没有采用flask框架。
但是两种部署框架都差不多都很方便,下面是部署的代码。
图片通过base64编码传入服务器,返回目标的位置信息。
同时进行了模型载入的优化,模型权重文件一次性加载入内存。
后面进行推理的时候就不需要对模型进行加载了,减少了模型推理时间,便于cpu推理。
很适用于线上要求时延较高的场景,代码如下:
fastapi接口接收图片代码:
模型加载配置文件代码:
json配置文件:
预测推理代码:
def predict(opt, model, img):
source, imgsz = \
opt['source'], opt['imgsz']
device = select_device(opt['device']) # 选择设备
half = device.type != 'cpu' # half precision only supported on CUDA
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
if half:
model.half() # to FP16
dataset = LoadImages(opt['source'], img_size=imgsz)
for path, img, im0s, _ in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# 前向推理
pred = model(img, augment=opt['augment'])[0]
# Apply NMS(非极大抑制)
pred = non_max_suppression(pred, opt['conf_thres'], opt['iou_thres'], classes=opt['classes'],
agnostic=opt['agnostic_nms'])
for i, det in enumerate(pred): # detections per image
p, s, im0 = path, '', im0s
if det is not None and len(det):
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
boxes_detected = [] # 检测结果
for *xyxy, conf, cls in reversed(det):
xyxy_list = (torch.tensor(xyxy).view(1, 4)).view(-1).tolist()
boxes_detected.append({
"clear_flag": 1,
"left_top_x": int(xyxy_list[0]) ,
"left_top_y": int(xyxy_list[1]) ,
"w": int(int(xyxy_list[2])-int(xyxy_list[0])) ,
"h": int(int(xyxy_list[3])-int(xyxy_list[1]))
})
results = {
"code":200,"msg":'success' ,"data": boxes_detected}
return results