TAOB.py

import argparse
import torch
import numpy as np
from ultralytics import YOLO, RTDETR
from ultralytics.cfg import get_cfg
from ultralytics.data import build_yolo_dataset, build_dataloader
from ultralytics.data.utils import check_det_dataset
from ultralytics.utils.metrics import box_iou, bbox_iou
from ultralytics.utils.ops import xywhn2xyxy, xywh2xyxy, xyxy2xywh, non_max_suppression
from PIL import Image, ImageDraw
import torchvision.transforms.v2 as v2
import os
import copy
from tqdm import tqdm
from pathlib import Path
from utils import (
    adapt_topk1,
    load_patch,
    get_batch_idx_2_bboxes,
    get_batch_idx_2_bboxesANDcls,
    merge_bboxes,
    print_colored_arguments,
    save_patch,
    PatchApplier,
    PatchApplyToBottom,
    PatchApplyToLeft,
    PatchApplyToRight,
    PatchApplyToTop,
    AugPipeLine,
    ColorJitter,
    DropoutMatrix,
    expand_normalized_bboxes
)
from utils import set_seed
from utils.augment import RGB2Gray, RandomShuffle, Resize
from utils.loss import TVLoss
from utils.misc import CosineAnnealingLR, MultiStepLR, get_available_device
from utils import ContextLoss
# torch.autograd.set_detect_anomaly(True) # only for debugging
class YoloAttack:
    def __init__(
        self,
        model_name,
        data,
        target_label,
        k1,
        k2,
        patch,
        epoch,
        rate,
        infersize,
        verbose,
        attack_target,
        bbox_resize_factor,
        args
    ):
        self.device = "cuda:0"
        self.args = args
        self.half = True
        self.model_name = model_name
        self.model_eval_name = model_name
        print(f"load model {model_name}")
        self.model = YOLO(self.model_name)
        self.model_eval = YOLO(self.model_name)
        self.decay = 0.937
        self.data = data
        self.k1 = k1
        self.k2 = k2
        self.patch = patch
        self.target_label = target_label
        self.infersize = infersize
        self.epoch = epoch
        self.bbox_resize_factor = bbox_resize_factor
        self.attack_target = attack_target
        device = next(self.model.parameters()).device
        self.get_dataloader()
        if args.place == "left":
            self.patch_transformer = PatchApplyToLeft(position=args.placePosition)
        elif args.place == "right":
            self.patch_transformer = PatchApplyToRight(position=args.placePosition)
        elif args.place == "top":
            self.patch_transformer = PatchApplyToTop(position=args.placePosition)
        elif args.place == "bottom":
            self.patch_transformer = PatchApplyToBottom(position=args.placePosition)
        self.patch_applyer = PatchApplier()
        self.tv_loss = TVLoss()
        self.rate = rate
        self.verbose = not verbose
        self.resize = Resize()
        self.lr_scheduler = CosineAnnealingLR(eta_max=0.01, eta_min=0, T_max=self.epoch)
        self.aug_pipeline = AugPipeLine()

        # self.aug_pipeline.add(RandomShuffle(25,25), p=0.5)
        self.aug_pipeline.add(
            ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1), p=0.5
        )
        self.aug_pipeline.add(DropoutMatrix(0.1), p=0.5)
        # self.aug_pipeline.add(RGB2Gray(), p=1)
        # self.lr_scheduler = MultiStepLR(0.01, [60, 90], 0.5)
    def preprocess(self, batch):
        batch["img"] = batch["img"].to(self.device, non_blocking=True)
        batch["img"] = (
            batch["img"].half() if self.half else batch["img"].float()
        ) / 255.0
        for k in ["batch_idx", "cls", "bboxes"]:
            batch[k] = batch[k].to(self.device)

        return batch

    def get_dataloader(self):
        cfg_path = self.data
        data = check_det_dataset(cfg_path)
        args = {
            "task": "detect",
            "data": cfg_path,
            "imgsz": 640,
            "single_cls": False,
            # "model": "/home/user/dk/project/fakeobj/runs/detect/personv8s_trans/weights/best.pt",
            "rect": False,
        }
        args = get_cfg(overrides=args)
        self.batch = 8 * 2
        # self.batch = 1
        dataset_train = build_yolo_dataset(
            args, data=data, img_path=data["train"], batch=self.batch, mode="val"
        )
        dataset_val = build_yolo_dataset(
            args, data=data, img_path=data["val"], batch=self.batch, mode="val"
        )
        loader = build_dataloader
        # TODO Bug fix loader = torch.utils.data.DataLoader
        loader_train = loader(dataset_train, batch=self.batch, workers=0, shuffle=True)
        loader_test = loader(dataset_val, batch=self.batch, workers=0, shuffle=False)
        self.loader_train = loader_train
        self.loader_test = loader_test

    def get_momentum(self, grad, momentum):
        return momentum * self.decay + grad / (
            grad.abs().mean(dim=(1, 2, 3), keepdim=True)
        )

    def evalAttackSuccessRate(self, patch, save_path, save_img=False, log_file=None):
        total_attack = 1
        success_attak = 1

        total_attack_untarget = 1
        success_attack_untarget = 1
        patch = patch.cuda()
        for i, batch in enumerate(self.loader_test):
            batch = self.preprocess(batch)
            data = batch["img"]  # b,c,h,w

            if save_path is not None:
                current_path = Path(save_path)
                if not current_path.exists():
                    current_path.mkdir(parents=True, exist_ok=True)

            batch_idx_2_bboxes_cls = get_batch_idx_2_bboxesANDcls(batch)
            data_idx = []
            for batch_idx in range(data.shape[0]):
                for j, (bbox, cls_label, should_attack) in enumerate(
                    batch_idx_2_bboxes_cls[batch_idx]
                ):
                    if cls_label == self.attack_target:
                        batch_i_bbox = bbox
                        batch_i_bbox = xywhn2xyxy(batch_i_bbox, w=640, h=640)
                        _, _, W, H = xyxy2xywh(batch_i_bbox)  
                        _, _, h, w = patch.shape  
                        H = int(H)
                        W = int(W)
                        H = max(20,H)
                        W = max(20,W)
                        # HINT 在这里对patch进行增广
                        patch_1 = self.resize(
                            patch, H, W, scale_factor_arr=[self.infersize]
                        )
                        # HINT 结束增广

                        pp, patch_bbox = self.patch_transformer(
                            patch_1, (640, 640), batch_i_bbox
                        )
                        if patch_bbox is None:
                            continue
                        data_idx.append(batch_idx)
                        data[batch_idx : batch_idx + 1, :, :, :] = self.patch_applyer(
                            pp, data[batch_idx : batch_idx + 1, :, :, :]
                        )

            patched_data = data

            with torch.inference_mode():
                results2 = self.model_eval.predict(
                    patched_data, verbose=False, conf=0.5, iou=0.7
                )
                for idx, result in enumerate(results2):
                    if idx not in data_idx:
                        continue
                    im_bgr = result.plot()  # BGR-order numpy array
                    im_rgb = Image.fromarray(im_bgr)  #
                    file_name = Path(batch["im_file"][idx]).name.replace(".jpg", ".jpg")
                    if save_img and save_path:
                        im_rgb.save(current_path.joinpath(f"{file_name}"))

                    for j, (bbox, cls_label, should_attack) in enumerate(
                        batch_idx_2_bboxes_cls[idx]
                    ):
                        if cls_label == self.attack_target:  # Confirm that the attack is on stopsign.
                            batch_i_bbox = xywhn2xyxy(bbox, w=640, h=640)  # [4]
                            xyxy_pred = result.boxes.xyxy  # type: ignore
                            if len(xyxy_pred) == 0:
                                continue
                            conf_pred = result.boxes.conf  # type: ignore
                            conf_cls = result.boxes.cls  # type: ignore
                            ious = box_iou(
                                batch_i_bbox.unsqueeze(0), xyxy_pred
                            ).squeeze(
                                0
                            )  # n*4
                            is_success = False
                            is_success_untarget = True
                            # The condition for an untargeted attack to succeed is that as long as a stop sign bounding box appears, it counts as a failure.
                            for k, iou in enumerate(ious, 0):
                                if conf_cls[k] in [self.attack_target]:
                                    is_success_untarget = False
                            if is_success_untarget:
                                total_attack_untarget += 1
                                success_attack_untarget += 1
                            else:
                                total_attack_untarget += 1
                            for k, iou in enumerate(ious, 0):
                                if iou > 0.5 and conf_cls[k] in [self.attack_target]:  # stopsign
                                    # total_attack+=1
                                    is_success = False
                                    # is_success_untarget = False
                                    break
                                if iou >= 0.1 and conf_cls[k] in [
                                    self.target_label
                                ]:  
                                    total_attack += 1
                                    success_attak += 1
                                    is_success = True
                                    break
                            if not is_success:
                                total_attack += 1


        print(
            f"untarget attack success rate:{success_attack_untarget/total_attack_untarget}"
        )
        print(f"target attack success rate:{success_attak/total_attack}")
        def get_args(args):
            args_dict = vars(args)
            # Concatenate in the format "a=1,b=2,c=3"
            args_str = ",".join([f"{k}={v}" for k, v in args_dict.items()])
            return args_str
        if log_file is not None:
            ASRU = f"{(success_attack_untarget/total_attack_untarget):.4f}"
            ASRT = f"{(success_attak/total_attack):.4f}"
            with open(f"{log_file}", "+a") as f:
                f.write(f"orimodel:{self.model_name} test_model:{self.model_eval_name}\n")
                f.write(
                    f"para:beta={args.beta} infersize={self.infersize} epoch={self.epoch} bbox_resize_factor={self.bbox_resize_factor} attack_target={self.attack_target} target={self.target_label}\n"
                )
                f.write(f"args:{get_args(self.args)},ASRU={ASRU},ASRT={ASRT}\n")
                f.write(f"ASR(target):{success_attak/total_attack}\n")
                f.write(
                    f"ASR(un target):{success_attack_untarget/total_attack_untarget}\n"
                )
                # orimodel->test_model:[beta,infersize,epoch,bbox_resize_factor,attack_target,target,ASR(U),ASR(T),placePosition,targetloss]
                statment = f"{self.model_name}->{self.model_eval_name}:[{args.beta},{self.infersize},{self.epoch},{self.bbox_resize_factor},{self.attack_target},{self.target_label},{success_attack_untarget/total_attack_untarget},{success_attak/total_attack},{self.args.placePosition},{self.args.targetLoss}]\n"
                f.write(statment)
                f.write("-" * 66)
                f.write("\n")
        ASR = f"{(success_attak/total_attack):.4f}"
        new_path = current_path.parent / (current_path.name + f"ASR_{ASR}")

        if current_path.exists() and not new_path.exists():
            current_path.rename(new_path)

        return [
            np.round(success_attak / total_attack, 3),
            np.round(success_attack_untarget / total_attack_untarget, 3),
        ]

    def hook_fn2_p5(self, m, i, o):
        # print("for:",o.shape)  # Print output shape of the model
        self.hook_data[0] = o

    def hook_fn5_p4(self, m, i, o):
        # print("for:",o.shape)  # Print output shape of the model
        self.hook_data[1] = o

    def hook_fn8_p3(self, m, i, o):
        self.hook_data[2] = o

    def attack(self):
        self.hook_data = [0,0,0]
        patch = self.patch.clone()  #BCHW
        patch = patch.cuda()
        patch.requires_grad = True

        # 注册hook
        self.model.model.model[-2].register_forward_hook(hook=self.hook_fn2_p5)  # type: ignore # Register hook to get output of model
        self.model.model.model[-5].register_forward_hook(hook=self.hook_fn5_p4)  # type: ignore # Register hook to get output of model
        self.model.model.model[-8].register_forward_hook(hook=self.hook_fn8_p3) # type: ignore

        for epoch in range(1, self.epoch):
            pbar = tqdm(range(len(self.loader_train.dataset) // self.batch), disable=self.verbose)  # type: ignore
            # Attack
            moment = 0
            loss_total_show = [[], [], [], []]
            for idx, batch in enumerate(self.loader_train):
                batch = self.preprocess(batch)
                data = batch["img"]  # b,c,h,w
                patch.requires_grad_(True)

                # Obtain the original clean features.
                with torch.no_grad():
                    dummy_ = self.model(data, embed=[-1], verbose=False)
                    self.hook_data_clean = copy.deepcopy(self.hook_data)
                    batch_idxes = batch['batch_idx'].unsqueeze(1).to(dtype=int)
                    bboxes = batch['bboxes']

                    if args.randomResize:
                        self.bbox_resize_factor = np.random.randint(50, 110) / 100.
                        random_bbox_resize_factor = np.random.randint(50, 110) / 100.
                        # self.bbox_resize_factor = -1
                        # random_float = lambda: np.random.randint(50, 110) / 100.
                        random_float = lambda: random_bbox_resize_factor
                        bboxes_expand = expand_normalized_bboxes(
                            bboxes,
                            image_width=640,
                            image_height=640,
                            expand_left_ratio=random_float(),
                            expand_right_ratio=random_float(),
                            expand_top_ratio=random_float(),
                            expand_bottom_ratio=random_float(),
                        )
                    else:
                        bboxes_expand = expand_normalized_bboxes(
                            bboxes,
                            image_width=640,
                            image_height=640,
                            expand_left_ratio=self.bbox_resize_factor,
                            expand_right_ratio=self.bbox_resize_factor,
                            expand_top_ratio=self.bbox_resize_factor,
                            expand_bottom_ratio=self.bbox_resize_factor,
                        ) 
                    bboxes_expand[:, 2:] = torch.clamp(bboxes[:, 2:], max=1.0)
                    batch_bboxes = xywhn2xyxy(bboxes_expand) #xywhn
                    self.combined_bboxes = torch.cat([batch_idxes, batch_bboxes], dim=1)

                # Align the image with its bboxes and cls.
                batch_idx_2_bboxes_cls = get_batch_idx_2_bboxesANDcls(batch)
                batch_idx_2_patch_bboxes_cls = {}
                for batch_idx in range(data.shape[0]):
                    for j, (bbox, cls_label, should_attack) in enumerate(
                        batch_idx_2_bboxes_cls[batch_idx]
                    ):
                        if cls_label in [self.attack_target]:
                            batch_i_bbox = bbox
                            # batch_i_bbox = batch_idx_2_bboxes_cls[batch_idx][0][0] 
                            batch_i_bbox = xywhn2xyxy(batch_i_bbox, w=640, h=640)
                            # HINT Attach a patch to each instance of an image.
                            _, _, W, H = xyxy2xywh(batch_i_bbox)  # bbox的高宽
                            _, _, h, w = patch.shape  # patch的高宽
                            H = int(H)
                            W = int(W)
                            H = max(20,H)
                            W = max(20,W)
                            # HINT Augment the patch here
                            if self.infersize <=0.2:
                                scale_factors = [np.random.uniform(0.1, 0.3)]
                            else:
                                scale_factors = [np.random.uniform(self.infersize-0.2,self.infersize+0.2)]
                            patch_1 = self.resize(
                                    patch, H, W, scale_factor_arr=scale_factors
                                )

                            # patch_1 = self.aug_pipeline(patch_1)
                            # HINT 结束增广
                            pp, patch_bbox = self.patch_transformer(
                                patch_1, (640, 640), batch_i_bbox
                            )
                            attack_patch = True
                            if patch_bbox is None: #If it's not attached, then don't attack that target box.
                                attack_patch = False

                            if batch_idx_2_patch_bboxes_cls.get(batch_idx):
                                batch_idx_2_patch_bboxes_cls[batch_idx].append(
                                    [patch_bbox, self.target_label, attack_patch]
                                )
                            else:
                                batch_idx_2_patch_bboxes_cls[batch_idx] = [
                                    # HINT patch's bbox, targeted attack class, whether to perform a targeted attack
                                    [patch_bbox, self.target_label, attack_patch]
                                ]
                            if patch_bbox is not None:
                                data[batch_idx : batch_idx + 1, :, :, :] = (
                                    self.patch_applyer(
                                        pp, data[batch_idx : batch_idx + 1, :, :, :]
                                    )
                                )
                                batch_idx_2_bboxes_cls[batch_idx][j][2] = True
                            # break

                patched_data = data

                #!Obtain the post-processing data of the Detect header and the raw output data of the Detect header.
                results = self.model(patched_data, embed=[-1], verbose=False)
  
                # yolov5,v8,v9,v11
                if (
                    "v5" in self.model_name
                    or "v8" in self.model_name
                    or "v9" in self.model_name
                    or "11" in self.model_name,
                ):
                    # The features after processing by the Detect head are generally in the shape of [batch, (bbox, conf), 8400].
                    feature = results[0]
                    xc = feature[:, 4:, :].amax(1) > 0.0
                    feature_select = feature[:, :, xc[0]]  #! 1 * 84 * 8400

                    # HINT 12*8400*(number_class)
                    feature_select_cls = feature_select[:, 4:, :].permute((0, 2, 1))
                    # HINT 12*8400*(4)
                    feature_select_bbox = feature_select[:, :4, :].permute((0, 2, 1))
                    if feature_select.sum() == 0:
                        pbar.update(1)
                        continue
                    tokp_indices_list = []
                    tokp_indices_list_patch = []
                    for batch_idx in range(data.shape[0]):
                        for j, (bbox, cls_label, should_attack) in enumerate(
                            batch_idx_2_bboxes_cls[batch_idx]
                        ):
                            # HINT 找到需要攻击的真实框与之对应的预测框
                            if should_attack:
                                batch_i_bbox = bbox
                                batch_i_bbox = xywhn2xyxy(batch_i_bbox, w=640, h=640)
                                _, _, H, W = xyxy2xywh(batch_i_bbox)
                                # Find all relevant bounding boxes and their corresponding confidence scores (the boxes are in xyxy format)
                                # batch_i_bbox([4])
                                # feature_select_bbox([1,4,8400])
                                # HINT 注意框的格式是xyxy,yolo默认的输出的格式为xywh
                                # ious = box_iou(
                                ious = box_iou(
                                    batch_i_bbox.unsqueeze(0) / 640.0,
                                    xywh2xyxy(feature_select_bbox[batch_idx]) / 640.0,
                                )
                                # HINT use DIOU [-1,1]

                                k = max(ious.shape)
                                select_logits_cls = feature_select_cls[
                                    batch_idx : batch_idx + 1, :, :
                                ]  # B * topk * 3(3class)
                                values_conf, _ = select_logits_cls.max(2)

                                topk_values, topk_indices = torch.topk(
                                    ious * values_conf, k=k, dim=1
                                )
                                # topk_values [1,k1]

                                select_logits_cls = feature_select_cls[
                                    batch_idx : batch_idx + 1, topk_indices[0], :
                                ]  # B * topk * 3(3class)
                                values, indices = select_logits_cls.max(2)
                                select_idx1 = (
                                    indices == self.attack_target
                                )  #!important attack label[stop sign]
                                select_idx2 = values >= 0.0  # conf filter
                                cond2 = select_idx1 & select_idx2
                                cond = topk_values > 0
                                cond = cond & cond2
                                topk_values = topk_values[cond].unsqueeze(0)
                                topk_indices = topk_indices[cond].unsqueeze(0)

                                k_filter = self.k1
                                k_filter = self.args.attackTopK
                                if topk_values.shape[1] < k_filter:
                                    k_filter = topk_values.shape[1]
                                if topk_values.sum() > 0:
                                    tokp_indices_list.append(
                                        (
                                            batch_idx,
                                            topk_values[:, :k_filter],
                                            topk_indices[:, :k_filter],
                                        )
                                    )
                                bbox_patch, bbox_cls_label, bbox_should_attack = (
                                    batch_idx_2_patch_bboxes_cls[batch_idx][j]
                                )
                                if not bbox_should_attack:
                                    continue

                                bbox_instance = (
                                    batch_i_bbox.detach()
                                    .clone()
                                    .cpu()
                                    .numpy()
                                    .astype(np.int32)
                                )
                                bbox_patch = merge_bboxes(
                                    bbox_patch, bbox_instance, padding=0
                                )

                                ious_patch = box_iou(
                                    torch.FloatTensor(bbox_patch)
                                    .unsqueeze(0)
                                    .to(device=batch_i_bbox.device)
                                    / 640.0,
                                    xywh2xyxy(feature_select_bbox[batch_idx]) / 640.0,
                                )
                                _, _, bboxH, bboxW = xyxy2xywh(
                                    torch.FloatTensor(bbox_patch)
                                )

                                k2 = self.k2
                                if ious_patch.shape[1] < k2:
                                    k2 = ious_patch.shape[1]
                                topk_values_patch, topk_indices_patch = torch.topk(
                                    ious_patch, k=k2, dim=1
                                )  # shape [1,k2]
                                mask = topk_values_patch > 0
                                topk_values_patch = topk_values_patch[mask].unsqueeze(0)
                                topk_indices_patch = topk_indices_patch[mask].unsqueeze(
                                    0
                                )

                                tokp_indices_list_patch.append(
                                    (batch_idx, topk_values_patch, topk_indices_patch)
                                )
                    loss_total_show_batch = [0.0, 0.0, 0.0, 0.0]
                    total_conf = torch.Tensor([0.0]).cuda()
                    # Process each image in a batch individually.
                    # HINT Make the target box disappear
                    for batch_idx, topk_values, topk_indices in tokp_indices_list:
                        select_logits_cls = feature_select_cls[
                            batch_idx : batch_idx + 1, topk_indices[0], :
                        ]  # B * topk * 3(3class)
                        select_logits_bbox = feature_select_bbox[
                            batch_idx : batch_idx + 1, topk_indices[0], :
                        ]
                        # Weighted by confidence and IOU.
                        # HINT Using 'max' in multi-class classification can lead to issues, as it might fail to identify the intended target class for an attack.
                        values, indices = select_logits_cls.max(2)
                        select_idx1 = indices == self.attack_target  #!important attack label[stop sign]
                        select_idx2 = values >= 0.0  # conf filter
                        select_idx = select_idx1 & select_idx2
                        if not select_idx.sum():
                            continue
                        max_k = select_idx.sum()
                        loss_topk = torch.topk(
                            topk_values[select_idx], max_k // 1, 0
                        ).values
                        loss_tk = torch.mean(loss_topk)
                        total_conf += loss_tk

                    loss_confs_target = torch.FloatTensor([0.0]).cuda()
                    loss_ious = torch.FloatTensor([0.0]).cuda()
                    loss_combine = torch.FloatTensor([0.0]).cuda()  # todo
                    for (
                        batch_idx,
                        topk_values_patch,
                        topk_indices_patch,
                    ) in tokp_indices_list_patch:
                        select_logits_cls = feature_select_cls[
                            batch_idx : batch_idx + 1, topk_indices_patch[0], :
                        ]  # B * topk * 3(3class)
                        # HINT target attack
                        target = torch.zeros_like(select_logits_cls).cuda()
                        target[:, :, self.target_label] = 1

                        loss_bce = torch.nn.BCELoss()(
                            select_logits_cls,  # 1 * N * 80
                            target,
                        )
                        loss_iou = (1 - topk_values_patch).mean()
                        loss_ious += loss_iou
                        loss_confs_target += loss_bce
                    if total_conf < 1e-10:
                        pbar.update(1)
                        continue
                    total_conf *= 1
                    loss_ious *= 1
                    loss_confs_target *= 2
                    loss_tv = self.tv_loss(patch) * 2

                    loss_total_show_batch[0] = total_conf.item()
                    loss_total_show_batch[1] = loss_ious.item()
                    loss_total_show_batch[2] = loss_confs_target.item()
                    loss_total_show_batch[3] = loss_tv.item()
                    if args.useAllFeatureLoss:
                        ctx_loss1 = torch.nn.CosineSimilarity(dim=1, eps=1e-6)(self.hook_data_clean[0], self.hook_data[0]).mean()+1
                        ctx_loss2 = torch.nn.CosineSimilarity(dim=1, eps=1e-6)(self.hook_data_clean[1], self.hook_data[1]).mean()+1
                        ctx_loss3 = torch.nn.CosineSimilarity(dim=1, eps=1e-6)(self.hook_data_clean[2], self.hook_data[2]).mean()+1
                    else:
                        ctx_loss1 = ContextLoss(self.combined_bboxes)(self.hook_data_clean[0], self.hook_data[0])
                        ctx_loss2 = ContextLoss(self.combined_bboxes)(self.hook_data_clean[1], self.hook_data[1])
                        ctx_loss3 = ContextLoss(self.combined_bboxes)(self.hook_data_clean[2], self.hook_data[2])
                    ctx_loss = ctx_loss1 + ctx_loss2 + ctx_loss3
                    target_loss = loss_ious + loss_confs_target
                    loss = (
                        total_conf
                        + args.beta * ctx_loss
                        + args.targetLoss * target_loss
                        + args.tvloss * loss_tv
                    )
   
                elif "v10" in self.model_name:
                    # The features processed by the Detect head are generally in the format [x, y, w, h, maxclassprob, class_index].
                    feature = results[0]
                elif "detr" in self.model_name:
                    pass
                else:
                    raise NotImplementedError
                grad = torch.autograd.grad(
                    loss, patch, retain_graph=False, create_graph=False
                )[0]
                # weight decay
                # grad = grad + 1e-6 * patch
                moment = self.get_momentum(grad, moment)
                # update perturbation
                alpha = self.lr_scheduler.get_lr(epoch)
                patch = patch - alpha * torch.sign(moment)

                # clip
                # patch = torch.clamp(patch, 0, 1).detach()
                patch = torch.clamp(patch, 0, 1)

                for i in range(len(loss_total_show_batch)):
                    loss_total_show_batch[i] = float(np.round(loss_total_show_batch[i], 3))
                    loss_total_show[i].append(loss_total_show_batch[i])
                pbar.set_description(
                    f"epoch={epoch} i={idx} ctx_loss={ctx_loss.item()} loss_batch={loss_total_show_batch} loss_total={[np.round(np.sum(x),2) for x in loss_total_show]}"
                )
                pbar.update(1)

            if epoch % 2000 == 0 and epoch != 0:
                save_patch(patch, epoch, self.model_name)
            self.patch = patch.detach().clone()


def parse_arguments():
    parser = argparse.ArgumentParser(description="Parse command line arguments")

    parser.add_argument("--data_path", type=str, default="/home/user/dk/dengkang_data/datasets/cocostopsign_refine_truelabel/cocostopsign_true.yaml", help="Path to the data")
    parser.add_argument(
        "--model_path", type=str, nargs="*",default=["yolov8s.pt"], help="Path to the model"
    )
    parser.add_argument("--target_label", nargs="*",type=int, default=[0], help="target_label")
    parser.add_argument(
        "--save_path", type=str, default="save", help="Path to save results"
    )
    parser.add_argument(
        "--devices", type=str, default="3", help='Devices to use (e.g., "cuda:0")'
    )    
    parser.add_argument(
        "--place", type=str, default="bottom", help='Place position of the patch'
    )
    parser.add_argument(
        "--placePosition", type=int, default=0, help='Place position of the patch'
    )
    parser.add_argument(
        "--patchAspectRatio", type=int, default=2, help=''
    )
    parser.add_argument(
        "--beta", type=float, default=0.05, help=''
    )
    parser.add_argument(
        "--targetLoss", type=float, default=2, help=''
    )
    parser.add_argument(
        "--tvloss", type=float, default=3, help=''
    )
    parser.add_argument(
        "--attackTopK", type=int, default=10, help=''
    )
    parser.add_argument("--infersize", type=float, nargs="*",default=[1.2], help="infersize")
    parser.add_argument("--save", action='store_true', default=True, help="Enable save")
    parser.add_argument("--randomResize", action='store_true', default=True, help="Enable random resize")
    parser.add_argument("--epoch", type=int, nargs="*",default=[10], help="epoch")
    parser.add_argument("--attack_target", type=int,default=11, help="attack_target,default 11[stop sign]")
    parser.add_argument("--verbose", action='store_true', default=False, help="verbose")
    parser.add_argument("--bbox_resize_factor_list",default=[-1], type=float, nargs="*", help="List of integers for k2")
    parser.add_argument("--useAllFeatureLoss", action='store_true', default=False, help="List of integers for k2")
    parser.add_argument(
        "--logfile",
        type=str,
        default="test.txt",
        help="Path to log file",
    )

    args = parser.parse_args()

    return args


if __name__ == "__main__":
    available_device = get_available_device()
    os.environ["CUDA_VISIBLE_DEVICES"] = available_device.split(":")[1]

    args = parse_arguments()

    if args.patchAspectRatio == 2:
        base_size_w = 200
        base_size_h = 100
    elif args.patchAspectRatio == 1:
        base_size_w = 100
        base_size_h = 100
    elif args.patchAspectRatio == -2:
        base_size_w = 100
        base_size_h = 200
    patch = load_patch("./patch_home/patch_hide.jpg", (base_size_w, base_size_h))  # hint WH
    print(f"patch shape:{patch.shape}")

    k1_list = [1] 
    k2_list = [13] 
    # args.verbose = False
   
    print_colored_arguments(args)
    for model in args.model_path:
        for epoch in args.epoch:
            for infersize in args.infersize:
                for bbox_resize_factor in args.bbox_resize_factor_list:
                    for target_label in args.target_label:
                        for k1 in k1_list:
                            for k2 in k2_list:
                                attack = YoloAttack(
                                    args=args,
                                    model_name=model,
                                    data=args.data_path,
                                    target_label=target_label,
                                    k1=k1,  
                                    k2=k2, 
                                    patch=patch,
                                    epoch=epoch,
                                    rate=None,
                                    infersize=infersize,
                                    verbose=args.verbose,
                                    attack_target=args.attack_target,
                                    bbox_resize_factor=bbox_resize_factor,
                                )
                                set_seed(114514)
                                attack.attack()
                                if args.save_path:
                                    save_path=f"./{args.save_path}/0711_{args.attack_target}to{target_label}_{model[:-3]}"
                                else:
                                    save_path=f"./save/0711_{args.attack_target}to{target_label}_{model[:-3]}",
                                acc = attack.evalAttackSuccessRate(
                                    attack.patch,
                                    save_path=save_path,
                                    save_img=args.save,
                                    log_file=args.logfile,
                                )
                                if args.save:
                                    save_patch(attack.patch, f"0417final_phy_{epoch}_{acc}", "yolov8s")