conda create -n TIONE-test python=3.8
conda activate TIONE-test

import os
import logging
import torch
from typing import Dict, List, Union
import numpy as np
import cv2
cv2.setNumThreads(5)
import tiinfer
import tiinfer.utils
import base64
import urllib

from utils.general import non_max_suppression, scale_coords
from torchvision import transforms

# 如果使用加速，需添加加速库引用：tiacc_inference
# import tiacc_inference

# 未使用加速时模型路径
PYTORCH_FILE = "model/yolov5s_ts.pt"
# 加速后模型路径
# TIACC_FILE = "model/tiacc.pt"


class YoloV5DetectModel(tiinfer.Model):
    def __init__(self, model_dir: str):
        super().__init__(model_dir)
        # workers: 平台参数，设置进程数，非必填，默认值：1
        self.workers = 1
        # 以下为自定义参数，根据模型实际需求定义
        self.model = None
        self.imgsz=640  # 输入图片大小
        self.conf_thres=0.05
        self.iou_thres=0.45  # for NMS
        self.max_det = 20
        self.use_gpu = torch.cuda.is_available()
        self.device = torch.device('cuda:0' if self.use_gpu else 'cpu')
        self.resize = transforms.Resize((self.imgsz, self.imgsz))
        self.normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

    # 加载模型
    def load(self) -> bool:
        try:
            # 拼装模型路径，self.model_dir: 当前工作目录；当为加速模型时，请将'PYTORCH_FILE'替换为'TIACC_FILE'
            model_file = os.path.join(self.model_dir,  PYTORCH_FILE)
            logging.info("model_file %s", model_file)
            # 加载模型
            self.model = torch.jit.load(model_file, map_location=self.device)
            # 设置类别信息
            self.det_obj_names = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 
            'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 
            'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 
            'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 
            'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 
            'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 
            'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 
            'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 
            'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 
            'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']
            logging.info("load model to device %s" % (self.device) )
        except Exception as e:  # pylint: disable=broad-except
            logging.error("load model failed: %s", e)
            return False
        return True

    # 预处理请求体
    def preprocess(self, request: Dict) -> Dict:
        try:
            img_data = request["image"]
            # image_to_cv2_mat能够正确处理本地图片、图片url及base64编码的图片内容，将之转换为opencv2中的矩阵
            img_mat = tiinfer.utils.image_to_cv2_mat(img_data)
            # 纬度转换
            img_mat = torch.Tensor(img_mat).permute((2, 0, 1)).to(self.device)

            # 如果在加速场景，图片尺寸必须在加速配置的tensor范围内，平台提供图片等比扩缩与填充接口：tiinfer.scale_fill_image方便用户使用。
            # 例如：加速模型tensor设置的范围为[720，2080]，等比扩缩图片的范围[1080,1999],等比扩缩后最小边不在tensor范围内，需要做补边操作，例如下面最小边需要补边到800。
            # min_size=1080; #图片扩缩小边最小范围值
            # max_size=1999; #图片扩缩大边最大范围值
            # padding_limit=800; #填充时的短边上限
            # image, self.scaleratio = tiinfer.scale_fill_image(image, min_size, max_size, padding_limit, Horizontal.RIGHT,Vertical.LOWER) 具体使用指引。
            
            # 数据标准化
            c, h, w = img_mat.shape
            img_tensor = self.normalize(self.resize(img_mat).type(torch.float32)/255)
            return {"img_tensor": img_tensor.unsqueeze(0), "original_dims": [w, h]}
        except Exception as e:
            logging.error("Preprocess failed: %s" % (e))
            return {"error": "preprocess failed"}

    # 推理执行，参数request是preprocess方法的返回结果
    def predict(self, request: Dict) -> Dict:
        with torch.no_grad():
            try:
                # 获取预测结果
                out = self.model(request["img_tensor"])
                # 可选操作，根据实际情况执行
                # 对预测结果进行非极大值抑制
                preds, nms_t = non_max_suppression(out, self.conf_thres, self.iou_thres, max_det=self.max_det)
                data = {
                    "pred_out": preds[0].to('cpu'),
                    "original_dims": request["original_dims"]
                }
                return data
            except Exception as e:
                logging.error("Failed to predict" % (e))
                request["predict_err"] = str(e)
                return request

    # 后处理，参数request是predict方法的返回结果
    def postprocess(self, request: Dict) -> Dict:
        try:
            # 可选操作，根据模型输出与实际需求对预测结果进行后处理
            pred_out = request["pred_out"]
            original_dim = request["original_dims"]
            # 调整图片形状
            pred_bbox = scale_coords([self.imgsz, self.imgsz], pred_out[:, :4], original_dim).round()
            # 输出参数，根据网络实际情况定义
            labels = []
            confs = []
            label_indx = []
            boxes = []
            for i in range(pred_out.shape[0]):
                obj = pred_out[i, :].tolist()
                boxes.append(pred_bbox[i, :].tolist())
                confs.append(obj[4])
                label_indx.append(int(obj[5]))
                labels.append(self.det_obj_names[int(obj[5])])
            image_result = {
                "det_boxes": boxes,
                "det_labels": labels,
                "det_labels_idx": label_indx,
                "det_scores": confs
            }
            return {
                "result": {
                    "det_objs": image_result
                }
            }
        except Exception as e:
            logging.error("Postprocess failed: %s" % (e))
            request.pop("image", '')
            request.pop("instances", '')
            request.pop("predictions", '')
            request["post_process_err"] = str(e)
            return request

python3 -m tiinfer --timeout 30000

curl -X POST -H "Content-Type: application/json" http://127.0.0.1:8501/v1/models/m:predict -d @test.json