[CVPR 2023] Understanding the Robustness of 3D Object Detection With Bird's-Eye-View Representations in Autonomous Driving
import mmcv
import torch
import numpy as np
import PIL.Image as Image
import torchvision.transforms as transforms
import torch.nn.functional as F
import torchvision
from torchvision.utils import save_image
import cv2
import time
import os
from extend.custom_func import *
from extend_common.img_check import img_diff_print
from extend_common.time_counter import time_counter

def single_gpu_test(model, data_loader, 
                    scattered_result_prefix=None, 
                    eps255=None
                    ):

    model.eval()
    results = []
    dataset = data_loader.dataset
    
    scattered_result_dir = scattered_result_prefix + 'eps'+eps255
    eps255 = float(eps255)
    

    
    os.makedirs(scattered_result_dir, exist_ok=True)
    device = model.src_device_obj
    
    time_test_flag = False
    prog_bar = mmcv.ProgressBar(len(dataset))
    for data_i, data_out in enumerate(data_loader):
        
        results_name = str(data_i) + '.pkl'
        scattered_result_path = os.path.join(scattered_result_dir, results_name)
        
        if os.path.exists(scattered_result_path):
            print(scattered_result_path, 'exists! pass!')
            prog_bar.update()
            continue
        else:
            mmcv.dump(' ', scattered_result_path)
            
        
        last_time = time.time()

        # 1. data processing(customed)
        data_out = custom_data_preprocess(data_out)
        img_metas, img_path_list, img_org_np, img_processed, gt_labels_3d = custom_data_work(data_out)
        img_tensor_ncam = custom_img_read_from_img_org(img_org_np, device)
        last_time = time_counter(last_time, 'data load', time_test_flag)


        # FGSM all region attack, do not need patch info
        
        imgs_noisy = img_tensor_ncam.clone().detach()
        orig_imgs_input = img_tensor_ncam.clone().detach()
        
        eps01 = eps255 / 255.
        lr = eps01
        max_step = 1

        # to avoid no_gt
        if gt_labels_3d.shape[0] != 0:
            for step in range(max_step):
                imgs_noisy = imgs_noisy.clone().detach().requires_grad_(True)

                ############ resize norm pad
                image_ready = custom_differentiable_transform(
                        img_tensor_rgb_6chw_0to1=imgs_noisy,
                        img_metas=img_metas,
                    )
                last_time = time_counter(last_time, 'img rsnmpd', time_test_flag)

                if image_ready.isnan().sum()>0:
                    print('nan in input image please check!')
                if data_i < 3 and step < 10:
                    img_diff_print(img_processed, image_ready,'img_processed','image_ready')

                
                data_give = custom_image_data_give(data_out, image_ready)
                result = model(return_loss=True, **data_give)
                last_time = time_counter(last_time, 'model forward', time_test_flag)
                
                loss = 0
                for key in result:
                    if 'loss' in key:
                        loss = loss + result[key]
                advloss = - loss
                advloss.backward()
                
                # attack.step img
                imgs_noisy_grad = imgs_noisy.grad.detach()
                imgs_noisy = imgs_noisy - lr * imgs_noisy_grad.sign()
                
                # attack.project img
                diff = imgs_noisy - orig_imgs_input
                diff = torch.clamp(diff, - eps01, eps01)
                imgs_noisy = torch.clamp(diff + orig_imgs_input, 0, 1)
                
                last_time = time_counter(last_time, 'model backward', time_test_flag)
                try:
                    print('attack step:', step, 
                        'model_loss:',round(float(loss),5),
                        'fgsm_eps:', round(float(eps01),5),
                        )
                except Exception as e:
                    print('print_func error:',e)
        else:
            # No gt bbox, directly test
            pass
            
        ############################################    
        # final eval！
        with torch.no_grad():
            imgs_noisy = imgs_noisy.clone().detach()
            ############ resize norm pad
            image_ready = custom_differentiable_transform(
                    img_tensor_rgb_6chw_0to1=imgs_noisy,
                    img_metas=img_metas,
                )
            last_time = time_counter(last_time, 'img rsnmpd', time_test_flag)
            if image_ready.isnan().sum()>0:
                print('nan in input image please check!')

            img_diff_print(img_processed, image_ready,'img_processed','image_ready')

            data_give = custom_image_data_give(data_out, image_ready)
            data_give = custom_data_postprocess_eval(data_give)
            result = model(return_loss=False, rescale=True, **data_give)
            result = custom_result_postprocess(result)

        last_time = time_counter(last_time, 'eval forward', time_test_flag)
        if 'pts_bbox' in result[0]:
            _scores = result[0]['pts_bbox']['scores_3d']
        elif 'img_bbox' in result[0]:
            _scores = result[0]['img_bbox']['scores_3d']
        else:
            _scores = result[0]['scores_3d']
        if len(_scores)>0:
            print('max conf:', round(float(_scores.max()),3))
        else:
            print('nothing detected')
 
        mmcv.dump(result, scattered_result_path)
        results.extend(result)
        prog_bar.update()

    return results

apis_common/test_fgsm_img.py
zzj403/BEV_Robust
import mmcv
import torch
import numpy as np
import PIL.Image as Image
import torchvision.transforms as transforms
import torch.nn.functional as F
import torchvision
from torchvision.utils import save_image
import cv2
import time
import os
import pickle
from extend.custom_func import *
from extend_common.img_check import img_diff_print
from extend_common.time_counter import time_counter
from extend_common.patch_apply import apply_patches_by_info
from extend_common.path_string_split import split_path_string_to_multiname



