模型推理文件

# -*- coding: utf-8 -*-   
"""   
用户自定义启动脚本，用于模型评测  
"""   
# 用户推理入口，函数名不限  
def detect(image_dir, output_dir):  
    # 用户处理脚本，此处省略
   # 示例推理过程，用户也可根据dataloader等完成推理过程 
    for img in os.listdir(image_dir): 
        img_path = os.path.join(image_dir, img) 
        img_data = cv2.imread(img_path) 
        detections = model.predict(img_data) 
   # 将单个样本的推理结果保存为如下的json文件 
    results = [] 
    for x1, y1, x2, y2, score, cls_pred in detections: 
        bbox = {} 
        bbox["xmin"] = int(x1) 
        bbox["ymin"] = int(y1) 
        bbox["xmax"] = int(x2) 
        bbox["ymax"] = int(y2) 
        results.append({ 
            "name": classes[int(cls_pred)], 
            "score": f"{score:0.4f}", 
            "bndbox": bbox
        }) 
    # 保证推理结果文件与图片（样本）同名存入结果路径（保存为Octopus-2d检测推理格式） 
    sample_dir = os.path.join(output_path, os.path.basename(image_path).split(".")[0])
    os.makedirs(sample_dir, exist_ok=True)
    json_file = os.path.join(sample_dir, os.path.basename(image_path).split(".")[0] + ".json")
    json.dump(results, open(json_file, "w")) 
# 推理脚本入口
if __name__ == '__main__':  
# 获取数据集路径与推理结果存放路径，用户根据命令行参数获取数据集位置和推理结果存储位置
    parser = argparse.ArgumentParser()  
    parser.add_argument("--image_dir", type=str, default=os.getenv("DATASET"")  
    parser.add_argument("--output_dir", type=str, default=os.getenv("INFER_RESULT")  
    args, unparsed = parser.parse_known_args()  
    # 接收数据路径合结果路径后，调用推理服务，根据用户实际推理方法或文件进行替换 
    detect(args.image_dir, args.output_dir)