import argparse

import json

import os

from pathlib import Path

from threading import Thread

import numpy as np

import torch

import yaml

from tqdm import tqdm

from models.experimental import attempt_load

from utils.datasets import create_dataloader

from utils.general import coco80_to_coco91_class, check_dataset, check_file, check_img_size, check_requirements, \

    box_iou, non_max_suppression, scale_coords, xyxy2xywh, xywh2xyxy, set_logging, increment_path, colorstr

from utils.metrics import ap_per_class, ConfusionMatrix, batch_pix_accuracy,batch_pix_accuracy_class, batch_intersection_union  # 后两个新增分割

from utils.plots import plot_images, output_to_target, plot_study_txt

from utils.torch_utils import select_device, time_synchronized

import torch.nn.functional as F

import cv2

from models.yolo import Model

from utils.loss import ComputeLoss, SegmentationLosses, SegFocalLoss, OhemCELoss, ProbOhemCrossEntropy2d,PoseLoss

import SegmentationDataset

"""

test_custom.py与test.py的区别仅在加载器上从Cityscapes改成了Custom

新版训练测试(loader的mode为"testval")可以把验证集长边resize到base-size输入到网络, 但mask仍然是原图尺寸, 以下代码自动把网络输出双线性插值到原图算指标

调用示例:

python test.py --data cityscapes_det.yaml --segdata ./data/citys --weights ./best.pt --img-size 640 --base-size 640

即相比原版yolov5多 --segdata 和 --base-size两个参数

"""

Cityscapes_COLORMAP = [

    # [128, 64, 128],

    [0,0,0],

    [244, 35, 232],

    [0, 0, 192],

    [70, 70, 70],

    [102, 102, 0],

    [190, 153, 153],

    [153, 153, 153],

    [250, 170, 30],

    [220, 220, 0],

    [107, 142, 35],

    [152, 251, 152],

    [0, 130, 180],

    [220, 20, 60],

    [255, 0, 0],

    [0, 0, 70],

    [0, 60, 100],

    [0, 80, 100],

    [0, 0, 230],

    [119, 11, 32],

]

def attempt_load(weights, map_location=None):

    # Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a

    model = Ensemble()

    for w in weights if isinstance(weights, list) else [weights]:

        attempt_download(w)

        ckpt = torch.load(w, map_location=map_location)  # load

        # model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval())  # FP32 model

        model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval())  # FP32 model

    # Compatibility updates

    for m in model.modules():

        if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:

            m.inplace = True  # pytorch 1.7.0 compatibility

        elif type(m) is Conv:

            m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatibility

    if len(model) == 1:

        return model[-1]  # return model

    else:

        print('Ensemble created with %s\n' % weights)

        for k in ['names', 'stride']:

            setattr(model, k, getattr(model[-1], k))

        return model  # return ensemble

def label2image(pred, COLORMAP=Cityscapes_COLORMAP):

    colormap = np.array(COLORMAP, dtype='uint8')

    X = pred.astype('int32')

    return colormap[X, :]