def single_gpu_test(model, data_loader,
                    patch_save_prefix=None, 
                    area_rate_str=None,
                    optim_lr=None
                    ):
    
    model.eval()
    dataset = data_loader.dataset
    device = model.src_device_obj
    
    patch_save_dir = patch_save_prefix +'_area'+area_rate_str+'_lr'+optim_lr
    os.makedirs(patch_save_dir, exist_ok=True)

    optim_lr = float(optim_lr)


    # 为每一个类别定义一个patch
    # define one patch for evey class
    class_names_list = [
        'car', 'truck', 'construction_vehicle', 
        'bus', 'trailer', 'barrier',
        'motorcycle', 'bicycle', 'pedestrian', 
        'traffic_cone'
    ]
    patch_w = 100
    patch_h = 100
    class_patches_tensor = torch.rand(len(class_names_list),3, patch_h, patch_w).to(device)
    class_patches_tensor.requires_grad_()
    optimizer = torch.optim.Adam([class_patches_tensor], lr=optim_lr)


    time_test_flag = False

    epoch_max = 3
    patch_info_list_database = {}

    # epoch 0 2 4 6 ... for train
    # epoch 1 3 5 7 ... for eval
    for epoch_d in range(epoch_max*2+1):
        epoch = int(epoch_d/2)
        patch_is_training = (epoch_d % 2 == 0)

        if patch_is_training:
            print('=============================')
            print('======= epoch',epoch,' train start =========')
            print('=============================')
        else:
            print('=============================')
            print('======= epoch',epoch,'eval start =========')
            print('=============================')
            results = []

        prog_bar = mmcv.ProgressBar(len(dataset))
        last_time = time.time()
        for data_i, data_out in enumerate(data_loader):


            #### 1. data processing(customed)
            data_out = custom_data_preprocess(data_out)
            img_metas, img_path_list, img_org_np, img_processed, gt_labels_3d = custom_data_work(data_out)
            img_tensor_ncam = custom_img_read_from_img_org(img_org_np, device)
            last_time = time_counter(last_time, 'data load', time_test_flag)
            cam_num = len(img_path_list)

            #### 2. read patch info from file/database
            if not str(data_i) in patch_info_list_database:
                patch_info_list = []
                for cams_i in range(cam_num):
                    img_path = img_path_list[cams_i]
                    file_name_valid_list = split_path_string_to_multiname(img_path)[-3:]
                    file_name_valid_list.insert(0, '/data/zijian/mycode/BEV_Robust/TransFusion/patch_info_2d3d3dt_square_dir/all')
                    info_path = os.path.join(*file_name_valid_list)
                    info_path = info_path.replace('.jpg', '.pkl')
                    info_i = pickle.load(open(info_path, 'rb'))
                    patch_info_list.append(info_i)
                patch_info_list_database[str(data_i)] = patch_info_list
            else:
                patch_info_list = patch_info_list_database[str(data_i)]
            last_time = time_counter(last_time, 'read pkl', time_test_flag)


            #### 3. apply patch
            patched_img_tensor_ncam = img_tensor_ncam.clone()
            # to avoid no_gt
            has_gt_flag = gt_labels_3d.shape[0] != 0
            if has_gt_flag:
                if patch_is_training:
                    # for patch training
                    for cams_i in range(cam_num):
                        patch_info_in_cami = patch_info_list[cams_i]
                        patched_img_tensor_ncam[cams_i] = apply_patches_by_info(
                            info=patch_info_in_cami, 
                            image=patched_img_tensor_ncam[cams_i], 
                            patch_book=class_patches_tensor,
                            area_str=area_rate_str,
                            )
                else:
                    with torch.no_grad():
                        # for patch eval
                        for cams_i in range(cam_num):
                            patch_info_in_cami = patch_info_list[cams_i]
                            patched_img_tensor_ncam[cams_i] = apply_patches_by_info(
                                info=patch_info_in_cami, 
                                image=patched_img_tensor_ncam[cams_i], 
                                patch_book=class_patches_tensor,
                                area_str=area_rate_str,
                                )
            else: # 没有gt 图像不做改变 if no gt donot change images
                pass

            if not has_gt_flag and patch_is_training: 
                # 训练时，无gt的图直接跳过
                # when training, img with no gt will be skip
                # 测试时，正常测试，不跳过
                # when evaluating, img with no gt will still be evaluated
                continue


            # save for watch
            if patch_is_training and data_i % 100 == 0:
                save_image(patched_img_tensor_ncam, os.path.join(patch_save_dir, str(data_i)+'.png'))
            


            #### 4. resize norm pad
            image_ready = custom_differentiable_transform(
                    img_tensor_rgb_6chw_0to1=patched_img_tensor_ncam,
                    img_metas=img_metas,
                )
            last_time = time_counter(last_time, 'img rsnmpd', time_test_flag)

            if image_ready.isnan().sum()>0:
                print('nan in input image please check!')
            if data_i < 10:
                img_diff_print(img_processed, image_ready,'img_processed','image_ready')


            #### 5. update patch or evaluate



            if patch_is_training: # 在训练 更新patch
                data_give = custom_image_data_give(data_out, image_ready)
                result = model(return_loss=True, **data_give)
                last_time = time_counter(last_time, 'model forward', time_test_flag)
                
                loss = 0
                for key in result:
                    if 'loss' in key:
                        loss = loss + result[key]
                advloss = - loss

                # attack.step img
                optimizer.zero_grad()
                advloss.backward()
                optimizer.step()

                # attack.project img
                class_patches_tensor.data = torch.clamp(class_patches_tensor, 0, 1)
                last_time = time_counter(last_time, 'model backward', time_test_flag)
                print('attack step:', data_i, 
                        'model_loss:',round(float(loss),5),
                        )


            else:
                with torch.no_grad():
                    data_give = custom_image_data_give(data_out, image_ready)
                    data_give = custom_data_postprocess_eval(data_give)
                    result = model(return_loss=False, rescale=True, **data_give)
                    result = custom_result_postprocess(result)
                    results.extend(result)
                last_time = time_counter(last_time, 'model forward', time_test_flag)

            prog_bar.update()


        #### After one (train or val) epoch_d
        if not patch_is_training:
            print(dataset.evaluate(results,)) # eval_kwargs 在DETR3d里面，不是必须用到
            # class patch is evaluated during training. All evaluation scores are saved in nohup-log.

        ##################################
        # save
        ##################################
        if patch_is_training:
            print('=============================')
            print('======= epoch',epoch,'save =========')
            print('=============================')
            save_class_patches_path = os.path.join(patch_save_dir, 'epoch_'+str(epoch)+'class_patches.pkl')
            pickle.dump(class_patches_tensor.cpu(), open(save_class_patches_path, 'wb'))

    return results



apis_common/test_patch_class.py
zzj403/BEV_Robust
import mmcv
import torch
import numpy as np
import PIL.Image as Image
import torchvision.transforms as transforms
import torch.nn.functional as F
import torchvision
from torchvision.utils import save_image
import cv2
import time
import os
from extend.custom_func import *
from extend_common.img_check import img_diff_print
from extend_common.time_counter import time_counter
from extend_common.path_string_split import split_path_string_to_multiname


def single_gpu_test(model, data_loader,
                    scattered_result_prefix=None, 
                    mask_code=None,
                    max_step=None
                    ):

    model.eval()
    results = []
    
    scattered_result_dir = scattered_result_prefix +'_area'+mask_code+'_step'+max_step
    max_step = int(max_step)
    
    os.makedirs(scattered_result_dir, exist_ok=True)
    device = model.src_device_obj
    
    time_test_flag = False
    prog_bar = mmcv.ProgressBar(len(data_loader.dataset))
    for data_i, data_out in enumerate(data_loader):
        
        results_name = str(data_i) + '.pkl'
        scattered_result_path = os.path.join(scattered_result_dir, results_name)
        
        if os.path.exists(scattered_result_path):
            print(scattered_result_path, 'exists! pass!')
            prog_bar.update()
            continue
        else:
            mmcv.dump(' ', scattered_result_path)
            
        
        last_time = time.time()

        # 1. data processing(customed)
        data_out = custom_data_preprocess(data_out)
        img_metas, img_path_list, img_org_np, img_processed, gt_labels_3d = custom_data_work(data_out)
        img_tensor_ncam = custom_img_read_from_img_org(img_org_np, device)
        last_time = time_counter(last_time, 'data load', time_test_flag)


        # 2. read mask
        # instance patch need read mask
        mask_tensor_list = []
        for img_path in img_path_list:
            file_name_valid_list = split_path_string_to_multiname(img_path)[-3:]
            file_name_valid_list.insert(0, '/data/zijian/mycode/BEV_Robust/TransFusion/patch_instance_mask_dir/'+mask_code)
            mask_path = os.path.join(*file_name_valid_list)
            mask_np = cv2.imread(mask_path)
            mask_tensor = torch.from_numpy(mask_np).permute(2,0,1)[[2,1,0]].float()/255.
            mask_tensor = (mask_tensor>0.5).float()
            mask_tensor_list.append(mask_tensor)
        mask_tensor_ncam = torch.stack(mask_tensor_list).to(device)
        
        


        orig_imgs_input = img_tensor_ncam.clone().detach()


        

        # to avoid no_gt
        if gt_labels_3d.shape[0] != 0:

            # random init noise layer
            noise_layer = torch.rand_like(orig_imgs_input)
            noise_layer.requires_grad_(True)
            optimizer = torch.optim.Adam([noise_layer], lr=0.1)


            for step in range(max_step):
                ############ apply patch
                imgs_noisy = torch.where(mask_tensor_ncam>0, noise_layer, orig_imgs_input)

                ############ resize norm pad
                image_ready = custom_differentiable_transform(
                        img_tensor_rgb_6chw_0to1=imgs_noisy,
                        img_metas=img_metas,
                    )
                last_time = time_counter(last_time, 'img rsnmpd', time_test_flag)

                if image_ready.isnan().sum()>0:
                    print('nan in input image please check!')
                if data_i < 3 and step < 10:
                    img_diff_print(img_processed, image_ready,'img_processed','image_ready')

                
                data_give = custom_image_data_give(data_out, image_ready)
                result = model(return_loss=True, **data_give)
                last_time = time_counter(last_time, 'model forward', time_test_flag)
                
                loss = 0
                for key in result:
                    if 'loss' in key:
                        loss = loss + result[key]
                advloss = - loss


                # attack.step img
                optimizer.zero_grad()
                advloss.backward()
                optimizer.step()

                # attack.project img
                noise_layer.data = torch.clamp(noise_layer, 0, 1)
                last_time = time_counter(last_time, 'model backward', time_test_flag)
                
                try:
                    print('attack step:', step, 
                        'model_loss:', round(float(loss),3),
                        )
                except Exception as e:
                    print('print_func error:',e)
        else:
            # No gt bbox, directly test
            pass
            
        ############################################    
        # final eval！
        with torch.no_grad():
            ############ apply patch
            noise_layer = noise_layer.clone().detach()
            imgs_noisy = torch.where(mask_tensor_ncam>0, noise_layer, orig_imgs_input)

            ############ resize norm pad
            image_ready = custom_differentiable_transform(
                    img_tensor_rgb_6chw_0to1=imgs_noisy,
                    img_metas=img_metas,
                )
            last_time = time_counter(last_time, 'img rsnmpd', time_test_flag)
            if image_ready.isnan().sum()>0:
                print('nan in input image please check!')

            img_diff_print(img_processed, image_ready,'img_processed','image_ready')

            data_give = custom_image_data_give(data_out, image_ready)
            data_give = custom_data_postprocess_eval(data_give)
            result = model(return_loss=False, rescale=True, **data_give)
            result = custom_result_postprocess(result)

        last_time = time_counter(last_time, 'eval forward', time_test_flag)

        if 'pts_bbox' in result[0]:
            _scores = result[0]['pts_bbox']['scores_3d']
        elif 'img_bbox' in result[0]:
            _scores = result[0]['img_bbox']['scores_3d']
        else:
            _scores = result[0]['scores_3d']
        if len(_scores)>0:
            print('max conf:', round(float(_scores.max()),3))
        else:
            print('nothing detected')

        mmcv.dump(result, scattered_result_path)
        results.extend(result)
        prog_bar.update()

    return results

