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

# If acceleration is used, add a reference to the acceleration library: tiacc_inference.
# import tiacc_inference

# Model path when acceleration is not used.
PYTORCH_FILE = "model/yolov5s_ts.pt"
# Model path after acceleration.
# TIACC_FILE = "model/tiacc.pt"


class YoloV5DetectModel(tiinfer.Model):
    def __init__(self, model_dir: str):
        super().__init__(model_dir)
        # workers: a platform parameter, used to set the number of processes. The parameter is optional and defaults to 1.
        self.workers = 1
        # The following parameters can be customized and defined based on model needs.
        self.model = None
        self.imgsz=640 # Input image size.
        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])

    # Load model.
    def load(self) -> bool:
        try:
            # Assemble model path. self.model_dir: current working directory. Replace 'PYTORCH_FILE' with 'TIACC_FILE' for an acceleration model.
            model_file = os.path.join(self.model_dir,  PYTORCH_FILE)
            logging.info("model_file %s", model_file)
            # Load model.
            self.model = torch.jit.load(model_file, map_location=self.device)
            # Set class information.
            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

    # Preprocess the request body.
    def preprocess(self, request: Dict) -> Dict:
        try:
            img_data = request["image"]
            # image_to_cv2_mat can correctly handle local images, image URLs, and base64-encoded image content, converting them into matrices in OpenCV 2.
            img_mat = tiinfer.utils.image_to_cv2_mat(img_data)
            # Latitude conversion
            img_mat = torch.Tensor(img_mat).permute((2, 0, 1)).to(self.device)

            # The image size should be within the tensor range of the acceleration configurations in an acceleration scenario. The platform provides interfaces for proportionally scaling and padding images: tiinfer.scale_fill_image.
            # For example: The range of tensor settings for the acceleration model is [720, 2080], and the range of proportional scaling of images is [1080, 1999]. If the minimum side after proportional scaling is not within the tensor range, edge padding is required. For example, the minimum side below needs to be padded to 800.
            # min_size=1080; # Minimum side value range for image scaling.
            # max_size=1999; # Maximum side value range for image scaling.
            # padding_limit=800; # Upper limit of the short side during padding.
            # Specific use guide for image, self.scaleratio = tiinfer.scale_fill_image(image, min_size, max_size, padding_limit, Horizontal.RIGHT,Vertical.LOWER).
            
            # Data standardization
            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"}

    # Perform inference. The request parameter is the returned result of the preprocess method.
    def predict(self, request: Dict) -> Dict:
        with torch.no_grad():
            try:
                # Get the forecast results.
                out = self.model(request["img_tensor"])
                # Optional. Executed according to actual conditions.
                # Perform non-maximum suppression on the forecast result.
                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

    # Post-processing. The request parameter is the returned result of the forecast method.
    def postprocess(self, request: Dict) -> Dict:
        try:
            # Optional. Perform post-processing on the forecast results based on the model output and actual needs.
            pred_out = request["pred_out"]
            original_dim = request["original_dims"]
            # Adjust the image shape.
            pred_bbox = scale_coords([self.imgsz, self.imgsz], pred_out[:, :4], original_dim).round()
            # Output parameters, defined based on actual network circumstances.
            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