def seg_validation(model, n_segcls, valloader, device, half_precision=True):

    # Fast test during the training

    def eval_batch(model, image, target, half):

 

        # torch.Size([64, 2, 544, 480])

        # torch.Size([64, 14, 544, 480])

        # outputs = gather(outputs, 0, dim=0)

        if n_segcls == 2:

            pred = outputs[1][0]  # 1是分割

        if n_segcls == 14:

            pred = outputs[1][1]

        target = target.to(device, non_blocking=True)

        pred = F.interpolate(pred, (target.shape[1], target.shape[2]), mode='bilinear', align_corners=True)

        maskpred  = F.interpolate(pred, (target.shape[1], target.shape[2]), mode='bilinear', align_corners=True)[0]

        mask = label2image(maskpred.max(axis=0)[1].cpu().numpy(), Cityscapes_COLORMAP)[:, :, ::-1]

        correct, labeled,total_pixel = batch_pix_accuracy_class(pred.data, target,n_segcls)

        inter, union = batch_intersection_union(pred.data, target, n_segcls)

        return correct, labeled, inter, union,total_pixel,mask

    half = device.type != 'cpu' and half_precision  # half precision only supported on CUDA

    if half:

        model.half()

    model.eval()

    total_inter, total_union, total_correct, total_label = 0, 0, 0, 0

    total_c = []

    tbar = tqdm(valloader, desc='\r')

    numtime = 1

    for i, (image, target) in enumerate(tbar):

        image = image.to(device, non_blocking=True)

        image = image.half() if half else image.float()

        with torch.no_grad():

            correct, labeled, inter, union ,total_pixel,mask= eval_batch(model, image, target, half)

           

            #将测试集中的图片进行可视化保存

            # imagename = (image[0].cpu().numpy().transpose(1,2,0) * 255).astype('uint8')

            # dst_fsd = cv2.addWeighted(mask, 0.4, imagename, 0.6, 0)

            # cv2.imwrite('/ai/zhdata/multiyolov5_point_v2/image/FSD/train_v30_zh_20240805_FSD/1/' + str(i) + '.jpg', dst_fsd)

            #将测试集中的图片进行可视化保存

        total_correct += correct

        total_label += labeled

        total_inter += inter

        total_union += union

        total_c += total_pixel

        pixAcc = 1.0 * total_correct / (np.spacing(1) + total_label)

        total_c = np.array(total_c)

        pixAcc_class = 1.0 * total_c[:,0] / (np.spacing(1) + total_c[:,1])

        IoU = 1.0 * total_inter / (np.spacing(1) + total_union)

        numtime += 1

        mIoU = IoU.mean()

        # tbar.set_description(

        #     'pixAcc: %.3f, mIoU: %.3f' % (pixAcc, mIoU))

       

        if i >= len(tbar) - 1:

            print('mIoU=    ',str(mIoU))

            for j in range(len(IoU)):

                print("class:" + str(j) + "result------" +str(IoU[j]) + "--ACC----" + str(pixAcc_class[j]))

        if len(IoU) == 14:

            tbar.set_description(

            'Acc: %.2f, mIoU: %.2f, 0: %.2f, 1: %.2f, 2: %.2f, 3: %.2f, 4: %.2f, 5: %.2f, 6: %.2f, 7: %.2f, 8: %.2f, 9: %.2f, 10: %.2f, 11: %.2f, 12: %.2f, 13: %.2f ' % (pixAcc, mIoU, IoU[0], IoU[1], IoU[2], IoU[3], IoU[4] \

                , IoU[5], IoU[6], IoU[7], IoU[8], IoU[9], IoU[10], IoU[11], IoU[12], IoU[13]))

            # tbar.set_description(

            # 'Acc: %.2f, Acc: %.2f, 0: %.2f, 1: %.2f, 2: %.2f, 3: %.2f, 4: %.2f, 5: %.2f, 6: %.2f, 7: %.2f, 8: %.2f, 9: %.2f, 10: %.2f, 11: %.2f, 12: %.2f, 13: %.2f ' % (pixAcc, mIoU, IoU[0], IoU[1], IoU[2], IoU[3], IoU[4] \

            #     , IoU[5], IoU[6], IoU[7], IoU[8], IoU[9], IoU[10], IoU[11], IoU[12], IoU[13]))          

            # tbar.set_description(

            # 'Acc: %.2f, mIoU: %.2f, 0: %.2f, 1: %.2f, 2: %.2f, 3: %.2f, 4: %.2f, 5: %.2f, 6: %.2f ' % (pixAcc, mIoU, IoU[0], IoU[1], IoU[2], IoU[3], IoU[4] \

            #     , IoU[5], IoU[6]))

        if len(IoU) == 2:

            tbar.set_description(

            'pixAcc: %.3f, mIoU: %.3f, class0: %.3f, class1: %.3f' % (pixAcc, mIoU, IoU[0], IoU[1]))

    return mIoU,IoU

 