apis_common/test_patch_instance.py
zzj403/BEV_Robust
import mmcv
import torch
import numpy as np
import PIL.Image as Image
import torchvision.transforms as transforms
import torch.nn.functional as F
import torchvision
from torchvision.utils import save_image
import cv2
import time
import os
import pickle
from extend.custom_func import *
from extend_common.img_check import img_diff_print
from extend_common.time_counter import time_counter
from extend_common.patch_apply import apply_patches_by_info_ovlp
from extend_common.path_string_split import split_path_string_to_multiname


def single_gpu_test(model, data_loader,
                    scattered_result_prefix=None,
                    area_rate_str=None,
                    optim_lr=None,
                    optim_step=None,
                    ):
    
    model.eval()
    device = model.src_device_obj
    
    scattered_result_dir = scattered_result_prefix + '_rate'+area_rate_str + '_lr' + str(optim_lr)+'_step' + str(optim_step)
    os.makedirs(scattered_result_dir, exist_ok=True)

    optim_lr = float(optim_lr)
    optim_step = int(optim_step)


    '''
        这里定义一下整体的流程
        for循环遍历
        为每一帧的每一个overlap的物体，定义一个patch的tensor，这个tensor会投影到两个图像上，
        所有的 overlap的物体的tensor投影完成之后，通过for循环，adam lr=0.1循环优化20step，
        当场检测这张攻击后的图，获得pkl文件，保存。
        如果没有overlap物体，则直接原图eval。

        为了保证2视角下，对抗优化的像素的一致性，保证两个视角下的patch视觉相似，patch的tensor尺寸设计为50x50
        
        总体上思路和 instance-level 的攻击类似
        在全部pkl保存完之后，使用 overlap专用的eval方法进行eval
    '''

    
    max_step = optim_step
    cam_num = 6
    results = []

    prog_bar = mmcv.ProgressBar(len(data_loader.dataset))
    last_time = time.time()
    time_test_flag = False
    for data_i, data_out in enumerate(data_loader):
        
        results_name = str(data_i) + '.pkl'
        scattered_result_path = os.path.join(scattered_result_dir, results_name)
        
        if os.path.exists(scattered_result_path):
            print(scattered_result_path, 'exists! pass!')
            prog_bar.update()
            continue
        else:
            mmcv.dump(' ', scattered_result_path)


        #### 1. data processing(customed)
        data_out = custom_data_preprocess(data_out)
        img_metas, img_path_list, img_org_np, img_processed, gt_labels_3d = custom_data_work(data_out)
        img_tensor_ncam = custom_img_read_from_img_org(img_org_np, device)
        last_time = time_counter(last_time, 'data load', time_test_flag)
        cam_num = len(img_path_list)



        #### 2. read patch info from file
        patch_info_list = []
        for cams_i in range(cam_num):
            img_path = img_path_list[cams_i]
            file_name_valid_list = split_path_string_to_multiname(img_path)[-3:]
            file_name_valid_list.insert(0, '/data/zijian/mycode/BEV_Robust/TransFusion/patch_info_2d3d3dt_square_dir/all')
            info_path = os.path.join(*file_name_valid_list)
            info_path = info_path.replace('.jpg', '.pkl')
            info_i = pickle.load(open(info_path, 'rb'))




            patch_info_list.append(info_i)
        last_time = time_counter(last_time, 'read pkl', time_test_flag)




        #### 3. judge overlap object exists or not
        # 同时存在gt和ovlp物体，才进行攻击， 否则直接eval
        attack_flag = False
        has_gt_flag = gt_labels_3d.shape[0] != 0
        if has_gt_flag:
            ovlp_flag = patch_info_list[0]['objects_info']['is_overlap']
            ovlp_num = ovlp_flag.sum()
            ovlp_exist = (ovlp_num > 0)
            if ovlp_exist:
                attack_flag = True
    
        if attack_flag:
            patch_w = 50
            patch_h = 50
            ovlp_inst_patch_tensor = torch.rand(ovlp_num, 3, patch_h, patch_w).to(device)
            ovlp_inst_patch_tensor.requires_grad_()
            optimizer = torch.optim.Adam([ovlp_inst_patch_tensor], lr=optim_lr)



            for step in range(max_step):
                ####  apply patch
                patched_img_tensor_ncam = img_tensor_ncam.clone()
                for cams_i in range(cam_num):
                    patch_info_in_cami = patch_info_list[cams_i]
                    patched_img_tensor_ncam[cams_i] = apply_patches_by_info_ovlp(
                        info=patch_info_in_cami, 
                        image=patched_img_tensor_ncam[cams_i], 
                        patch_book=ovlp_inst_patch_tensor,
                        area_str=area_rate_str,
                        )
                
                ####  resize norm pad
                image_ready = custom_differentiable_transform(
                        img_tensor_rgb_6chw_0to1=patched_img_tensor_ncam,
                        img_metas=img_metas,
                    )
                last_time = time_counter(last_time, 'img rsnmpd', time_test_flag)
                if image_ready.isnan().sum()>0:
                    print('nan in input image please check!')
                if data_i < 3 and step < 10:
                    img_diff_print(img_processed, image_ready,'img_processed','image_ready')
                

                data_give = custom_image_data_give(data_out, image_ready)
                result = model(return_loss=True, **data_give)
                last_time = time_counter(last_time, 'model forward', time_test_flag)
                loss = 0
                for key in result:
                    if 'loss' in key:
                        loss = loss + result[key]
                advloss = - loss

                # attack.step img
                optimizer.zero_grad()
                advloss.backward()
                optimizer.step()

                # attack.project img
                ovlp_inst_patch_tensor.data = torch.clamp(ovlp_inst_patch_tensor, 0, 1)
                last_time = time_counter(last_time, 'model backward', time_test_flag)
                print('attack step:', step, 
                        'model_loss:',round(float(loss),5),
                        )


            #########################
            ##### 攻击结束，最后粘贴patch，准备eval
            ####  apply patch
            patched_img_tensor_ncam = img_tensor_ncam.clone()
            for cams_i in range(cam_num):
                patch_info_in_cami = patch_info_list[cams_i]
                patched_img_tensor_ncam[cams_i] = apply_patches_by_info_ovlp(
                    info=patch_info_in_cami, 
                    image=patched_img_tensor_ncam[cams_i], 
                    patch_book=ovlp_inst_patch_tensor,
                    area_str=area_rate_str,
                    )
            if data_i < 30:
                save_image(patched_img_tensor_ncam, os.path.join(scattered_result_dir, str(data_i)+'.png'))
        else:
            # 同时存在gt和ovlp物体，才进行攻击， 否则直接eval
            patched_img_tensor_ncam = img_tensor_ncam.clone()



        ############################################    
        # 结尾 eval！
        with torch.no_grad():
            ############ resize norm pad
            image_ready = custom_differentiable_transform(
                    img_tensor_rgb_6chw_0to1=patched_img_tensor_ncam,
                    img_metas=img_metas,
                )
            last_time = time_counter(last_time, 'img rsnmpd', time_test_flag)
            if image_ready.isnan().sum()>0:
                print('nan in input image please check!')
            img_diff_print(img_processed, image_ready,'img_processed','image_ready')

            data_give = custom_image_data_give(data_out, image_ready)
            data_give = custom_data_postprocess_eval(data_give)
            result = model(return_loss=False, rescale=True, **data_give)
            result = custom_result_postprocess(result)
            results.extend(result)
            mmcv.dump(result, scattered_result_path)

        prog_bar.update()
    return results




apis_common/test_patch_overlap.py
zzj403/BEV_Robust
import mmcv
import torch
import numpy as np
import PIL.Image as Image
import torchvision.transforms as transforms
import torch.nn.functional as F
import torchvision
from torchvision.utils import save_image
import cv2
import time
import os
import pickle
from extend.custom_func import *
from extend_common.img_check import img_diff_print
from extend_common.time_counter import time_counter
from extend_common.patch_apply import apply_patches_by_info_4side
from extend_common.path_string_split import split_path_string_to_multiname
from extend_common.get_scene_start_idx import get_scene_start_idx


