#!/usr/bin/env python3
import os
# from hobot_dnn import pyeasy_dnn as dnn
# from hobot_vio import libsrcampy as srcampy
import onnxruntime
import numpy as np
import cv2
import colorsys
from time import time
# detection model class names
from fcos_core.utils.env import setup_environment # noqa F401 isort:skip
import argparse
import os
import torch
from fcos_core.config import cfg
from fcos_core.data import make_data_loader
from fcos_core.engine.inference import inference
from fcos_core.modeling.detector import build_detection_model
from fcos_core.utils.checkpoint import DetectronCheckpointer
from fcos_core.utils.collect_env import collect_env_info
from fcos_core.utils.comm import synchronize, get_rank
from fcos_core.utils.logger import setup_logger
from fcos_core.utils.miscellaneous import mkdir
from collections import OrderedDict
# parser = argparse.ArgumentParser(description="Export model to the onnx format")
# parser.add_argument(
# "--config-file",
# default="/home/dl/projects/fcos_effici/configs/fcos/fcos_imprv_R_50_FPN_1x.yaml",
# metavar="FILE",
# help="path to config file",
# )
# parser.add_argument(
# "--output",
# default="/home/dl/projects/fcos_effici/fcos.onnx",
# metavar="FILE",
# help="path to the output onnx file",
# )
# parser.add_argument(
# "opts",
# help="Modify config options using the command-line",
# default=None,
# nargs=argparse.REMAINDER,
# )
#
# args = parser.parse_args()
#
# cfg.merge_from_file(args.config_file)
# cfg.merge_from_list(args.opts)
# cfg.freeze()
#
# assert cfg.MODEL.FCOS_ON, "This script is only tested for the detector FCOS."
#
# save_dir = ""
# logger = setup_logger("fcos_core", save_dir, get_rank())
# logger.info(cfg)
#
# logger.info("Collecting env info (might take some time)")
# logger.info("\n" + collect_env_info())
#
# model = build_detection_model(cfg)
# model.to(cfg.MODEL.DEVICE)
#
# output_dir = cfg.OUTPUT_DIR
# checkpointer = DetectronCheckpointer(cfg, model, save_dir=output_dir)
# _ = checkpointer.load(cfg.MODEL.WEIGHT)
#
# onnx_model = torch.nn.Sequential(OrderedDict([
# ('backbone', model.backbone),
# ('heads', model.rpn.head),
# ]))
def get_classes():
return np.array(["no-park", "park"])
# return np.array(["tissue", "sock", "shoe", "cable", "bin"])
# bgr格式图片转换成 NV12格式
def bgr2nv12_opencv(image):
height, width = image.shape[0], image.shape[1]
area = height * width
yuv420p = cv2.cvtColor(
image, cv2.COLOR_BGR2YUV_I420).reshape((area * 3 // 2,))
y = yuv420p[:area]
uv_planar = yuv420p[area:].reshape((2, area // 4))
uv_packed = uv_planar.transpose((1, 0)).reshape((area // 2,))
nv12 = np.zeros_like(yuv420p)
nv12[:height * width] = y
nv12[height * width:] = uv_packed
return nv12
#postprocess函数接受模型输出(model_output)、模型输入的高度和宽度(model_hw_shape)，以及其他可选参数。
def postprocess(model_output,
model_hw_shape,
origin_image=None,
origin_img_shape=None,
score_threshold=0.5,
nms_threshold=0.6,
dump_image=False):
input_height = model_hw_shape[0]
input_width = model_hw_shape[1]
if origin_image is not None:
origin_image_shape = origin_image.shape[0:2]
else:
origin_image_shape = origin_img_shape
#它首先根据输入图像是否存在来确定原始图像的形状(origin_image_shape)。然后，它调用decode函数对模型输出进行解码，得到预测的边界框。
prediction_bbox = decode(outputs=model_output,
score_threshold=score_threshold,
origin_shape=origin_image_shape,
input_size=512)
#使用非极大值抑制（NMS）算法对边界框进行筛选，去除重叠的边界框。然后，它将边界框转换为NumPy数组，并根据置信度分数对边界框进行排序。
prediction_bbox = nms(prediction_bbox, iou_threshold=nms_threshold)
#指定了dump_image为True并且原始图像存在，则调用draw_bboxs`函数在原始图像上绘制边界框。最后，它返回预测的边界框。
prediction_bbox = np.array(prediction_bbox)
topk = min(prediction_bbox.shape[0], 1000)
if topk != 0:
idx = np.argpartition(prediction_bbox[..., 4], -topk)[-topk:]
prediction_bbox = prediction_bbox[idx]
if dump_image and origin_image is not None:
draw_bboxs(origin_image, prediction_bbox)
return prediction_bbox
def draw_bboxs(image, bboxes, gt_classes_index=None, classes=get_classes()):
"""draw the bboxes in the original image
"""
num_classes = len(classes)
image_h, image_w, channel = image.shape
hsv_tuples = [(1.0 * x / num_classes, 1., 1.) for x in range(num_classes)]
colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
colors = list(
map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
colors))
fontScale = 0.5
bbox_thick = int(0.6 * (image_h + image_w) / 600)
for i, bbox in enumerate(bboxes):
coor = np.array(bbox[:4], dtype=np.int32)
if gt_classes_index == None:
class_index = int(bbox[5])
score = bbox[4]
else:
class_index = gt_classes_index[i]
score = 1
bbox_color = colors[class_index]
c1, c2 = (coor[0], coor[1]), (coor[2], coor[3])
cv2.rectangle(image, c1, c2, bbox_color, bbox_thick)
classes_name = classes[class_index]
bbox_mess = '%s: %.2f' % (classes_name, score)
t_size = cv2.getTextSize(bbox_mess,
0,
fontScale,
thickness=bbox_thick // 2)[0]
cv2.rectangle(image, c1, (c1[0] + t_size[0], c1[1] - t_size[1] - 3),
bbox_color, -1)
cv2.putText(image,
bbox_mess, (c1[0], c1[1] - 2),
cv2.FONT_HERSHEY_SIMPLEX,
fontScale, (0, 0, 0),
bbox_thick // 2,
lineType=cv2.LINE_AA)
print("{} is in the picture with confidence:{:.4f}".format(
classes_name, score))
# cv2.imwrite("demo.jpg", image)
return image
def decode(outputs, score_threshold, origin_shape, input_size=512):
#接受points和distance作为输入，计算边界框的四个坐标（左上角和右下角）
def _distance2bbox(points, distance):
x1 = points[..., 0] - distance[0][0]
y1 = points[..., 1] - distance[0][1]
x2 = points[..., 0] + distance[0][2]
y2 = points[..., 1] + distance[0][3]
return np.stack([x1, y1, x2, y2], -1)
def _scores(cls, ce):
cls = 1 / (1 + np.exp(-cls))#经过sigmoid处理分类得分
ce = 1 / (1 + np.exp(-ce))#sigmoid处理置信度得分
return np.sqrt(ce * cls)#得到分数
#用于对边界框进行处理。bbox（边界框）、stride（步长）、
# origin_shape（原始图像形状）和input_size（输入大小）。
# 首先，根据边界框的大小创建网格坐标（yv和xv），根据步长和坐标
# 计算出边界框的坐标。接下来，根据输入图像的形状和输入大小计算出缩放比例。
# 最后，将边界框乘以缩放比例进行调整，以适应原始图像的大小。
def _bbox(bbox, stride, origin_shape, input_size):
# l t r b | h, w = t, r
h, w = bbox.shape[2:4]
#这行代码 yv, xv = np.meshgrid(np.arange(h), np.arange(w)) 使用了 NumPy 中的 np.meshgrid 函数来创建一个坐标网格。
#np.arange(h) 生成一个从 0 到 h-1 的一维数组，表示垂直轴上的值。类似地，np.arange(w) 生成一个从 0 到 w-1 的一维数组，表示水平轴上的值。
#np.meshgrid 函数接受这两个一维数组作为输入，并返回两个二维�
