【无标题】

import csv
import os
from ultralytics import YOLO
import cv2
import pandas as pd
import numpy as np
model = YOLO("runs/detect/train/weights/best.pt")
paht_img = './2-0-1-50-2023-11-29-06-26-01_LM'
jpgname1 = os.listdir(paht_img)
jpgname = sorted(jpgname1, key=lambda x: os.path.getmtime(os.path.join(paht_img, x)))
#print(jpgname)
name_list = [os.path.join(paht_img, name) for name in jpgname]
compute_lc = []
start_lc = 1.052#初始里程
guizhen_conut = 0#轨枕计数
guizhen = [0, 0]
gz_last = 0
guizhen_mowei = 0#每张图片的最
dangqian_lc = 0#图像中轨枕的里程
file_data_1 = '2-0-1-50-2023-11-29-06-26-01_LM.csv'
file_data_2 = '2-0-1-50-2023-11-29-06-26-01_LMque_bomo.csv'
if os.path.isfile(file_data_1):
    os.remove(file_data_1)
if os.path.isfile(file_data_2):
    os.remove(file_data_2)
import shutil
shutil.rmtree(paht_img.replace('LM', 'LMceshi'), ignore_errors=True)
os.makedirs(paht_img.replace('LM', 'LMceshi'), exist_ok=True)