def single_gpu_test(model, data_loader,
                    scattered_result_prefix=None, 
                    area_rate_str=None,
                    optim_lr=None,
                    optim_step=None,
                    index_min = None, 
                    index_max = None,
                    ):
    
    model.eval()
    dataset = data_loader.dataset
    device = model.src_device_obj
    
    scattered_result_dir = scattered_result_prefix +'_area'+area_rate_str+'_lr'+optim_lr+'_step' + optim_step
    os.makedirs(scattered_result_dir, exist_ok=True)

    optim_lr = float(optim_lr)
    optim_step = int(optim_step)

    scene_start_idx_list = get_scene_start_idx()
    max_epoch_local = optim_step
    
    patch_info_list_database = {}
    time_test_flag = False


    results = []
    prog_bar = mmcv.ProgressBar(len(dataset))
    last_time = time.time()
    for data_i, data_out in enumerate(data_loader):
        if data_i < index_min:
            prog_bar.update()
            continue
        if data_i > index_max:
            break
        
        #### 1. data processing(customed)
        data_out = custom_data_preprocess(data_out)
        _, img_path_list, _, _, _ = custom_data_work(data_out)
        last_time = time_counter(last_time, 'data load', time_test_flag)
        cam_num = len(img_path_list)


        #### 2. read patch info from file/database
        if not str(data_i) in patch_info_list_database:
            patch_info_list = []
            for cams_i in range(cam_num):
                img_path = img_path_list[cams_i]
                file_name_valid_list = split_path_string_to_multiname(img_path)[-3:]
                file_name_valid_list.insert(0, '/data/zijian/mycode/BEV_Robust/TransFusion/patch_info_2d3d3dt_square_dir/all')
                info_path = os.path.join(*file_name_valid_list)
                info_path = info_path.replace('.jpg', '.pkl')
                info_i = pickle.load(open(info_path, 'rb'))
                patch_info_list.append(info_i)
            patch_info_list_database[str(data_i)] = patch_info_list
        else:
            patch_info_list = patch_info_list_database[str(data_i)]
        last_time = time_counter(last_time, 'read pkl', time_test_flag)

        
        '''
            由于我们要一个场景（大概40帧左右），一起进行攻击
            所以我需要先遍历数据集，把这一个场景的数据先拿出来，统计里面instance的数量，构建一个 patch 库
            然后再在读取出的这一个场景的数据里做攻击

            如果是场景的第0帧
            则开始遍历当前场景，直到下一个第0帧的出现，这时候暂存下一个第0帧
            遍历场景时，存下所有的注释信息，
            并从之前存好的 patch info 中 获取 instance_token
        '''


        scene_start_here_flag = (data_i in scene_start_idx_list)
        
        go_to_training_flag = False
        
        if data_i == 0:
            # 第0帧
            # start new
            data_in_scene_list = []
            patch_info_in_scene_list = []
            data_i_list = []
            data_in_scene_list.append(data_out)
            patch_info_in_scene_list.append(patch_info_list)
            data_i_list.append(data_i)
        elif scene_start_here_flag and data_i > 0:
            # 之后的每一个首帧
            # 存一个连续场景的全部 data 和 patch_info
            # end old
            try:
                data_in_scene_list_full = data_in_scene_list
                patch_info_in_scene_list_full = patch_info_in_scene_list
                data_i_list_full = data_i_list
                go_to_training_flag = True
            except:
                print('start from data_i:', data_i)
            # start new
            data_in_scene_list = []
            patch_info_in_scene_list = []
            data_i_list = []
            data_in_scene_list.append(data_out)
            patch_info_in_scene_list.append(patch_info_list)
            data_i_list.append(data_i)
        elif data_i == len(dataset)-1:
            data_in_scene_list.append(data_out)
            patch_info_in_scene_list.append(patch_info_list)
            data_i_list.append(data_i)
            # 最后一帧
            # end old
            data_in_scene_list_full = data_in_scene_list
            patch_info_in_scene_list_full = patch_info_in_scene_list
            data_i_list_full = data_i_list
            go_to_training_flag = True
        else:
            data_in_scene_list.append(data_out)
            patch_info_in_scene_list.append(patch_info_list)
            data_i_list.append(data_i)
        prog_bar.update()
        
        if go_to_training_flag:
            # local dataset: data_in_scene_list_full
            # local dataset: patch_info_in_scene_list_full
            # local dataset: data_i_list_full
            scene_length = len(data_in_scene_list_full)
            
            ###### 1.构建patch库 Establish local-scene patchbook
            # 每个物体的4个面，都放patch，
            # patchtensor的形状, 由实际的patchsize确定，兼容正方形patch
            instance_token_list = []
            patch_4side_book_list = []
            for i_local in range(scene_length):
                # 1.把数据拿出来，处理数据
                data_local = data_in_scene_list_full[i_local]
                patch_info_local = patch_info_in_scene_list_full[i_local]
                _, _, _, _, gt_labels_3d = custom_data_work(data_local)
                # 2.判断有没有gt
                # 防止出现 no_gt
                has_gt_flag = (gt_labels_3d.shape[0] != 0) and (type(patch_info_local[0]) != str)
                if has_gt_flag:
                    scene_name = patch_info_local[0]['scene_info']['scene_name']
                    instance_tokens_i = patch_info_local[0]['objects_info']['instance_tokens']
                    for inst_tk_idx in range(len(instance_tokens_i)):
                        instance_token = instance_tokens_i[inst_tk_idx]
                        if not instance_token in instance_token_list:
                            # 添加patch 
                            # 根据最先出现的patch，标注的信息，添加4个patch
                            for j_cam_1frame in range(cam_num):
                                if patch_info_local[j_cam_1frame]['patch_visible_bigger'][inst_tk_idx]:
                                    # 如果可以被，当前的camera看到，则添加，否则不添加
                                    patch_3d_wh = patch_info_local[j_cam_1frame]['patch_3d_temporal']['patch_3d_wh'][inst_tk_idx]
                                    patch_3d_wh_use = patch_3d_wh[area_rate_str]

                                    patch_4side_ = []
                                    for j_side in range(4):
                                        patch_w_real, patch_h_real = patch_3d_wh_use[j_side]
                                        # 遵循每1m 100pix的密度
                                        patch_w_tensor = int(patch_w_real*100)
                                        patch_h_tensor = int(patch_h_real*100)
                                        patch_jside_ = torch.rand(3, patch_h_tensor, patch_w_tensor).to(device)
                                        patch_jside_.requires_grad_()
                                        patch_4side_.append(patch_jside_)

                                    instance_token_list.append(instance_token)
                                    patch_4side_book_list.extend(patch_4side_)

            # 为这些patch定义 优化器
            optimizer = torch.optim.Adam(patch_4side_book_list, lr=optim_lr)

            # 以后每一次取用，都需要，结合instance_token_list获取 token对应的index，再用

            
            for epoch_local in range(max_epoch_local):
                print('scene_name:', scene_name,'start epoch_local', epoch_local,'training')
                for i_local in range(scene_length):

                    ##############  把数据拿出来，处理数据 Take out the data and process the data
                    data_local = data_in_scene_list_full[i_local]
                    patch_info_local = patch_info_in_scene_list_full[i_local]
                    img_metas, img_path_list, img_org_np, img_processed, gt_labels_3d = custom_data_work(data_local)
                    img_tensor_ncam = custom_img_read_from_img_org(img_org_np, device)
                    last_time = time_counter(last_time, 'data process', time_test_flag)

                    ############## apply patch
                    patched_img_tensor_ncam = img_tensor_ncam.clone()
                    # in case of no_gt
                    has_gt_flag = (gt_labels_3d.shape[0] != 0) and (type(patch_info_local[0]) != str)
                    if has_gt_flag:
                        # apply patch
                        for cams_i in range(cam_num):
                            patch_info_in_cami = patch_info_local[cams_i]
                            patched_img_tensor_ncam[cams_i] = apply_patches_by_info_4side(
                                info=patch_info_in_cami, 
                                image=patched_img_tensor_ncam[cams_i], 
                                instance_token_book=instance_token_list,
                                patch_book_4side=patch_4side_book_list,
                                area_str=area_rate_str,
                                )
                            # patched_img_tensor_ncam[cams_i] = (patched_img_tensor_ncam[cams_i] + patch_4side_book_list[0].mean()/1000).clamp(0,1)
                    else: # no gt，图像不做改变，也不必优化patch
                        continue

                    last_time = time_counter(last_time, 'apply patch', time_test_flag)

                    ############ resize norm pad
                    image_ready = custom_differentiable_transform(
                            img_tensor_rgb_6chw_0to1=patched_img_tensor_ncam,
                            img_metas=img_metas,
                        )
                    last_time = time_counter(last_time, 'img rsnmpd', time_test_flag)


                    if image_ready.isnan().sum()>0:
                        print('nan in input image please check!')

                    data_i_actual = data_i_list_full[i_local]
                    if data_i_actual < 100 and epoch_local < 3 and i_local < 3:
                        img_diff_print(img_processed, image_ready,'img_processed','image_ready')


                    data_give = custom_image_data_give(data_local, image_ready)
                    result = model(return_loss=True, **data_give) # 经过model， data中的img会被修改为[6,3,H,W]
                    last_time = time_counter(last_time, 'model forward', time_test_flag)
                    loss = 0
                    for key in result:
                        if 'loss' in key:
                            loss = loss + result[key]
                    advloss = - loss
                    optimizer.zero_grad()
                    advloss.backward()
                    optimizer.step()
                    optimizer.zero_grad()

                    last_time = time_counter(last_time, 'model backward', time_test_flag)

                    for _patch_i in range(len(patch_4side_book_list)):
                        patch_4side_book_list[_patch_i].data = torch.clamp(patch_4side_book_list[_patch_i], 0, 1)
                    last_time = time_counter(last_time, 'patch clamp', time_test_flag)
                    print('attack step:', i_local, 
                            'model_loss:',round(float(loss),5),
                            )
            #########################
            ##### 攻击结束，最后再遍历一遍，粘贴patch，eval
            print('scene_name:', scene_name,'start eval')
            prog_bar_local_eval = mmcv.ProgressBar(scene_length)
            with torch.no_grad():
                for i_local in range(scene_length):

                    #################  把数据拿出来，处理数据
                    data_local = data_in_scene_list_full[i_local]
                    patch_info_local = patch_info_in_scene_list_full[i_local]
                    img_metas, img_path_list, img_org_np, img_processed, gt_labels_3d = custom_data_work(data_local)
                    img_tensor_ncam = custom_img_read_from_img_org(img_org_np, device)

                    ################  安装patch
                    patched_img_tensor_ncam = img_tensor_ncam.clone()
                    # 防止出现 no_gt
                    has_gt_flag = (gt_labels_3d.shape[0] != 0) and (type(patch_info_local[0]) != str)
                    if has_gt_flag:
                        # apply patch
                        for cams_i in range(cam_num):
                            patch_info_in_cami = patch_info_local[cams_i]
                            patched_img_tensor_ncam[cams_i] = apply_patches_by_info_4side(
                                info=patch_info_in_cami, 
                                image=patched_img_tensor_ncam[cams_i], 
                                instance_token_book=instance_token_list,
                                patch_book_4side=patch_4side_book_list,
                                area_str=area_rate_str,
                                )
                    else: # 没有gt，图像不做改变，直接eval
                        pass

                    ############ resize norm pad
                    image_ready = custom_differentiable_transform(
                            img_tensor_rgb_6chw_0to1=patched_img_tensor_ncam,
                            img_metas=img_metas,
                        )
                    last_time = time_counter(last_time, 'img rsnmpd', time_test_flag)
                    if image_ready.isnan().sum()>0:
                        print('nan in input image please check!')
                    if i_local < 3:
                        img_diff_print(img_processed, image_ready,'img_processed','image_ready')

                    data_give = custom_image_data_give(data_local, image_ready)
                    data_give = custom_data_postprocess_eval(data_give)
                    result = model(return_loss=False, rescale=True, **data_give)
                    result = custom_result_postprocess(result)
                    results.extend(result)

                    data_i_actual = data_i_list_full[i_local]
                    scattered_result_path = os.path.join(scattered_result_dir, str(data_i_actual)+'.pkl')
                    mmcv.dump(result, scattered_result_path)
                    if data_i_actual < 100:
                        save_image(patched_img_tensor_ncam, os.path.join(scattered_result_dir, str(data_i_actual)+'.png'))
                    prog_bar_local_eval.update()
                print()
    return results




apis_common/test_patch_temporal.py
zzj403/BEV_Robust
import mmcv
import torch
import numpy as np
import PIL.Image as Image
import torchvision.transforms as transforms
import torch.nn.functional as F
import torchvision
from torchvision.utils import save_image
import cv2
import time
import os
from extend.custom_func import *
from extend_common.img_check import img_diff_print
from extend_common.time_counter import time_counter

def single_gpu_test(model, data_loader,
                    scattered_result_prefix=None, 
                    eps255=None, 
                    max_step=None
                    ):
   
    model.eval()
    results = []
    
    scattered_result_dir = scattered_result_prefix +'_eps'+eps255+'_step'+max_step
    eps255 = float(eps255)
    max_step = int(max_step)
    
    os.makedirs(scattered_result_dir, exist_ok=True)
    device = model.src_device_obj
    
    time_test_flag = False
    prog_bar = mmcv.ProgressBar(len(data_loader.dataset))
    for data_i, data_out in enumerate(data_loader):
        
        results_name = str(data_i) + '.pkl'
        scattered_result_path = os.path.join(scattered_result_dir, results_name)
        
        if os.path.exists(scattered_result_path):
            print(scattered_result_path, 'exists! pass!')
            prog_bar.update()
            continue
        else:
            mmcv.dump(' ', scattered_result_path)
            
        
        last_time = time.time()

        # 1. data processing(customed)
        data_out = custom_data_preprocess(data_out)
        img_metas, img_path_list, img_org_np, img_processed, gt_labels_3d = custom_data_work(data_out)
        img_tensor_ncam = custom_img_read_from_img_org(img_org_np, device)
        last_time = time_counter(last_time, 'data load', time_test_flag)


        # PGD all region attack, do not need patch info

        imgs_noisy = img_tensor_ncam.clone().detach()
        orig_imgs_input = img_tensor_ncam.clone().detach()

        
        eps01 = eps255 / 255.
        lr = eps01 / (max_step-2)

        # pgd 是随机初始化的 random start
        delta = torch.rand_like(orig_imgs_input)*2*eps01 - eps01
        imgs_noisy = orig_imgs_input + delta
        imgs_noisy = torch.clamp(imgs_noisy, 0, 1)
        

        # to avoid no_gt
        if gt_labels_3d.shape[0] != 0:
            for step in range(max_step):
                imgs_noisy = imgs_noisy.clone().detach().requires_grad_(True)

                ############ resize norm pad
                image_ready = custom_differentiable_transform(
                        img_tensor_rgb_6chw_0to1=imgs_noisy,
                        img_metas=img_metas,
                    )
                last_time = time_counter(last_time, 'img rsnmpd', time_test_flag)

                if image_ready.isnan().sum()>0:
                    print('nan in input image please check!')
                if data_i < 3 and step < 10:
                    img_diff_print(img_processed, image_ready,'img_processed','image_ready')

                
                data_give = custom_image_data_give(data_out, image_ready)
                result = model(return_loss=True, **data_give)
                last_time = time_counter(last_time, 'model forward', time_test_flag)
                
                loss = 0
                for key in result:
                    if 'loss' in key:
                        loss = loss + result[key]
                advloss = - loss
                advloss.backward()
                
                # attack.step img
                imgs_noisy_grad = imgs_noisy.grad.detach()
                imgs_noisy = imgs_noisy - lr * imgs_noisy_grad.sign()
                
                # attack.project img
                diff = imgs_noisy - orig_imgs_input
                diff = torch.clamp(diff, - eps01, eps01)
                imgs_noisy = torch.clamp(diff + orig_imgs_input, 0, 1)
                
                
                last_time = time_counter(last_time, 'model backward', time_test_flag)
                try:
                    print('attack step:', step, 
                        'model_loss:', round(float(loss),3),
                        'pgd_eps:', round(float(eps01),5),
                        'pgd_lr:', round(float(lr),5),
                        )
                except Exception as e:
                    print('print_func error:',e)
        else:
            # No gt bbox, directly test
            pass
            
        ############################################    
        # final eval！
        with torch.no_grad():
            ############ resize norm pad
            image_ready = custom_differentiable_transform(
                    img_tensor_rgb_6chw_0to1=imgs_noisy,
                    img_metas=img_metas,
                )
            last_time = time_counter(last_time, 'img rsnmpd', time_test_flag)
            if image_ready.isnan().sum()>0:
                print('nan in input image please check!')

            img_diff_print(img_processed, image_ready,'img_processed','image_ready')

            data_give = custom_image_data_give(data_out, image_ready)
            data_give = custom_data_postprocess_eval(data_give)
            result = model(return_loss=False, rescale=True, **data_give)
            result = custom_result_postprocess(result)
        imgs_noisy = imgs_noisy.clone().detach()

        last_time = time_counter(last_time, 'eval forward', time_test_flag)
        if 'pts_bbox' in result[0]:
            _scores = result[0]['pts_bbox']['scores_3d']
        elif 'img_bbox' in result[0]:
            _scores = result[0]['img_bbox']['scores_3d']
        else:
            _scores = result[0]['scores_3d']
        if len(_scores)>0:
            print('max conf:', round(float(_scores.max()),3))
        else:
            print('nothing detected')

        mmcv.dump(result, scattered_result_path)
        results.extend(result)
        prog_bar.update()

    return results