def segtest(weights, root="data/citys", batch_size=16, half_precision=True, n_segcls=19, base_size=2048):  # 会使用原始尺寸测, 未考虑尺寸对不齐, 图片尺寸应为32倍数

    device = select_device(opt.device, batch_size=batch_size)

    model = attempt_load(weights, map_location=device)  # load FP32 model

    testvalloader = SegmentationDataset.get_rm_loader(root, batch_size=batch_size, split="val", mode="val", workers=8, base_size=base_size) #get_custom_loader

    # testvalloader = SegmentationDataset.get_citys_loader(root, batch_size=batch_size, split="val", mode="val", workers=4, base_size=1024, crop_size=1024)

    seg_validation(model, n_segcls, testvalloader, device, half_precision)

   

def segtest_fsd(weights, root="data/citys", batch_size=16, half_precision=True, n_segcls=19, base_size=2048):  # 会使用原始尺寸测, 未考虑尺寸对不齐, 图片尺寸应为32倍数

    device = select_device(opt.device, batch_size=batch_size)

    model = attempt_load(weights, map_location=device)  # load FP32 model

    testvalloader = SegmentationDataset.get_custom_loader(root, batch_size=batch_size, split="val", mode="val", workers=8, base_size=base_size) #get_custom_loader

    # testvalloader = SegmentationDataset.get_citys_loader(root, batch_size=batch_size, split="val", mode="val", workers=4, base_size=1024, crop_size=1024)

    seg_validation(model, n_segcls, testvalloader, device, half_precision)

 

def test(data,

         weights=None,

         batch_size=64,

         imgsz=640,

         conf_thres=0.001,

         iou_thres=0.6,  # for NMS

         save_json=False,

         single_cls=False,

         augment=False,

         verbose=False,

         model=None,

         dataloader=None,

         save_dir=Path(''),  # for saving images

         save_txt=False,  # for auto-labelling

         save_hybrid=False,  # for hybrid auto-labelling

         save_conf=False,  # save auto-label confidences

         plots=False,

         wandb_logger=None,

         compute_loss=False,

         half_precision=True,

         is_coco=False):

    # Initialize/load model and set device

    training = model is not None

    if training:  # called by train.py

        device = next(model.parameters()).device  # get model device

    else:  # called directly

        set_logging()

        device = select_device(opt.device, batch_size=batch_size)

        # Directories

        save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))  # increment run

        (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir

        # Load model

        model = attempt_load(weights, map_location=device)  # load FP32 model

        # model = Model(opt.cfg, ch=3, nc=3).to('cpu')  # create

        gs = max(int(model.stride.max()), 32)  # grid size (max stride)

        imgsz = check_img_size(imgsz, s=gs)  # check img_size

        # Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99

        # if device.type != 'cpu' and torch.cuda.device_count() > 1:

        #     model = nn.DataParallel(model)

    # Half

    half = device.type != 'cpu' and half_precision  # half precision only supported on CUDA

    # half = False

    if half:

        model.half()

    # Configure

    model.eval()

    if isinstance(data, str):

        is_coco = data.endswith('coco.yaml')

        with open(data) as f:

            data = yaml.load(f, Loader=yaml.SafeLoader)

    check_dataset(data)  # check

    nc = 1 if single_cls else int(data['nc'])  # number of classes

    iouv = torch.linspace(0.5, 0.95, 10).to(device)  # iou vector for mAP@0.5:0.95

    niou = iouv.numel()

    # Logging

    log_imgs = 0

    if wandb_logger and wandb_logger.wandb:

        log_imgs = min(wandb_logger.log_imgs, 100)

    # Dataloader

    if not training:

        if device.type != 'cpu':

            model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))  # run once

        task = opt.task if opt.task in ('train', 'val', 'test') else 'val'  # path to train/val/test images

        dataloader = create_dataloader(data[task], imgsz, batch_size, gs, opt, pad=0.5, rect=True,

                                       prefix=colorstr(f'{task}: '))[0]

