#include <iostream>
#include <fstream>
#include <opencv2/opencv.hpp>
#include "yolov5.h"
#include "utils.h"

std::string model_type="yolov5";
std::string task_type="classify";
std::string task_mode="image";
std::string kmodel_path="yolov5n_cls.kmodel";
std::string image_path="test.jpg";
float conf_thres=0.35;
float nms_thres=0.65;
float mask_thres=0.5;
std::vector<std::string> labels={"apple","banana","orange"};
int debug_mode=2;

void setup() {
  // put your setup code here, to run once:
  cv::Mat ori_img = cv::imread(image_path);
  FrameSize image_wh={ori_img.cols,ori_img.rows};
  runtime_tensor image_tensor=mat_to_tensor(ori_img);
  std::vector<YOLOBbox> yolo_results;
  Yolov5 yolov5(task_type,task_mode,kmodel_path,image_wh,labels,conf_thres,nms_thres,mask_thres,debug_mode);
  yolov5.run(image_tensor,yolo_results);
  yolov5.draw_results(ori_img,yolo_results);
  cv::imwrite("result_"+model_type+"_"+task_type+".jpg",ori_img);
}

void loop() {
  // put your main code here, to run repeatedly:

}