apis_common/test_pgd_img.py
zzj403/BEV_Robust
import mmcv
import torch
import numpy as np
import PIL.Image as Image
import torchvision.transforms as transforms
import torch.nn.functional as F
import torchvision
from torchvision.utils import save_image
import cv2
import time
import os
from extend.custom_func import *
from extend_common.img_check import img_diff_print
from extend_common.time_counter import time_counter


def single_gpu_test(model, data_loader,
                    scattered_result_prefix=None, 
                    img_eps255=None, 
                    point_eps_m=None, 
                    max_step=None,
                    ):

    model.eval()
    results = []
    dataset = data_loader.dataset
    
    scattered_result_dir = scattered_result_prefix +'_imgeps'+img_eps255 +'_pointeps'+point_eps_m +'_step'+max_step
    img_eps255 = float(img_eps255)
    point_eps_m = float(point_eps_m)
    max_step = int(max_step)
    
    os.makedirs(scattered_result_dir, exist_ok=True)
    device = model.src_device_obj
    
    time_test_flag = False
    prog_bar = mmcv.ProgressBar(len(dataset))
    for data_i, data_out in enumerate(data_loader):
        
        results_name = str(data_i) + '.pkl'
        scattered_result_path = os.path.join(scattered_result_dir, results_name)
        
        if os.path.exists(scattered_result_path):
            print(scattered_result_path, 'exists! pass!')
            prog_bar.update()
            continue
        else:
            mmcv.dump(' ', scattered_result_path)
            
        
        last_time = time.time()

        # 1. data processing(customed)
        data_out = custom_data_preprocess(data_out)
        img_metas, img_path_list, img_org_np, img_processed, gt_labels_3d, points_tensor = custom_data_work_point(data_out)
        img_tensor_ncam = custom_img_read_from_img_org(img_org_np, device)
        last_time = time_counter(last_time, 'data load', time_test_flag)


        # PGD all region attack, do not need patch info

        orig_imgs_input = img_tensor_ncam.clone().detach()
        orig_pts_input = points_tensor.clone().detach()

        
        img_eps01 = img_eps255 / 255.
        point_eps_m = point_eps_m
        img_lr = img_eps01 / (max_step-2)
        point_lr = point_eps_m / (max_step-2)

        # pgd random start
        delta = torch.rand_like(orig_imgs_input) * 2 * img_eps01 - img_eps01
        imgs_noisy = orig_imgs_input + delta
        imgs_noisy = torch.clamp(imgs_noisy, 0, 1)

        
        # pgd random start
        delta = torch.rand_like(orig_pts_input) * 2 * point_eps_m - point_eps_m
        pts_noisy = orig_pts_input + delta
        # points do not need clamp to [0,1], 
        # but we only change xyz-channel of points
        pts_noisy[:,3:] = orig_pts_input[:,3:]


        

        # to avoid no_gt
        if gt_labels_3d.shape[0] != 0:
            for step in range(max_step):
                imgs_noisy = imgs_noisy.clone().detach().requires_grad_(True)
                pts_noisy  = pts_noisy.clone().detach().requires_grad_(True)

                ############ resize norm pad
                image_ready = custom_differentiable_transform(
                        img_tensor_rgb_6chw_0to1=imgs_noisy,
                        img_metas=img_metas,
                    )
                last_time = time_counter(last_time, 'img rsnmpd', time_test_flag)
                if image_ready.isnan().sum()>0:
                    print('nan in input image please check!')
                if data_i < 3 and step < 10:
                    img_diff_print(img_processed, image_ready,'img_processed','image_ready')

                points_ready = pts_noisy

                data_give = custom_image_data_give_point(data_out, image_ready, points_ready)
                result = model(return_loss=True, **data_give)
                last_time = time_counter(last_time, 'model forward', time_test_flag)
                
                loss = 0
                for key in result:
                    if 'loss' in key:
                        loss = loss + result[key]
                advloss = - loss
                advloss.backward()
                
                # attack.step img
                imgs_noisy_grad = imgs_noisy.grad.detach()
                imgs_noisy = imgs_noisy - img_lr * imgs_noisy_grad.sign()
                # attack.project img
                diff = imgs_noisy - orig_imgs_input
                diff = torch.clamp(diff, -img_eps01, img_eps01)
                imgs_noisy = torch.clamp(diff + orig_imgs_input, 0, 1)


                # attack.step point
                pts_noisy_grad = pts_noisy.grad.detach()
                pts_noisy = pts_noisy - point_lr * pts_noisy_grad.sign()
                # attack.project point
                diff = pts_noisy - orig_pts_input
                diff = torch.clamp(diff, - point_eps_m, point_eps_m)
                pts_noisy = diff + orig_pts_input
                # the channel other than xyz cannot be changed
                pts_noisy[:,3:] = orig_pts_input[:,3:]

                
                last_time = time_counter(last_time, 'model backward', time_test_flag)

                try:
                    print('step:', step, 
                        'model_loss:', round(float(loss),3),
                        'img_eps:', round(float(img_eps01),5),
                        'img_lr:', round(float(img_lr),5),
                        'point_eps:', round(float(point_eps_m),5),
                        'point_lr:', round(float(point_lr),5),
                        )
                except Exception as e:
                    print('print_func error:',e)
        else:
            # No gt bbox, directly test
            pass
            
        ############################################    
        # final eval！
        with torch.no_grad():
            imgs_noisy = imgs_noisy.clone().detach()
            ############ resize norm pad
            image_ready = custom_differentiable_transform(
                    img_tensor_rgb_6chw_0to1=imgs_noisy,
                    img_metas=img_metas,
                )
            last_time = time_counter(last_time, 'img rsnmpd', time_test_flag)
            if image_ready.isnan().sum()>0:
                print('nan in input image please check!')
            img_diff_print(img_processed, image_ready,'img_processed','image_ready')

            points_ready = pts_noisy

            data_give = custom_image_data_give_point(data_out, image_ready, points_ready)
            data_give = custom_data_postprocess_eval(data_give)
            result = model(return_loss=False, rescale=True, **data_give)
            result = custom_result_postprocess(result)
            results.extend(result)
            mmcv.dump(result, scattered_result_path)

        last_time = time_counter(last_time, 'eval forward', time_test_flag)
        if 'pts_bbox' in result[0]:
            print('max conf:', round(float(result[0]['pts_bbox']['scores_3d'].max()),3))
        else:
            print('max conf:', round(float(result[0]['scores_3d'].max()),3))

        mmcv.dump(result, scattered_result_path)
        results.extend(result)
        prog_bar.update()

    return results

apis_common/test_pgd_imgpoint.py
zzj403/BEV_Robust
def get_scene_start_idx():
    print('scene_start_idx can only be used for no-shuffled nus-3d dataset')
    print('CANNOT be used for nus-mono3d dataste! Because of different scene order!')
    scene_start_idx_list = [
        0, 39, 79, 119, 158, 197, 236, 276, 316, 356, 396,
        436, 475, 515, 555, 595, 635, 675, 715, 755, 795, 835,
        875, 915, 955, 994, 1034, 1073, 1112, 1151, 1190, 1229,
        1269, 1309, 1349, 1389, 1428, 1467, 1507, 1547, 1586, 1626,
        1666, 1706, 1746, 1785, 1825, 1864, 1903, 1942, 1982, 2023,
        2064, 2105, 2145, 2186, 2227, 2267, 2307, 2347, 2388, 2428,
        2469, 2510, 2550, 2590, 2631, 2671, 2712, 2753, 2794, 2834,
        2874, 2915, 2956, 2996, 3036, 3076, 3116, 3156, 3196, 3236,
        3276, 3317, 3357, 3398, 3438, 3478, 3518, 3558, 3598, 3638,
        3679, 3720, 3760, 3801, 3841, 3882, 3922, 3962, 4002, 4042,
        4082, 4123, 4164, 4205, 4245, 4286, 4327, 4368, 4408, 4448,
        4489, 4529, 4569, 4609, 4649, 4689, 4729, 4769, 4810, 4851,
        4892, 4933, 4974, 5015, 5056, 5096, 5136, 5176, 5217, 5257,
        5297, 5337, 5377, 5417, 5457, 5497, 5538, 5578, 5618, 5658,
        5698, 5738, 5778, 5818, 5858, 5899, 5939, 5979
        ]
    return scene_start_idx_list

if __name__ == '__main__':
    print(get_scene_start_idx())

extend_common/get_scene_start_idx.py
zzj403/BEV_Robust

import torch