#/ai/DataSets/OD_FSD_zh/psd_v2.0/data/txt_dataset/ann_zh/test_img_list_v5.txt

#data[task]

#/ai/DataSets/OD_FSD_zh/psd_v2.0/data/83_psd_20240723_占用属性优化/alldata/test_img_list.txt

    seen = 0

    confusion_matrix = ConfusionMatrix(nc=nc)

    names = {k: v for k, v in enumerate(model.names if hasattr(model, 'names') else model.module.names)}

    coco91class = coco80_to_coco91_class()

    s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Labels', 'P', 'R', 'mAP@.5', 'mAP@.5:.95')

    p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0.

    loss = torch.zeros(3, device=device)

    jdict, stats, ap, ap_class, wandb_images = [], [], [], [], []

    numtotal = 0

    numocc,numvip,numwoman,numdisabled,numcharging = 0,0,0,0,0

    for batch_i, (img, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s)):

        img = img.to(device, non_blocking=True)

        img = img.half() if half else img.float()  # uint8 to fp16/32

        img /= 255.0  # 0 - 255 to 0.0 - 1.0

        targets = targets.to(device)

        nb, _, height, width = img.shape  # batch size, channels, height, width

        with torch.no_grad():

            # Run model

            t = time_synchronized()

            out, train_out = model(img, augment=augment)[0]  # inference and training outputs 修改[0]新模型输出[0]是检测

            t0 += time_synchronized() - t

            # Compute loss

            if compute_loss:

                # Hyperparameters 配置超参数

                with open(opt.hyp) as f:

                    hyp = yaml.load(f, Loader=yaml.SafeLoader)  # load hyps

                model.hyp = hyp

                model.gr = 1.0

                compute_loss = PoseLoss(model)

                loss += compute_loss([x.float() for x in train_out], targets)[1][:3]  # box, obj, cls

            # Run NMS

            targets[:, 2:6] *= torch.Tensor([width, height, width, height]).to(device)  # to pixels

            lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else []  # for autolabelling

            t = time_synchronized()

            out = non_max_suppression(out, conf_thres=conf_thres, iou_thres=iou_thres, labels=lb, multi_label=False)

            t1 += time_synchronized() - t

        # Statistics per image

        for si, pred in enumerate(out):

            perd_slot = pred[:, 14:]

            pred = torch.cat((pred[:, :5], pred[:, 13:14]), 1)

            labels = targets[targets[:, 0] == si, 1:]

            nl = len(labels)

            tcls = labels[:, 0].tolist() if nl else []  # target class

            tocc = labels[:, 13].tolist() if nl else []

            tvip = labels[:, 14].tolist() if nl else []

            twoman = labels[:, 15].tolist() if nl else []

            tdisable = labels[:, 16].tolist() if nl else []

            tcharging = labels[:, 17].tolist() if nl else []

            path = Path(paths[si])

            seen += 1

            if len(pred) == 0:

                if nl:

                    stats.append((torch.zeros(0, niou, dtype=torch.bool), torch.Tensor(), torch.Tensor(), tcls))

                continue

            # Predictions

            predn = pred.clone()

            scale_coords(img[si].shape[1:], predn[:, :4], shapes[si][0], shapes[si][1])  # native-space pred

            # Append to text file

            if save_txt:

                gn = torch.tensor(shapes[si][0])[[1, 0, 1, 0]]  # normalization gain whwh

                for *xyxy, conf, cls in predn.tolist():

                    xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh

                    line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format

                    with open(save_dir / 'labels' / (path.stem + '.txt'), 'a') as f:

                        f.write(('%g ' * len(line)).rstrip() % line + '\n')

            # W&B logging - Media Panel Plots

            if len(wandb_images) < log_imgs and wandb_logger.current_epoch > 0:  # Check for test operation

                if wandb_logger.current_epoch % wandb_logger.bbox_interval == 0:

                    box_data = [{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},

                                 "class_id": int(cls),

                                 "box_caption": "%s %.3f" % (names[cls], conf),

                                 "scores": {"class_score": conf},

                                 "domain": "pixel"} for *xyxy, conf, cls in pred.tolist()]

                    boxes = {"predictions": {"box_data": box_data, "class_labels": names}}  # inference-space

                    wandb_images.append(wandb_logger.wandb.Image(img[si], boxes=boxes, caption=path.name))

            wandb_logger.log_training_progress(predn, path, names) if wandb_logger and wandb_logger.wandb_run else None

            # Append to pycocotools JSON dictionary

            if save_json:

                # [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ...

                image_id = int(path.stem) if path.stem.isnumeric() else path.stem

                box = xyxy2xywh(predn[:, :4])  # xywh

                box[:, :2] -= box[:, 2:] / 2  # xy center to top-left corner

                for p, b in zip(pred.tolist(), box.tolist()):

                    jdict.append({'image_id': image_id,

                                  'category_id': coco91class[int(p[5])] if is_coco else int(p[5]),

                                  'bbox': [round(x, 3) for x in b],

                                  'score': round(p[4], 5)})

            # Assign all predictions as incorrect

            correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool, device=device)

            if nl:

                detected = []  # target indices

                tcls_tensor = labels[:, 0]

                # target boxes

                tbox = xywh2xyxy(labels[:, 1:5])

                scale_coords(img[si].shape[1:], tbox, shapes[si][0], shapes[si][1])  # native-space labels

                if plots:

                    confusion_matrix.process_batch(predn, torch.cat((labels[:, 0:1], tbox), 1))

                # Per target class

                for cls in torch.unique(tcls_tensor):

                    ti = (cls == tcls_tensor).nonzero(as_tuple=False).view(-1)  # prediction indices

                    pi = (cls == pred[:, 5]).nonzero(as_tuple=False).view(-1)  # target indices

                    # Search for detections

                    if pi.shape[0]:

                        # Prediction to target ious

                        ious, i = box_iou(predn[pi, :4], tbox[ti]).max(1)  # best ious, indices

                        # Append detections

                        detected_set = set()

                        for j in (ious > iouv[0]).nonzero(as_tuple=False):

                            d = ti[i[j]]  # detected target

                            if d.item() not in detected_set:

                                detected_set.add(d.item())

                                detected.append(d)

                                correct[pi[j]] = ious[j] > iouv  # iou_thres is 1xn                                                                                                                                        

                                if len(detected) == nl:  # all targets already located in image

                                    break

            # Append statistics (correct, conf, pcls, tcls)

            stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls))

        # Plot images

        if plots and batch_i < 3:

            f = save_dir / f'test_batch{batch_i}_labels.jpg'  # labels

            Thread(target=plot_images, args=(img, targets, paths, f, names), daemon=True).start()

            f = save_dir / f'test_batch{batch_i}_pred.jpg'  # predictions

            Thread(target=plot_images, args=(img, output_to_target(out), paths, f, names), daemon=True).start()

    # Compute statistics

    stats = [np.concatenate(x, 0) for x in zip(*stats)]  # to numpy

    if len(stats) and stats[0].any():

        p, r, ap, f1, ap_class = ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names)

        ap50, ap = ap[:, 0], ap.mean(1)  # AP@0.5, AP@0.5:0.95

        mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()

        nt = np.bincount(stats[3].astype(np.int64), minlength=nc)  # number of targets per class

    else:

        nt = torch.zeros(1)

    # Print results

    pf = '%20s' + '%12i' * 2 + '%12.3g' * 4  # print format

    print(pf % ('all', seen, nt.sum(), mp, mr, map50, map))

    # Print results per class

    if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats):

        for i, c in enumerate(ap_class):

            print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))

    # Print speeds

    t = tuple(x / seen * 1E3 for x in (t0, t1, t0 + t1)) + (imgsz, imgsz, batch_size)  # tuple

    if not training:

        print('Speed: %.1f/%.1f/%.1f ms inference/NMS/total per %gx%g image at batch-size %g' % t)

    # Plots

    if plots:

        confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))

        if wandb_logger and wandb_logger.wandb:

            val_batches = [wandb_logger.wandb.Image(str(f), caption=f.name) for f in sorted(save_dir.glob('test*.jpg'))]

            wandb_logger.log({"Validation": val_batches})

    if wandb_images:

        wandb_logger.log({"Bounding Box Debugger/Images": wandb_images})

    # Save JSON

    if save_json and len(jdict):

        w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else ''  # weights

        anno_json = '../coco/annotations/instances_val2017.json'  # annotations json

        pred_json = str(save_dir / f"{w}_predictions.json")  # predictions json

        print('\nEvaluating pycocotools mAP... saving %s...' % pred_json)

        with open(pred_json, 'w') as f:

            json.dump(jdict, f)

        try:  # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb

            from pycocotools.coco import COCO

            from pycocotools.cocoeval import COCOeval

            anno = COCO(anno_json)  # init annotations api

            pred = anno.loadRes(pred_json)  # init predictions api

            eval = COCOeval(anno, pred, 'bbox')

            if is_coco:

                eval.params.imgIds = [int(Path(x).stem) for x in dataloader.dataset.img_files]  # image IDs to evaluate

            eval.evaluate()

            eval.accumulate()

            eval.summarize()

            map, map50 = eval.stats[:2]  # update results (mAP@0.5:0.95, mAP@0.5)

        except Exception as e:

            print(f'pycocotools unable to run: {e}')

    # Return results

    model.float()  # for training

    if not training:

        s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''

        print(f"Results saved to {save_dir}{s}")

    maps = np.zeros(nc) + map

    for i, c in enumerate(ap_class):

        maps[c] = ap[i]

    return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, t

 