file = open(file_data_1, mode='a', newline='')#记录数据在csv
writer = csv.writer(file)
file_qb = open(file_data_2, mode='a', newline='')#记录数据在csv
writer_qb = csv.writer(file_qb)
for i,n in enumerate(name_list):
    results = model(n, device=0, conf=0.5, save=True)
    print(n)
    imgsave = n.replace('LM', 'LMceshi')
    img = cv2.imread(n)
    bq_flag = 0
    for result in results:
        boxes = result.boxes  # Boxes object for bbox outputs
        xbox1 = boxes.xywh.cpu().numpy()
        xcls = boxes.cls.cpu().numpy()
        xconf1 = boxes.conf.cpu().numpy()
        print('xbox44:', xbox1, 'xcls44:', xcls, 'xconf44:', xconf1)
        if (xbox1 is None):
            start_lc += 0.002
            continue
        #缺陷和波磨计数#######################################
        for j_qb, nn_qb in enumerate(xcls):
            if nn_qb == float(1.0):
                writer_qb.writerow([round(start_lc + xbox1[j_qb][1] * (2000 / 1870) * 0.000001, 6), 0, xconf1[j_qb]])
            elif nn_qb == float(2.0):
                writer_qb.writerow([0, round(start_lc + xbox1[j_qb][1] * (2000 / 1870) * 0.000001, 6), xconf1[j_qb]])
        ####################################################
        # 获取轨枕和标签数据
        box_sort = xbox1[xcls==float(0.0), :]
        xcls_sort1 = xcls[xcls==float(0.0)]
        if float(3.0) in xcls:
            box_sort = np.append(box_sort, xbox1[xcls==float(3.0), :], axis=0)
            len_cont = np.sum(xcls==float(3.0))
            print(len_cont)
            data_append = []
            for i in range(len_cont):
                data_append.append(float(3.0))
            xcls_sort = np.append(xcls_sort1, data_append)
            print(xcls_sort)
        else:
            xcls_sort = xcls_sort1

        box_count = box_sort.shape[0] #统计检测个数
        idex = np.lexsort([box_sort[:, 1], ])
        box_sort = box_sort[idex, :]
        print(box_sort, xcls_sort)
        xcls_sort = xcls_sort[idex]

        for j, nn in enumerate(xcls_sort):
            # 经过标签点处里程取整all(xconf1[(xcls==float(3.0))]) >= float(0.9)
            # 经过标签处，标签置信度大于0.75，则为判断为标签
            if nn == float(3.0) and any(x>=float(0.75) for x in xconf1[(xcls==float(3.0))]) :
                print(xconf1[(xcls==float(3.0))])
                x1 = int(box_sort[j][0] - box_sort[j][2] / 2)
                y1 = int(box_sort[j][1] - box_sort[j][3] / 2)
                x2 = int(box_sort[j][0] + box_sort[j][2] / 2)
                y2 = int(box_sort[j][1] + box_sort[j][3] / 2)
                flag_bq = box_sort[j][1] * (2000 / 1870) * 0.000001
                if box_sort[j][0] < float(1600.0):
                    continue
                #轨枕归零和里程增加
                # 经过标签里程置整，加上标签位置，重置当前轨枕位置   start_lc_lj = round(guizhen_conut/1670)
                start_lc = round(start_lc)
                bq_flag = 1
                guizhen_conut = 0
                cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), thickness=3, lineType=cv2.LINE_AA)
                cv2.putText(img, 'GZ:' + str(guizhen_conut) + '  ' + 'KM:' + str(start_lc), (500, y2),
                            cv2.FONT_HERSHEY_PLAIN, 3, (0, 0, 255), thickness=3, lineType=cv2.LINE_AA)
                #cv2.imwrite(imgsave, img)
            # 轨枕记录处理
            if nn == float(0.0):
                # 未存在标签时
                if bq_flag == 0:
                    # 轨枕前后相差小于200不统计轨枕数增加
                    if (abs(gz_last - start_lc - box_sort[j][1] * (2000 / 1870) * 0.000001) < float(0.0002)):
                        print(gz_last, round(start_lc + box_sort[j][1] * (2000 / 1870) * 0.000001, 6))
                        continue
                    guizhen_conut += 1  # 遇到类别为0轨枕的进行＋1
                    x1 = int(box_sort[j][0] - box_sort[j][2] / 2)
                    y1 = int(box_sort[j][1] - box_sort[j][3] / 2)
                    x2 = int(box_sort[j][0] + box_sort[j][2] / 2)
                    y2 = int(box_sort[j][1] + box_sort[j][3] / 2)
                    dangqian_lc = round(start_lc + box_sort[j][1] * (2000 / 1870) * 0.000001, 6)  # 取6位小数，精度1mm
                    writer.writerow([guizhen_conut, dangqian_lc, xconf1[(xcls == float(3.0))]])
                    cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), thickness=3, lineType=cv2.LINE_AA)
                    cv2.putText(img, 'GZ:' + str(guizhen_conut) + '  ' + 'KM:' + str(dangqian_lc), (x2, y2),
                                cv2.FONT_HERSHEY_PLAIN, 3, (0, 0, 255), thickness=3, lineType=cv2.LINE_AA)
                    gz_last = dangqian_lc
                # 存在标签时
                else:
                    guizhen_conut += 1  # 遇到类别为0轨枕的进行＋1
                    x1 = int(box_sort[j][0] - box_sort[j][2] / 2)
                    y1 = int(box_sort[j][1] - box_sort[j][3] / 2)
                    x2 = int(box_sort[j][0] + box_sort[j][2] / 2)
                    y2 = int(box_sort[j][1] + box_sort[j][3] / 2)
                    dangqian_lc = round(start_lc - flag_bq + box_sort[j][1] * (2000 / 1870) * 0.000001, 6)#取6位小数，精度1mm
                    writer.writerow([guizhen_conut, dangqian_lc, xconf1[(xcls == float(3.0))]])
                    cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), thickness=3, lineType=cv2.LINE_AA)
                    cv2.putText(img, 'GZ:' + str(guizhen_conut) + '  ' + 'KM:' + str(dangqian_lc), (x2, y2),
                                cv2.FONT_HERSHEY_PLAIN, 3, (0, 0, 255), thickness=3, lineType=cv2.LINE_AA)
                    gz_last = dangqian_lc
                    guizhen_mowei = (box_sort[j][1] + box_sort[j][3] / 2) * (2000 / 1870)
    if bq_flag == 0:
        start_lc = start_lc + 0.002
    else:
        start_lc = start_lc + 0.002-flag_bq
    cv2.imwrite(imgsave, img)
##############################