def img_diff_print(img1, img2, img1_name, img2_name):
    assert len(img1.shape)==len(img2.shape), 'imgtensor shape length must be the same'
    assert img1.shape==img2.shape, 'imgtensor shape must be the same'

    name_len = max(len(img1_name), len(img2_name))
    print(
        '\n'+img1_name.rjust(name_len,' ')+' range R:',round(float(img1[...,0,:,:].min()),3), round(float(img1[...,0,:,:].max()),3), 
        '\n'+img2_name.rjust(name_len,' ')+' range R:',round(float(img2[...,0,:,:].min()),3), round(float(img2[...,0,:,:].max()),3), 
        '\n'+img1_name.rjust(name_len,' ')+' range G:',round(float(img1[...,1,:,:].min()),3), round(float(img1[...,1,:,:].max()),3), 
        '\n'+img2_name.rjust(name_len,' ')+' range G:',round(float(img2[...,1,:,:].min()),3), round(float(img2[...,1,:,:].max()),3), 
        '\n'+img1_name.rjust(name_len,' ')+' range B:',round(float(img1[...,2,:,:].min()),3), round(float(img1[...,2,:,:].max()),3), 
        '\n'+img2_name.rjust(name_len,' ')+' range B:',round(float(img2[...,2,:,:].min()),3), round(float(img2[...,2,:,:].max()),3), 
        '\n'+img1_name.rjust(name_len,' ')+' shape:', img1.shape,
        '\n'+img2_name.rjust(name_len,' ')+' shape:', img2.shape,
    )

if __name__ == '__main__':
    a = torch.rand(3,10,10)
    b = torch.rand(3,10,10)
    
    img_diff_print(a,b,'aaa', 'd1asweq')
extend_common/img_check.py
zzj403/BEV_Robust
import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn.functional as F
from torchvision.utils import save_image


def proj_3d_to_2d(points_3d_batch, lidar2img, img_org_shape):
    points_4d_batch = torch.cat([points_3d_batch, torch.ones_like(points_3d_batch[:,0:1])], dim=-1)
    points_2d_batch = points_4d_batch @ lidar2img.t()
    points_2d_batch_depth = points_2d_batch[:,2]
    points_2d_batch_depth_positive = torch.clamp(points_2d_batch_depth, min=1e-5)
    points_2d_batch[:,0] /= points_2d_batch_depth_positive
    points_2d_batch[:,1] /= points_2d_batch_depth_positive
    on_the_image_depth = points_2d_batch_depth > 0
    coor_x, coor_y = points_2d_batch[:,0], points_2d_batch[:,1]
    h, w = img_org_shape
    on_the_image = (coor_x > 0) * (coor_x < w) * (coor_y > 0) * (coor_y < h)
    on_the_image = on_the_image * on_the_image_depth
    return points_2d_batch[:,:2], on_the_image



def safe_img_patch_apply(image, patch, x1, y1):
    assert len(image.shape) == len(patch.shape)
    assert image.shape[0] == patch.shape[0]

    try:
        w_patch = patch.shape[-1]
        h_patch = patch.shape[-2]
        x2 = x1 + w_patch
        y2 = y1 + h_patch
        w_img = image.shape[-1]
        h_img = image.shape[-2]

        x1_use = x1
        x2_use = x2
        y1_use = y1
        y2_use = y2
        patch_use = patch

        if x1 < 0:
            x1_use = 0
            patch_use = patch_use[:, :, -x1:]
        elif x1 > w_img-1:
            return image

        if x2 > w_img:
            x2_use = w_img
            patch_use = patch_use[:, :, :-(x2 - w_img)]
        elif x2 < 0:
            return image

        if y1 < 0:
            y1_use = 0
            patch_use = patch_use[:, -y1:, :]
        elif y1 > h_img - 1:
            return image

        if y2 > h_img:
            y2_use = h_img
            patch_use = patch_use[:, :-(y2 - h_img), :]
        elif y2 < 0:
            return image

        image[:, y1_use:y2_use, x1_use:x2_use] = patch_use
    except:
        return image
        print('safe_img_patch_apply failed, skip this patch')
    return image




def apply_patches_by_info(info, image, patch_book, area_str):
    local_classname_list = info['objects_info']['class_name']
    local_classidx_list = info['objects_info']['class_idx']

    # 重新排序，按照画面深度uvd的d排序，以保证从远到近粘贴
    _d = info['objects_info']['uvd_center'][:, 2]
    d_index = _d.sort(descending=True)[1]

    d_index_list = d_index.tolist()

    for i in d_index_list:
        if info['patch_visible_bigger'][i]: # 可见
            assert info['patch_2d']['patch_vertices'][i] != None
            use_patch_vertices_uv = info['patch_2d']['patch_vertices'][i][area_str][:,:2]
            use_patch_classname = local_classidx_list[i]
            use_patch = patch_book[use_patch_classname]
            patch_use = use_patch
            patch_h, patch_w = patch_use.shape[-2:]
            img_h, img_w = image.shape[-2:]

            if len(use_patch_vertices_uv.unique()) <= 4:
                # 低级抠图
                u1 = int(use_patch_vertices_uv[:,0].min())
                u2 = int(use_patch_vertices_uv[:,0].max())
                v1 = int(use_patch_vertices_uv[:,1].min())
                v2 = int(use_patch_vertices_uv[:,1].max())

                need_patch_h = v2 - v1
                need_patch_w = u2 - u1

                if need_patch_h > 0 and need_patch_w > 0:
                    patch_use = F.interpolate(patch_use[None], (need_patch_h, need_patch_w), mode='bilinear', align_corners=False)[0]
                    image = safe_img_patch_apply(
                        image=image, 
                        patch=patch_use, 
                        x1=u1, 
                        y1=v1
                    )
            else:
                # 高级变换

                patch_pad = F.pad(
                    patch_use, (0, img_w-patch_w, 0, img_h-patch_h),
                    )
                patch_mask_pad = torch.zeros_like(image)
                patch_mask_pad[:,:patch_h,:patch_w] = 1
                startpoints = [
                    (0, 0), 
                    (patch_w-1, 0),
                    (patch_w-1, patch_h-1),
                    (0, patch_h-1),
                ]
                endpoints = use_patch_vertices_uv
                patch_pad_trans = transforms.functional.perspective(
                                        patch_pad,
                                        startpoints=startpoints,
                                        endpoints=endpoints
                                        )
                patch_mask_pad_trans = transforms.functional.perspective(
                                        patch_mask_pad,
                                        startpoints=startpoints,
                                        endpoints=endpoints
                                        )
                image = torch.where(patch_mask_pad_trans > 0.5, patch_pad_trans, image)
    return image



def apply_patches_by_info_ovlp(info, image, patch_book, area_str):
    local_classname_list = info['objects_info']['class_name']
    local_classidx_list = info['objects_info']['class_idx']
    
    local_ovlp_flag = info['objects_info']['is_overlap']

    # 重新排序，按照画面深度uvd的d排序，以保证从远到近粘贴
    _d = info['objects_info']['uvd_center'][:, 2]
    d_index = _d.sort(descending=True)[1]
    
    # 需要判断，当前的obj是ovlp的第几个，粘贴对应的patch

    d_index_list = d_index.tolist()

    for i in d_index_list:
        if info['patch_visible_bigger'][i] and local_ovlp_flag[i]: # 可见 且 重叠
            assert info['patch_3d']['patch_vertices'][i] != None
            use_patch_vertices_uv = info['patch_3d']['patch_vertices'][i][area_str][:,:2]
            # 需要判断，当前的obj是ovlp的第几个，粘贴对应的patch
            use_patch_idx = local_ovlp_flag[:(i+1)].sum() - 1
            use_patch = patch_book[use_patch_idx]
            patch_use = use_patch
            patch_h, patch_w = patch_use.shape[-2:]
            img_h, img_w = image.shape[-2:]

            if len(use_patch_vertices_uv.unique()) <= 4:
                # 低级抠图
                u1 = int(use_patch_vertices_uv[:,0].min())
                u2 = int(use_patch_vertices_uv[:,0].max())
                v1 = int(use_patch_vertices_uv[:,1].min())
                v2 = int(use_patch_vertices_uv[:,1].max())

                need_patch_h = v2 - v1
                need_patch_w = u2 - u1

                if need_patch_h > 0 and need_patch_w > 0:
                    patch_use = F.interpolate(patch_use[None], (need_patch_h, need_patch_w), mode='bilinear', align_corners=False)[0]
                    image = safe_img_patch_apply(
                        image=image, 
                        patch=patch_use, 
                        x1=u1, 
                        y1=v1
                    )
            else:
                # 高级变换

                patch_pad = F.pad(
                    patch_use, (0, img_w-patch_w, 0, img_h-patch_h),
                    )
                patch_mask_pad = torch.zeros_like(image)
                patch_mask_pad[:,:patch_h,:patch_w] = 1
                startpoints = [
                    (0, 0), 
                    (patch_w-1, 0),
                    (patch_w-1, patch_h-1),
                    (0, patch_h-1),
                ]
                endpoints = use_patch_vertices_uv
                patch_pad_trans = transforms.functional.perspective(
                                        patch_pad,
                                        startpoints=startpoints,
                                        endpoints=endpoints
                                        )
                patch_mask_pad_trans = transforms.functional.perspective(
                                        patch_mask_pad,
                                        startpoints=startpoints,
                                        endpoints=endpoints
                                        )
                image = torch.where(patch_mask_pad_trans > 0.5, patch_pad_trans, image)
    return image