if __name__ == '__main__':

    parser = argparse.ArgumentParser(prog='test.py')

    parser.add_argument('--cfg', type=str, default='models/yolov5s_custom_seg.yaml', help='model.yaml path')#

    parser.add_argument('--hyp', type=str, default='data/hyp.scratch.yaml', help='hyperparameters path')

    #runs/train/exp52/weights/exp52_last_114_v2.0.21_20241129.pt

    #runs/train/exp4/weights/best.pt

    parser.add_argument('--weights', nargs='+', type=str, default="/ai/ypli/multiyolov5_point_v2/runs/train/exp52/weights_old/exp52_last_115_20250311_v2.0.25.pt", help='model.pt path(s)') #'runs/train/exp51/weights/last.pt'

    parser.add_argument('--data', type=str, default='data/custom.yaml', help='*.data path')

    path = ["/ai/DataSets/TopViewMultiTaskPerc_xmlin/cvat/test_fsd.txt"]

        # path = ["/ai/DataSets/TopViewMultiTaskPerc_xmlin/freeSpace/annotations/ann_zh/val_cpp_v29_20250806_gt.txt"]

    # path = ['/ai/DataSets/TopViewMultiTaskPerc_xmlin/freeSpace/annotations/ann_zh/val_zh_v23_20250312_gt.txt']

    # path_rm = ['/ai/DataSets/TopViewMultiTaskPerc_xmlin/roadmarking/annotations/ann_zh/val_v24_fullLabel0103_gt_v1.txt']

    path_rm = ["/ai/DataSets/TopViewMultiTaskPerc_xmlin/roadmarking/annotations/ann_zh/new_dataset/val_v36_20250806_gt.txt"]

        # path_rm = ["/ai/DataSets/TopViewMultiTaskPerc_xmlin/roadmarking/annotations/ann_zh/new_dataset/val_v36_20250806_gt.txt"]

    # path_rm = ['/ai/DataSets/TopViewMultiTaskPerc_xmlin/roadmarking/annotations/ann_zh/new_dataset/val_v30_20250305_gt.txt']

    # path_rm = ['/ai/DataSets/TopViewMultiTaskPerc_xmlin/roadmarking/annotations/20240411_生态园and保定对向路沿/2task_rm_20240405_wuluhong/2task_rm_20240405_wuluhong_gt.txt']

    parser.add_argument('--segdata', default=path, help='root path of segmentation data')#type=list,

    parser.add_argument('--segdata_rm', default=path_rm, help='root path of segmentation data')#type=list,

    parser.add_argument('--batch-size', type=int, default=8, help='size of each image batch')

    parser.add_argument('--img-size', type=int, default=544, help='inference size (pixels)')

    parser.add_argument('--base-size', type=int, default=544, help='long side of segtest image you want to input network')

    parser.add_argument('--conf-thres', type=float, default=0.001, help='object confidence threshold')

    parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')

    parser.add_argument('--task', default='val', help='train, val, test, speed or study')

    parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')

    parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')

    parser.add_argument('--augment', action='store_true', help='augmented inference')

    parser.add_argument('--verbose', action='store_true', help='report mAP by class')

    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')

    parser.add_argument('--save-hybrid', action='store_true', help='save label+prediction hybrid results to *.txt')

    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')

    parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file')

    parser.add_argument('--project', default='runs/test', help='save to project/name')

    parser.add_argument('--name', default='exp', help='save to project/name')

    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')

    opt = parser.parse_args()

    opt.save_json |= opt.data.endswith('coco.yaml')

    opt.data = check_file(opt.data)  # check file

    print(opt)

    check_requirements()

    # if opt.task in ('train', 'val', 'test'):  # run normally

    #     print(1111)

    #     test(opt.data,

    #          opt.weights,

    #          opt.batch_size,

    #          opt.img_size,

    #          opt.conf_thres,

    #          opt.iou_thres,

    #          opt.save_json,

    #          opt.single_cls,

    #          opt.augment,

    #          opt.verbose,

    #          save_txt=opt.save_txt | opt.save_hybrid,

    #          save_hybrid=opt.save_hybrid,

    #          save_conf=opt.save_conf,

    #          )

    # elif opt.task == 'speed':  # speed benchmarks

    #     for w in opt.weights:

    #         test(opt.data, w, opt.batch_size, opt.img_size, 0.25, 0.45, save_json=False, plots=False)

    # elif opt.task == 'study':  # run over a range of settings and save/plot

    #     # python test.py --task study --data coco.yaml --iou 0.7 --weights yolov5s.pt yolov5m.pt yolov5l.pt yolov5x.pt

    #     x = list(range(256, 1536 + 128, 128))  # x axis (image sizes)

    #     for w in opt.weights:

    #         f = f'study_{Path(opt.data).stem}_{Path(w).stem}.txt'  # filename to save to

    #         y = []  # y axis

    #         for i in x:  # img-size

    #             print(f'\nRunning {f} point {i}...')

    #             r, _, t = test(opt.data, w, opt.batch_size, i, opt.conf_thres, opt.iou_thres, opt.save_json,

    #                            plots=False)

    #             y.append(r + t)  # results and times

    #         np.savetxt(f, y, fmt='%10.4g')  # save

    #     os.system('zip -r study.zip study_*.txt')

    #     plot_study_txt(x=x)  # plot

    segtest_fsd(root=opt.segdata, weights=opt.weights, batch_size=64, n_segcls=2, base_size=opt.base_size)  # 19 for cityscapes

    segtest(root=opt.segdata_rm, weights=opt.weights, batch_size=64, n_segcls=14, base_size=opt.base_size)  # 19 for cityscapes