Algorithm-Repo/SourceCode/MachineLearningUtils/VINNOFramework/YOLO/ultralytics/trainsdk/TrainSdk.py

import http.client
import shutil
import os


# 获取训练文件夹中的文件数量
# token:训练用的授权码
def get_file_count(token):
    uri = "localhost:8888"
    conn = http.client.HTTPConnection(uri)
    conn.request("GET", "/GetFolderFileCount?token=" + token)
    res = conn.getresponse()
    if res.status == 200:
        count_data = res.read(4)
        #print(count_data)
        count = int.from_bytes(count_data, byteorder='little', signed=True)
        return count
    else:
        reason = str(res.reason)
        if len(reason) > 0:
            reason = ", reason: " + reason
        raise Exception("Get file count failed, could be url error or token error. Code: " + str(res.status) + reason)


# 获取文件和其对应的标注数据，名称，是否是视频
# token:训练用的授权码
# index:文件序号
def get_labeled_file(token, index):
    uri = "localhost:8888"
    conn = http.client.HTTPConnection(uri)
    conn.request("GET", "/GetTrainFile?token=" + token + "&index=" + str(index))
    res = conn.getresponse()
    if res.status == 200:
        file_size_data = res.read(4)
        file_size = int.from_bytes(file_size_data, byteorder='little', signed=True)
        file_data = res.read(file_size)
        label_size_data = res.read(4)
        label_size = int.from_bytes(label_size_data, byteorder='little', signed=True)
        label_data = res.read(label_size).decode("utf-8")
        file_name_size_data = res.read(4)
        file_name_size = int.from_bytes(file_name_size_data, byteorder='little', signed=True)
        file_name = res.read(file_name_size).decode("utf-8")
        file_isVideo_size_data = res.read(4)
        file_isVideo_size = int.from_bytes(file_isVideo_size_data, byteorder='little', signed=True)
        file_isVideo = res.read(file_isVideo_size).decode()
        return file_data, label_data, file_name, file_isVideo
    else:
        reason = str(res.reason)
        if len(reason) > 0:
            reason = ", reason: " + reason
        raise Exception(
            "Get file and label failed, could be url error or token error. Code: " + str(res.status) + reason)


# 将训练结果的模型保存到系统可识别的路径下
# trainedFile: 要保存的训练结果
def save_output_model(trainedFile: object) -> object:
    uri = "localhost:8888"
    conn = http.client.HTTPConnection(uri)
    conn.request("GET", "/GetModelOutputFolder")
    res = conn.getresponse()
    if res.status == 200:
        outputFolder = res.readlines()[0].decode("utf-8")
        sourceFolder, sourceFileName = os.path.split(trainedFile)
        outputFile = os.path.join(outputFolder, sourceFileName)
        shutil.copy(trainedFile, outputFile)
    else:
        reason = str(res.reason)
        if len(reason) > 0:
            reason = ", reason: " + reason
        raise Exception(
            "Save output model failed, could be url error or token error. Code: " + str(res.status) + reason)


# 获取测试文件夹中的文件数量
# token:授权码
def get_test_file_count(token):
    uri = "localhost:8888"
    conn = http.client.HTTPConnection(uri)
    conn.request("GET", "/GetTestFolderFileCount?token=" + token)
    res = conn.getresponse()
    if res.status == 200:
        count_data = res.read(4)
        #print(count_data)
        count = int.from_bytes(count_data, byteorder='little', signed=True)
        return count
    else:
        reason = str(res.reason)
        if len(reason) > 0:
            reason = ", reason: " + reason
        raise Exception(
            "Get test file count failed, could be url error or token error. Code: " + str(res.status) + reason)


# 获取测试文件和其对应的标注数据，名称，是否是视频
# token:授权码
# index:文件序号
def get_test_labeled_file(token, index):
    uri = "localhost:8888"
    conn = http.client.HTTPConnection(uri)
    conn.request("GET", "/GetTestFile?token=" + token + "&index=" + str(index))
    res = conn.getresponse()
    if res.status == 200:
        file_size_data = res.read(4)
        file_size = int.from_bytes(file_size_data, byteorder='little', signed=True)
        file_data = res.read(file_size)
        label_size_data = res.read(4)
        label_size = int.from_bytes(label_size_data, byteorder='little', signed=True)
        label_data = res.read(label_size).decode("utf-8")
        file_name_size_data = res.read(4)
        file_name_size = int.from_bytes(file_name_size_data, byteorder='little', signed=True)
        file_name = res.read(file_name_size).decode("utf-8")
        file_isVideo_size_data = res.read(4)
        file_isVideo_size = int.from_bytes(file_isVideo_size_data, byteorder='little', signed=True)
        file_isVideo = res.read(file_isVideo_size).decode()
        return file_data, label_data, file_name, file_isVideo
    else:
        reason = str(res.reason)
        if len(reason) > 0:
            reason = ", reason: " + reason
        raise Exception(
            "Get test file and label failed, could be url error or token error. Code: " + str(res.status) + reason)