def apply_patches_by_info_4side(info, image, instance_token_book, patch_book_4side, area_str):

    # 重新排序，按照画面深度uvd的d排序，以保证从远到近粘贴
    _d = info['objects_info']['uvd_center'][:, 2]
    d_index = _d.sort(descending=True)[1]

    d_index_list = d_index.tolist()

    for i_obj in d_index_list:
        if info['patch_visible_bigger'][i_obj]: # 可见
            assert info['patch_3d_temporal']['patch_vertices'][i_obj] != None
            use_patch_vertices_uv_4side = info['patch_3d_temporal']['patch_vertices'][i_obj][area_str]
            use_patch_vertices_uv_side_visible = info['patch_3d_temporal']['boxside_visiblie'][i_obj]

            for side_j in range(4):
                try:
                    if use_patch_vertices_uv_side_visible[side_j]:

                        use_patch_instance_token = info['objects_info']['instance_tokens'][i_obj]
                        # 注意！ 每一个token对应4个patch，所以，
                        # 找到index之后要*4，
                        # 再加 0,1,2,3
                        patch_4_start_idx = instance_token_book.index(use_patch_instance_token) * 4
                        use_patch_vertices_uv = use_patch_vertices_uv_4side[side_j][:,:2].cuda()

                        use_patch = patch_book_4side[patch_4_start_idx + side_j]
                        patch_use = use_patch
                        patch_h, patch_w = torch.tensor(patch_use.shape[-2:]).cuda().long()
                        img_h, img_w = torch.tensor(image.shape[-2:]).cuda().long()

                        if len(use_patch_vertices_uv.unique()) <= 4:
                            # 低级抠图
                            u1 = (use_patch_vertices_uv[:,0].min()).long()
                            u2 = (use_patch_vertices_uv[:,0].max()).long()
                            v1 = (use_patch_vertices_uv[:,1].min()).long()
                            v2 = (use_patch_vertices_uv[:,1].max()).long()

                            need_patch_h = v2 - v1
                            need_patch_w = u2 - u1

                            if need_patch_h > 0 and need_patch_w > 0:
                                patch_use = F.interpolate(patch_use[None], (need_patch_h, need_patch_w), mode='bilinear', align_corners=False)[0]
                                image = safe_img_patch_apply(
                                    image=image, 
                                    patch=patch_use, 
                                    x1=u1, 
                                    y1=v1
                                )
                        else:
                            endpoints = use_patch_vertices_uv
                            # 判断 endpoints 是否都在图像中，如果是，使用节约显存版本
                            # 留一点边
                            edge_width = 2
                            u_in = (edge_width < endpoints[:,0]) * (endpoints[:,0] < img_w - edge_width - 1)
                            v_in = (edge_width < endpoints[:,1]) * (endpoints[:,1] < img_h - edge_width - 1)

                            endpoints_allinimage = u_in.all() * v_in.all()


                            need_old_version = ~endpoints_allinimage

                            if need_old_version == False:
                                # 高级变换 (节约显存版本)
                                # 思路： 只使用小图像进行变换，变换后抠图粘到大图中
                                u_min = endpoints[:,0].min()
                                u_max = endpoints[:,0].max()
                                v_min = endpoints[:,1].min()
                                v_max = endpoints[:,1].max()

                                u_min_int = (u_min.floor() - 1).long()
                                u_max_int = (u_max.ceil()  + 1).long()
                                v_min_int = (v_min.floor() - 1).long()
                                v_max_int = (v_max.ceil()  + 1).long()

                                small_region_w = u_max_int - u_min_int
                                small_region_h = v_max_int - v_min_int

                                # patch粘贴区域，可能比patch更大，也可能更小，我们取大值进行绘图

                                small_region_w_share = max(small_region_w, patch_w)
                                small_region_h_share = max(small_region_h, patch_h)
                                u_max_int_real = u_min_int + small_region_w_share
                                v_max_int_real = v_min_int + small_region_h_share
                                
                                patch_pad = patch_use.new_zeros(3, small_region_h_share, small_region_w_share)
                                patch_pad[:,:patch_h,:patch_w] = patch_use

                                # 注意，如果patch更大，贴回到原图的时候，上下左右边界就都不同了，要重新计算

                                # 注意，如果patch更大，可能patch左上角对齐粘贴处，尺寸就又超出去了，需要回到老方法
                                need_old_version = False
                                if u_max_int_real > img_w - edge_width - 1 or v_max_int_real > img_h - edge_width - 1:
                                    need_old_version = True


                            if need_old_version == False:
                                patch_mask_pad = torch.zeros_like(patch_pad)
                                patch_mask_pad[:,:patch_h,:patch_w] = 1
                                startpoints = torch.cuda.FloatTensor([
                                    (0, 0), 
                                    (patch_w-1, 0),
                                    (patch_w-1, patch_h-1),
                                    (0, patch_h-1),
                                ])

                                endpoints_small_region = torch.zeros_like(endpoints)
                                endpoints_small_region[:,0] = endpoints[:,0] - u_min_int
                                endpoints_small_region[:,1] = endpoints[:,1] - v_min_int

                                patch_pad_trans = transforms.functional.perspective(
                                                        patch_pad,
                                                        startpoints=startpoints.cpu(),
                                                        endpoints=endpoints_small_region.cpu()
                                                        ) # 输出居然有可能 = 1 + 1.19e-7 大于1了！
                                patch_mask_pad_trans = transforms.functional.perspective(
                                                        patch_mask_pad,
                                                        startpoints=startpoints.cpu(),
                                                        endpoints=endpoints_small_region.cpu()
                                                        )
                                small_region = torch.where(
                                    patch_mask_pad_trans > 0.5, 
                                    patch_pad_trans, 
                                    image[:, v_min_int:v_max_int_real, u_min_int:u_max_int_real]
                                    )

                                image[:, v_min_int:v_max_int_real, u_min_int:u_max_int_real] = small_region.clamp(0,1)
                            else:
                                # 高级变换 (原始版本)
                                patch_pad = F.pad(
                                    patch_use, (0, img_w-patch_w, 0, img_h-patch_h),
                                    )
                                patch_mask_pad = torch.zeros_like(image)
                                patch_mask_pad[:,:patch_h,:patch_w] = 1
                                startpoints = torch.cuda.FloatTensor([
                                    (0, 0), 
                                    (patch_w-1, 0),
                                    (patch_w-1, patch_h-1),
                                    (0, patch_h-1),
                                ])
                                endpoints = use_patch_vertices_uv
                                patch_pad_trans = transforms.functional.perspective(
                                                        patch_pad,
                                                        startpoints=startpoints.cpu(),
                                                        endpoints=endpoints.cpu()
                                                        )
                                patch_mask_pad_trans = transforms.functional.perspective(
                                                        patch_mask_pad,
                                                        startpoints=startpoints.cpu(),
                                                        endpoints=endpoints.cpu()
                                                        )
                                image = torch.where(patch_mask_pad_trans > 0.5, patch_pad_trans, image).clamp(0,1)
                except:
                    print('error happend, skip this side\'s patch')
    return image



extend_common/patch_apply.py
zzj403/BEV_Robust
import os


def split_path_string_to_multiname(path_string):
    # 使用os.path.split来拆分路径并创建一个列表
    folder_list = []
    while True:
        head, tail = os.path.split(path_string)
        if not tail:
            break
        folder_list.insert(0, tail)  # 将文件夹名插入到列表的开头
        path_string = head
    return folder_list

if __name__ == '__main__':
    folder_list = split_path_to_multiname("./folder1/folder2/folder3/1.png")
    print(folder_list)

extend_common/path_string_split.py
zzj403/BEV_Robust
import time

def time_counter(last_time, motion_str, show_flag):
    now_time = time.time()
    if show_flag:
        print(motion_str, ':', str(round(now_time - last_time,7)),'s')
    last_time = time.time()
    return last_time

extend_common/time_counter.py
zzj403/BEV_Robust