cycle_train.py

import os
import torch
import numpy as np
import imageio

from config import initial
from dataset.load_llff import load_llff_data
from run_nerf import *

from view_convert import *

import cv2

import pdb

from generate_mask.mask_generate import *

import math

#os.environ["CUDA_VISIBLE_DEVICES"] = '0'
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#config规定种子
#np.random.seed(0)
DEBUG = False

torch.autograd.set_detect_anomaly(True)

def get_patch_rays(H, W, K, size, downsample, pose, valid_pixels):

    rays_o, rays_d = get_rays(H, W, K, torch.Tensor(pose)) #([1920, 1080, 3])
 
    # 确立凸多边形边界 y是竖着的，第0维
    points = np.argwhere(valid_pixels == 0)
    min_y, min_x = points.min(axis=0)
    max_y, max_x = points.max(axis=0)

    #在边界内选择一个正方形
    start_x_range = (min_x, max_x - size + 1)  #(610, 780)
    start_y_range = (min_y, max_y - size + 1)  #(471, 1352)

    if start_x_range[0] > start_x_range[1]:
        start_x = start_x_range[0]
    else:
        start_x = np.random.randint(start_x_range[0], start_x_range[1]+1)
    
    if start_y_range[0] > start_y_range[1]:
        start_y = start_y_range[1]
    else:
        start_y = np.random.randint(start_y_range[0], start_y_range[1]+1)
    
    rays_o = rays_o[start_y:start_y+size:downsample, start_x:start_x+size:downsample]
    rays_d = rays_d[start_y:start_y+size:downsample, start_x:start_x+size:downsample]
    valid_pixels = valid_pixels[start_y:start_y+size:downsample, start_x:start_x+size:downsample]

    patch_rays = torch.stack([rays_o, rays_d], 0)

    sc_mask = 1-valid_pixels # 0,1翻转

    return patch_rays, sc_mask

def turn2d(rgb, mask):
    # 根据mask，reshape成二维图像后，mask为0的点都改成0
    if rgb.shape[0] != mask.shape[0]*mask.shape[1] :
        raise ValueError("The length of rgb does not match the number of elements in the mask.")
    output = rgb.reshape([mask.shape[0], mask.shape[1], 3])
    output[mask==0] = 0
    return output

def train_one_nerf_step(args, i, render_kwargs_train, render_kwargs_test, optimizer, i_train, i_test, images, poses, H, W, K, valid_pixels, model_name, hwf, clip_model=None, target_embs=None, render_poses=None):
    global_step = i-1
    if args.no_batching:
        if i < 500 and args.dataset_type=='blender' and np.sum(valid_pixels[0] == 0)!=0 and args.use_sideposes: 
            #print("using side poses")
            #img_i = np.random.choice(np.concatenate([i_train[:2], i_train[-2:]])) # 这也是随机的
            # pdb.set_trace()
            img_i = np.random.choice([0])
        else:
            #在一开始先选择两边的训练
            img_i = np.random.choice(i_train)
        target = images[img_i]
        target = torch.Tensor(target).to(device)   #(H, W, 3)
        valid_pixel = valid_pixels[img_i]
        pose = poses[img_i, :3, :4]
        
        # N_rand 是 一个iter 抽取的光线数量
        if args.N_rand is not None:
            rays_o, rays_d = get_rays(H, W, K, torch.Tensor(pose))  # (H, W, 3), (H, W, 3)
             # 只从中心采样， 这个可以尝试用在blender中试一下先训练中心
            if i < args.precrop_iters:
                dH = int(H//2 * args.precrop_frac)
                dW = int(W//2 * args.precrop_frac)
                coords = torch.stack(
                    torch.meshgrid(
                        torch.linspace(H//2 - dH, H//2 + dH - 1, 2*dH), 
                        torch.linspace(W//2 - dW, W//2 + dW - 1, 2*dW)
                    ), -1)
                # 中心化的时候，pixel也要中心化 这个写死了只按0.5
                valid_pixel = valid_pixel[H // 2 - dH : H // 2 + dH, W // 2 - dW : W // 2 + dW]        
                if i == 5:
                        print(f"[Config] Center cropping of size {2*dH} x {2*dW} is enabled until iter {args.precrop_iters}") 
            else:
                coords = torch.stack(torch.meshgrid(torch.linspace(0, H-1, H), torch.linspace(0, W-1, W)), -1)  # (H, W, 2) 设置好了坐标


            coords = coords[valid_pixel != 0] # (valid_num, 2)
            #print("coords.shape:", coords.shape) 
            if coords.shape[0] < args.N_rand:
                select_inds = np.random.choice(coords.shape[0], size=[args.N_rand], replace=True) # 可以取相同元素
            else:
                select_inds = np.random.choice(coords.shape[0], size=[args.N_rand], replace=False)
            select_coords = coords[select_inds].long()
            rays_o = rays_o[select_coords[:, 0], select_coords[:, 1]]
            rays_d = rays_d[select_coords[:, 0], select_coords[:, 1]]
            batch_rays = torch.stack([rays_o, rays_d], 0)     # (2, N_rand, 3)
            target_s = target[select_coords[:, 0], select_coords[:, 1]]  # (N_rand, 3)

            #print(batch_rays.shape, target_s.shape)

    else:
        return
    
    torch.set_printoptions(threshold=np.inf)
    rgb, disp, acc, extras = render(H, W, K, chunk=args.chunk, rays=batch_rays,
                                                verbose=i < 10, retraw=True,
                                                **render_kwargs_train)

    optimizer.zero_grad()
    img_loss = img2mse(rgb, target_s)
    # print(img_loss)
    trans = extras['raw'][...,-1]
    loss = img_loss
    psnr = mse2psnr(img_loss)

    if math.isinf(psnr.item()):

        pdb.set_trace()

    if torch.isnan(loss).any():
        #torch.set_printoptions(threshold=np.inf)
        print("loss has nan:learning_rate",args.lrate)
        print("iter %d loss nan"%i)
        print(loss)
        print(rgb, target_s)
        print("batchrays:",batch_rays)
        return

    # extras中是一些不太在意的返回值
    if 'rgb0' in extras:
        img_loss0 = img2mse(extras['rgb0'], target_s)
        loss = loss + img_loss0
        psnr0 = mse2psnr(img_loss0)

    if args.use_semantic_loss and global_step%args.sc_loss_every == 0 and global_step>args.begin_sc_loss and global_step<args.stop_sc_loss:
        # 就用当前的img_i
        target_embedding = target_embs[img_i] # [1,512]
        # 光线是二维的，对应的rgb也可以是二维的，导出所有光线，然后渲染之后，再对mask以外置零吧
        sc_patch_rays, sc_mask =  get_patch_rays(H, W, K, size=args.random_ray_size, downsample=args.random_ray_downsample, pose=pose, valid_pixels=valid_i)

        # torch.Size([5625, 3])
        rgb, disp, acc, extras = render(H, W, K, chunk=args.chunk, rays=sc_patch_rays,
                                                verbose=i < 10, retraw=True,
                                                **render_kwargs_train)
        # import pdb
        # pdb.set_trace()
        f_image = rgb
        #f_image = torch.expand_dims(torch.transpose(f_image,[2,0,1]), 0)
        f_image = f_image.permute(2, 0, 1)
        f_image = f_image.unsqueeze(0)
        src_embedding = clip_model.get_image_features(pixel_values=clip_utils.preprocess_for_CLIP_tensor(f_image, sc_mask))
        #src_embedding /= torch.norm(src_embedding, p=2)# 这是就地操作，要避免这样的问题
        src_embedding_norm = src_embedding/torch.norm(src_embedding, p=2)
        sc_loss = 1 - torch.sum(src_embedding_norm * target_embedding)

        if 'rgb0' in extras:
            # c_image = turn2d(extras['rgb0'], sc_mask)
            c_image = extras['rgb0']
            c_image = c_image.permute(2, 0, 1)
            c_image = c_image.unsqueeze(0)
            # c_image = np.expand_dims(np.transpose(c_image,[2,0,1]), 0)
            src_embedding_f = clip_model.get_image_features(pixel_values=clip_utils.preprocess_for_CLIP_tensor(c_image, sc_mask))
            src_embedding_f_norm = src_embedding_f/torch.norm(src_embedding_f, p=2)  
            sc_loss += 1 - torch.sum(src_embedding_f_norm * target_embedding)

        loss += sc_loss*args.sc_loss_mult


    loss.backward()
    optimizer.step()

    decay_rate = 0.1 #0.1
    decay_steps = args.lrate_decay * 1000
    new_lrate = args.lrate * (decay_rate ** (global_step / decay_steps))
    for param_group in optimizer.param_groups:
        param_group['lr'] = new_lrate


    if i%args.i_weights==0:
        path = os.path.join(args.basedir, args.expname, model_name,'{:06d}.tar'.format(i))
        torch.save({
            'global_step': global_step,
            'network_fn_state_dict': render_kwargs_train['network_fn'].state_dict(),
            'network_fine_state_dict': render_kwargs_train['network_fine'].state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
        }, path)
        print('Saved checkpoints at', path)

    if i%args.i_video==0 and i > 0:
        # Turn on testing mode
        with torch.no_grad():
            #  在这里进行了修改，使得渲染角度为原本 pose 改为renderposes，看最终渲染效果
            if render_poses is not None:
                rgbs, disps = render_path(render_poses, hwf, K, args.chunk, render_kwargs_test)
            #训练可以挑光线训练，但渲染要统一渲染，因为不知道渲染出来的像素位置
            else:
                rgbs, disps = render_path(poses, hwf, K, args.chunk, render_kwargs_test)
        print('Done, saving', rgbs.shape, disps.shape)
        moviebase = os.path.join(args.basedir, args.expname, model_name, '{}_spiral_{:06d}_'.format(args.expname, i))
        imageio.mimwrite(moviebase + 'rgb.mp4', to8b(rgbs), fps=15, quality=8)
        imageio.mimwrite(moviebase + 'disp.mp4', to8b(disps / np.max(disps)), fps=30, quality=8)

    if (i%args.i_testset==0 and i > 0) or i == 500:
        testsavedir = os.path.join(args.basedir, args.expname, model_name, 'testset_{:06d}'.format(i))
        os.makedirs(testsavedir, exist_ok=True)
        print('test all poses shape', poses.shape)
        with torch.no_grad():
            render_path(torch.Tensor(poses).to(device), hwf, K, args.chunk, render_kwargs_test, gt_imgs=images, savedir=testsavedir)
        print('Saved test set')

    if i == 1:
        with open(os.path.join(args.basedir, args.expname, model_name, "train_log.txt"), "w") as log_file:
            log_file.write("")

    if i%args.i_print==0:
        log_message = f"[TRAIN] Iter: {i} Loss: {loss.item()}  PSNR: {psnr.item()}"

        # 使用 tqdm 进行输出
        tqdm.write(log_message)
        
        # 将相同的信息写入到文件
        with open(os.path.join(args.basedir, args.expname, model_name, "train_log.txt"), "a") as log_file:  # 'a' 表示追加模式
            log_file.write(log_message + "\n")  # 每条信息后加换行符
    return

# 对于单个物体，在总体nerf训练的基础上
def render_with_intermediate(args, poses, hwf, K, model_name, render_kwargs_test, valid=None, cycle=None):
    # both_valid存在时，只生成mask并集的像素；不存在时，就是渲染完整的

    # 将 渲染returnall相关函数转移过来，注意返回值的对齐
    with torch.no_grad():
        intermediate_base = os.path.join(args.basedir, args.expname, "intermediate_"+model_name)

        if cycle is not None:
            intermediate_base = intermediate_base + f"_cycle_{cycle}"

        if os.path.exists(intermediate_base):
            print("intermediate path exists, no more generate")
            return intermediate_base
        else:
            os.makedirs(intermediate_base, exist_ok=True)
            print("generating intermediate...")
        # 自动有None
        rgbs, disps, depths, zs, c2ws, K = render_path_returnall(poses, hwf, K, args.chunk, render_kwargs_test, valids=valid)
        print('Done, saving intermediate', rgbs.shape, disps.shape, zs.shape)

        rgbs_path = os.path.join(intermediate_base, "rgbs")
        disps_path = os.path.join(intermediate_base, "disps")
        depths_path = os.path.join(intermediate_base, "depths")
        zs_path = os.path.join(intermediate_base, "zs")
        poses_path = os.path.join(intermediate_base, "poses")

        os.makedirs(rgbs_path, exist_ok=True)
        os.makedirs(disps_path, exist_ok=True)
        os.makedirs(depths_path, exist_ok=True)
        os.makedirs(zs_path, exist_ok=True)
        os.makedirs(poses_path, exist_ok=True)

        for j in range(rgbs.shape[0]):
            imageio.imwrite(os.path.join(rgbs_path, '{:03d}.png'.format(j)), to8b(rgbs[j]))

            np.save(os.path.join(disps_path, '{:06d}.npy'.format(j)), disps[j])
            np.save(os.path.join(depths_path, '{:06d}.npy'.format(j)), depths[j])
            # 可视化保存
            imageio.imwrite(os.path.join(disps_path, '{:06d}.png'.format(j)), to8b(disps[j]/np.max(disps[j])))
            imageio.imwrite(os.path.join(depths_path, '{:06d}.png'.format(j)), to8b(depths[j]/np.max(depths[j])))
            #np.save(os.path.join(zs_path, '{:06d}.npy'.format(j)), zs[j])
            np.savetxt(os.path.join(poses_path, '{:06d}.txt'.format(j)), c2ws[j])  
        np.savetxt(os.path.join(intermediate_base, 'intrinsics.txt'), K) 
    
    print("saved intermediate")

    return intermediate_base


def train(args):
    # 训练包括若干个物体的单个nerf，以及整个场景的nerf（todo）
    K = None
    if args.dataset_type == 'llff':
        # 增加一项 gt_labels， 读入图像的语义信息, 蕴藏在npy文件里
        # images 格式 (61, 544, 960, 3)为0-1的浮点数, gt_labels格式为(61, 544, 960) 都是numpy格式
        images, poses, bds, render_poses, i_test, _ = load_llff_data(args.datadir, args.factor,
                                                                  recenter=True, bd_factor=.75,
                                                                  spherify=args.spherify, load_msk=False)

        hwf = poses[0,:3,-1]
        poses = poses[:,:3,:4]
        print('Loaded llff', images.shape, render_poses.shape, hwf, args.datadir)

        if not isinstance(i_test, list):
            i_test = [i_test]
        
        if args.llffhold > 0:
            print('Auto LLFF holdout,', args.llffhold)
            i_test = np.arange(images.shape[0])[::args.llffhold]

        i_val = i_test
        i_train = np.array([i for i in np.arange(int(images.shape[0])) if
                        (i not in i_test and i not in i_val)])
        
        print('DEFINING BOUNDS')
        # 暂时不用管
        if args.no_ndc: 
            near = np.ndarray.min(bds) * .9
            far = np.ndarray.max(bds) * 1.
             
        else:
            near = 0.
            far = 1.
        print('NEAR FAR', near, far)

    elif args.dataset_type == 'blender':
        images, poses, render_poses, hwf, i_split = load_blender_data(args.datadir, args.half_res, args.testskip, args.drop_middle_poses)

        # blender 的 poses
        print('Loaded blender', images.shape, render_poses.shape, hwf, args.datadir)

        if args.blender_pose_noise:
            print("add little noise to poses")
            poses = add_noise_to_poses(poses=poses)
        
        i_train, i_val, i_test = i_split


        i_test = i_train 

        near = 1.
        far = 9.

        # 记录下来透明度不为0的像素点
        blender_valid = (images[..., -1] != 0).astype(int)

        if args.white_bkgd:
            images = images[...,:3]*images[...,-1:] + (1.-images[...,-1:])
        else:
            images = images[...,:3]

    # Cast intrinsics to right types
    H, W, focal = hwf
    H, W = int(H), int(W)
    hwf = [H, W, focal]

    if K is None:
        K = np.array([
            [focal, 0, 0.5*W],
            [0, focal, 0.5*H],
            [0, 0, 1]
        ])

    if args.render_test:
        render_poses = np.array(poses[i_test])

    basedir = args.basedir
    expname = args.expname
    os.makedirs(os.path.join(basedir, expname), exist_ok=True)
    f = os.path.join(basedir, expname, 'args.txt')
    with open(f, 'w') as file:
        for arg in sorted(vars(args)):
            attr = getattr(args, arg)
            file.write('{} = {}\n'.format(arg, attr))
    if args.config is not None:
        f = os.path.join(basedir, expname, 'config.txt')
        with open(f, 'w') as file:
            file.write(open(args.config, 'r').read())

    bds_dict = {
        'near' : near,
        'far' : far,
    }

    # 只保存 render poses
    np.save(os.path.join(args.datadir, 'render_poses.npy'), render_poses)
    print("saved render_poses:", render_poses.shape)

    # 只有全景 和 去掉遮挡物 两个nerf， 先训全景 创建nerf这里要改
    scene_model_name = "scene_nerf"
    scene_model_dir = os.path.join(args.basedir, args.expname, scene_model_name)
    os.makedirs(scene_model_dir, exist_ok=True)
    scene_render_kwargs_train, scene_render_kwargs_test, scene_start, scene_grad_vars, scene_optimizer = create_nerf(args, scene_model_name, device)
    print("加载scene模型 start:", scene_start)
    scene_render_kwargs_train.update(bds_dict)
    scene_render_kwargs_test.update(bds_dict)
    scene_start = scene_start + 2 # 加2接着训不会有bug

    # 训练的参数设置
    N_rand = args.N_rand
    use_batching = not args.no_batching
    poses = torch.Tensor(poses).to(device)
    render_poses = torch.Tensor(render_poses).to(device)

    
    # 如果只渲染，对于scene 来说，权当渲染一下，看看当前模型的训练水平, 加个参数吧
    if args.render_scene and args.render_only:
        with torch.no_grad():
            # 这里的render，就是用的处理后的renderposes

            print('SCENE RENDER ONLY')

            savedir = os.path.join(basedir, expname, scene_model_name, 'renderonly')
            os.makedirs(savedir, exist_ok=True)
            print('test poses shape', render_poses.shape, render_poses[1])

            rgbs, disps = render_path(render_poses, hwf, K, args.chunk, scene_render_kwargs_test)
            print('Done, saving returnall', rgbs.shape, disps.shape)
            
            imageio.mimwrite(os.path.join(savedir, 'rgb.mp4'), to8b(rgbs), fps=15, quality=8)
            imageio.mimwrite(os.path.join(savedir, 'disp.mp4'), to8b(disps / np.max(disps)), fps=30, quality=8)

        return

    scene_N_iters = 100000 + 1
    print('Begin Scene')
    print('TRAIN views are', i_train)
    print('TEST views are', i_test)
    print('VAL views are', i_val)

    # 对于scene_nerf 使用 全为1的mask
    scene_valid = np.ones(images.shape[:3])

    global_step = scene_start
    for i in trange(global_step, scene_N_iters):
        time0 = time.time()

        train_one_nerf_step(args, i, scene_render_kwargs_train, scene_render_kwargs_test, scene_optimizer
                        , i_train, i_test, images, poses, H, W, K, scene_valid, scene_model_name, hwf, render_poses=render_poses)

        dt = time.time()-time0

        global_step += 1

    # 训练完后，生成一下rgb和深度，用于和点的convert
    time_before_render_scene = time.time()
    scene_intermediate_path = render_with_intermediate(args, poses, hwf, K, scene_model_name ,render_kwargs_test=scene_render_kwargs_test)
    time_after_render_scene = time.time()
    print(f"render scene time : {time_after_render_scene - time_before_render_scene}")

    print("hello!!!!!!!!!!!!!")
    time_before_maskpro = time.time()
    # 对于mask的提示点
    mask_coords_idx = args.mask_coords_idx
    # 虽然factor是llff得参数，但这里为了统一，修改factor默认为1， half_res时也应该/2
    if args.half_res:
        args.factor = 2
    if args.occ_mask_coords:
        occ_mask_coords = np.array(args.occ_mask_coords)//args.factor
        print(f"occ_mask_coords:{occ_mask_coords}")
    else:
        occ_mask_coords = np.array([])
        print("no occ_mask_coords")

    if args.occed_mask_coords:
        occed_mask_coords = np.array(args.occed_mask_coords)//args.factor
        print(f"occed_mask_coords:{occed_mask_coords}")
    else:
        occed_mask_coords = np.array([])
        print("no occed_mask_coords")

    mask_path = os.path.join(basedir, expname, "mask")
    masks = []
    # 目录存在且不为空 现在mask中的点包括0，1（遮挡物），2（被遮挡物）
    # 只能写成大if else 来暂时区分 blneder 和 llff 的mask了 
    if args.dataset_type == 'llff':
        if os.path.exists(mask_path) and bool(os.listdir(mask_path)):
            mask_files = sorted([f for f in os.listdir(mask_path) if f.endswith('.npy')])
            masks = [np.load(os.path.join(mask_path, f)) for f in mask_files]
            masks = np.stack(masks, axis=0)
            if masks.shape == images.shape[:3]:
                print(f"load masks:{masks.shape}")
            else:
                # 可以对提示点进行补救
                segmentor = ImageSegmentor()
                occ_labels = np.ones(occ_mask_coords.shape[0], dtype=np.uint8)
                occed_labels = np.full(occed_mask_coords.shape[0], 2, dtype=np.uint8)
                mask_idx = segmentor.create_combined_mask(images[mask_coords_idx],occ_points=occ_mask_coords, occ_labels=occ_labels
                                                    , occed_points=occed_mask_coords, occed_labels=occed_labels)
                #增加一个判断 如果 occed_mask_coords为空 则所有mask不为1的地方都为2
                if occed_mask_coords.shape[0]==0:
                    mask_idx[mask_idx!=1] = 2
                np.save(os.path.join(mask_path, f"{mask_coords_idx:06d}.npy"), mask_idx)
                visualize_mask_on_image(images[mask_coords_idx], mask_idx, np.vstack((occ_mask_coords, occed_mask_coords)), os.path.join(mask_path, f"{mask_coords_idx:06d}.png"))
        else:
            # 要做很多工作
            os.makedirs(mask_path, exist_ok=True)
            
            segmentor = ImageSegmentor()

            depth_path = os.path.join(scene_intermediate_path, "depths")

            # import pdb
            # pdb.set_trace()

            #先整出源视图的mask
            occ_labels = np.ones(occ_mask_coords.shape[0], dtype=np.uint8)
            occed_labels = np.full(occed_mask_coords.shape[0], 2, dtype=np.uint8)

            # 先得到原始图的mask
            mask_idx = segmentor.create_combined_mask(images[mask_coords_idx],occ_points=occ_mask_coords, occ_labels=occ_labels
                                                    , occed_points=occed_mask_coords, occed_labels=occed_labels)
            
            # 这里也先暂时一样 
            if occed_mask_coords.shape[0]==0:
                    mask_idx[mask_idx!=1] = 2
            
            # 将从原始图像开始，递归处理周围视角
            masks = [None] * images.shape[0]
            masks[mask_coords_idx] = mask_idx

            np.save(os.path.join(mask_path, f"{mask_coords_idx:06d}.npy"), mask_idx)
            visualize_mask_on_image(images[mask_coords_idx], mask_idx, np.vstack((occ_mask_coords, occed_mask_coords)), os.path.join(mask_path, f"{mask_coords_idx:06d}.png"))
            

            process_view(poses, images, mask_coords_idx, masks, segmentor, mask_path, depth_path, K, args.occ_prompt_points, args.occed_prompt_points, occ_mask_coords)
    
    # 只用判断遮挡物的范围，直接用标记点+直接warp
    elif args.dataset_type == 'blender':
        if os.path.exists(mask_path) and bool(os.listdir(mask_path)):
            mask_files = sorted([f for f in os.listdir(mask_path) if f.endswith('.npy')])
            masks = [np.load(os.path.join(mask_path, f)) for f in mask_files]
            masks = np.stack(masks, axis=0)
            if masks.shape != images.shape[:3]:
                print("Shapes do not match. Masks shape:", masks.shape, "Images shape:", images.shape)
            else:
                print(f"load masks:{masks.shape}")
        else:
            os.makedirs(mask_path, exist_ok=True)
            masks = []
            segmentor = ImageSegmentor()
            occ_labels = np.ones(occ_mask_coords.shape[0], dtype=np.uint8)
            # 源视图的深度
            depth_idx = np.load(os.path.join(scene_intermediate_path, "depths", f"{mask_coords_idx:06d}.npy"))
            c2w_idx = poses[mask_coords_idx].cpu().numpy()
            if c2w_idx.shape[0] == 3:
                c2w_idx = np.concatenate((c2w_idx, np.array([[0,0,0,1]])), axis=0)
            c2w_idx = convert_pose(c2w_idx)
            w2c_idx = np.linalg.inv(c2w_idx)
            mask_idx, _ = segmentor.segment_image(images[mask_coords_idx], occ_mask_coords, occ_labels, args.separate_mask)
            for i in range(images.shape[0]):
                if i != mask_coords_idx:
                    c2w_i = poses[i].cpu().numpy()
                    if c2w_i.shape[0] == 3:
                        c2w_i = np.concatenate((c2w_i, np.array([[0,0,0,1]])), axis=0)
                    c2w_i = convert_pose(c2w_i)
                    w2c_i = np.linalg.inv(c2w_i)
                    #occ_target_coords = np.stack([trans_3d_points(c[0],c[1],depth_idx[c[1],c[0]],c2w_idx,w2c_i,K)[0] for c in occ_mask_coords], axis=0)
                    # 如果有-1 就保持不变
                    occ_target_coords = np.stack([
                        trans_3d_points(c[0], c[1], depth_idx[c[1], c[0]], c2w_idx, w2c_i, K)[0] if c[0] != -1 and c[1] != -1 else np.array(c)
                        for c in occ_mask_coords
                    ], axis=0)
                    mask_i,_ = segmentor.segment_image(images[i], occ_target_coords, occ_labels, args.separate_mask)
                    # 这里返回的mask_i是bool
                else:
                    mask_i = mask_idx
                    occ_target_coords = occ_mask_coords
                # 排除法获得被遮挡物的mask
                mask_i = mask_i.astype(np.int8)
                # 将mask改造为int，可以有1，2
                mask_i[(blender_valid[i] == 1) & (mask_i == 0)] = 2
                # import pdb
                # pdb.set_trace()
                masks.append(mask_i)
                np.save(os.path.join(mask_path, f"{i:06d}.npy"), mask_i)
                visualize_mask_on_image(images[i], mask_i, occ_target_coords, os.path.join(mask_path, f"{i:06d}.png"))
    # 对遮挡物mask扩张, 然后反转 之后再说 
    masks = np.stack(masks, axis=0)
    # 将mask解析成 遮挡mask和被遮挡mask，首先还得将遮挡mask 膨胀一下
    occ_masks = (masks == 1).astype(np.uint8)

    # 感觉分割的总体还是比较好的，尤其是遮挡物，不用膨胀太多
    kernel = np.ones((3, 3), np.uint8)
    for i in range(occ_masks.shape[0]):
        occ_masks[i] = cv2.dilate(occ_masks[i], kernel, iterations=3)
        #不膨胀做个对比
        pass

    
    # 剔除遮挡mask膨胀的部分
    occed_masks = ((masks == 2) & (occ_masks == 0)).astype(np.uint8)

    # 这里指参与训练donerf的像素
    valid_train_pixel = np.where(occ_masks==0, 1, 0)
    if args.only_train_occed:
        valid_train_pixel = occed_masks

    time_after_maskpro = time.time()
    print(f"time of mask pro : {time_after_maskpro - time_before_maskpro}")


    # 创建do_nerf
    do_model_name = "do_nerf"
    do_model_dir = os.path.join(args.basedir, args.expname, do_model_name)
    os.makedirs(do_model_dir, exist_ok=True)
    do_render_kwargs_train, do_render_kwargs_test, do_start, do_grad_vars, do_optimizer = create_nerf(args, do_model_name, device)
    print("加载do模型 start:", do_start)
    do_render_kwargs_train.update(bds_dict)
    do_render_kwargs_test.update(bds_dict)
    do_start = do_start + 2 # 加2接着训不会有bug
    
    do_N_iters = 200000 + 1
    print('Begin Do')
    print('TRAIN views are', i_train)
    print('TEST views are', i_test)
    print('VAL views are', i_val)

    global_step = do_start - 1
    for i in trange(global_step, do_N_iters):
        time0 = time.time()

        if args.dataset_type == 'blender' and args.train_occed_first and i < 500:
            train_one_nerf_step(args, i, do_render_kwargs_train, do_render_kwargs_test, do_optimizer
                            , i_train, i_test, images, poses, H, W, K, occed_masks, do_model_name, hwf, render_poses=render_poses)
        else:
            train_one_nerf_step(args, i, do_render_kwargs_train, do_render_kwargs_test, do_optimizer
                            , i_train, i_test, images, poses, H, W, K, valid_train_pixel, do_model_name, hwf, render_poses=render_poses)
        
        if global_step % args.i_cycle == 0:
            intermediate_except_path = render_with_intermediate(args, poses, hwf, K, do_model_name, do_render_kwargs_test, cycle=global_step//args.i_cycle)
            
            # 这里对除了遮挡物的所有范围都进行warp，其实是不可行的
            # 通过warp不断补全
            images, occed_masks = view_convert(intermediate_except_path, occed_masks, images=images,cycle=global_step//args.i_cycle)
            valid_train_pixel = valid_train_pixel | occed_masks
            if args.only_train_occed:
                valid_train_pixel = occed_masks

        dt = time.time()-time0
        global_step += 1
    return


    # 对每一个物体，单独记录图片，设有S个物体
    single_imgs = []

    # 对每一个物体，单独记录图片中的每个像素是否有效  (S, N, H , W )
    # 直接默认0是遮挡物， 1 是被遮挡， 2是全1
    single_valid = []

    # 记录以下所有物体的有效范围，不至于出现两个相重合物体的mask被腐蚀后出现间隔
    both_valid = np.zeros_like(gt_labels)
    print("!!!!!!!!!!!!!\nboth_valid size: ", both_valid.shape)
    both_valid[gt_labels != 0] = 1
    kernel_size = 3
    kernel = np.ones((kernel_size, kernel_size), np.int8)
    for i in range(both_valid.shape[0]):
        both_valid[i] = cv2.erode(both_valid[i], kernel, iterations=2)

    # 给图片加上透明度，然后分别写入
    for i in range(args.class_num):

        
        alpha = gt_labels.copy()
        alpha[alpha != i+1] = 0
        alpha[alpha == i+1] = 1

        # 对这个mask腐蚀缩减一点
        eroded_alpha = np.zeros_like(alpha)
        kernel_size = 5
        for ii in range(eroded_alpha.shape[0]):
            a = alpha[ii]
            kernel = np.ones((kernel_size, kernel_size), np.int8)
            eroded_a = cv2.erode(a, kernel, iterations=5)
            eroded_alpha[ii] = eroded_a
        alpha = eroded_alpha
        
        # 对每个物体的有效像素位
        single_valid.append(alpha.copy())
        alpha = np.expand_dims(alpha, axis=3)  #使得维度与img相同
        #print(np.where(alpha == 1))
        alpha[alpha==1] = 255 #获得不透明度
        #print("alpha shape", alpha.shape)
        
        #print("sin_imgs[0] shape", type(sin_imgs[0,:,:,:3]))

        # 保存一下初始的样子，也合理。
        savedir = os.path.join(args.basedir, args.expname, "single_img_class{:02d}".format(i))

        if savedir is not None and not os.path.exists(savedir):
            sin_imgs = np.concatenate((images*255, alpha), axis=3) #将不透明度写入第四个通道,这里的images归一化了
            os.makedirs(savedir, exist_ok=True)
            for j in range(sin_imgs.shape[0]):
                filename = os.path.join(savedir, '{:03d}.png'.format(j))
                imageio.imwrite(filename, np.uint8(sin_imgs[j]))

    # 现在认为第一个模型是训练除了遮挡物之外的所有Nerf，将mask倒置。对于每个元素，如果条件为真，则从第二个参数中选择值；否则，从第三个参数中选择
    # 这样直接转换会训练到腐蚀过的边缘, 给膨胀一下,再反转
    occ_masks = gt_labels.copy()
    occ_masks[occ_masks == 2] = 0
    kernel = np.ones((3, 3), np.uint8)
    for ii in range(occ_masks.shape[0]):
        occ_masks[ii] = cv2.dilate(occ_masks[ii], kernel, iterations=3)
    single_valid[0] = np.where(occ_masks==0, 1, 0)

    #single_valid[0] = np.where(single_valid[0]==0, 1, 0)

    # 给全景Nerf一个全1的valid
    single_valid.append(np.ones(single_valid[-1].shape,dtype=int))

    
    
    
    # 创建多个nerf模型,实际上，这里只创建第一个
    render_kwargs_train_s = []
    render_kwargs_test_s = []
    starts = []  #记录每个模型已经训练的次数
    optimizer_s = []

    # 只创建exceptnerf
    for i in range(1):
        class_model_dir = os.path.join(args.basedir, args.expname, "class{:02d}".format(i))
        os.makedirs(class_model_dir, exist_ok=True)

        # 只看渲染效果的话，其它模型没必要加载,省点内存
        render_kwargs_train, render_kwargs_test, start_i, grad_vars, optimizer = create_nerf(args, i, device)

        # 如果是被遮挡物体的模型，那么训练的时候就不加噪声（除了遮挡物体之外的Nerf，也不加噪声试试
        if i in [1]:
            render_kwargs_train["raw_noise_std"] = 0
        print("加载模型 start_{}:".format(i), start_i)

        bds_dict = {
        'near' : near,
        'far' : far,
        }
        render_kwargs_train.update(bds_dict)
        render_kwargs_test.update(bds_dict)

        render_kwargs_train_s.append(render_kwargs_train)
        render_kwargs_test_s.append(render_kwargs_test)

        # 记录模型下次训练的次数，只有+2不出bug，如果是新模型，也不差这一次
        starts.append(start_i+2)
        optimizer_s.append(optimizer)

    N_rand = args.N_rand
    use_batching = not args.no_batching

    #print("render_kwargs_train_s",render_kwargs_train_s)

    poses = torch.Tensor(poses).to(device)

    if args.render_only:

        with torch.no_grad():
            render_poses = torch.Tensor(render_poses).to(device)

            print('RENDER ONLY')

            savedir = os.path.join(basedir, expname, 'class01', 'renderonly_occled')
            os.makedirs(savedir, exist_ok=True)
            print('test poses shape', render_poses.shape, render_poses[1])

            rgbs, disps = render_path(render_poses, hwf, K, args.chunk, render_kwargs_test_s[1])
            print('Done, saving returnall', rgbs.shape, disps.shape)
            
            imageio.mimwrite(os.path.join(savedir, 'rgb.mp4'), to8b(rgbs), fps=15, quality=8)
            imageio.mimwrite(os.path.join(savedir, 'disp.mp4'), to8b(disps / np.max(disps)), fps=30, quality=8)

        return
    
    # 初始化clip模型
    if args.use_semantic_loss:
        clip_model, _ = clip_utils.init_CLIP(args.clip_model_name, device=device)
        print(f'semantic loss ACTIVATED, CLIP is set up '
              f'(sc_loss_mult: {args.sc_loss_mult})')
    else:
        clip_model = None
        print('semantic loss DEACTIVATED, CLIP is set to None')
        

    # 总轮数还是200000
    N_iters = 200000 + 1
    # N_iters = 50000 + 1
    # 只有训练了20000 之后，才会保存tar
    print('Begin')
    print('TRAIN views are', i_train)
    print('TEST views are', i_test)
    print('VAL views are', i_val)

    # 提取训练图像中被遮挡物的语义信息
    target_embs = []
    if args.use_semantic_loss and clip_model is not None:
        for i,img in enumerate(images):
            #pdb.set_trace()
            img = np.expand_dims(np.transpose(img,[2,0,1]), 0)
            emb = clip_model.get_image_features(pixel_values = clip_utils.preprocess_for_CLIP(img, single_valid[1][i]))
            normalized_emb = emb / torch.norm(emb, p=2)
            target_embs.append(normalized_emb.detach())
            # target_embs.append(emb/torch.norm(emb, p=2)) 这里作为监督信号，不参与反向传播
        # target_embs = np.concatenate(target_embs, 0)  这是Tensor，就不做这些了

    # 一般认为有了tar后，对应的图像已经生成了，不需要再生成一次了
    # 从 199999 到 200001
    global_step = min(starts)
    for i in trange(global_step, N_iters):
        # print("now iter:",i, "final iter:", N_iters-1)
        # 加载200000的模型tar后，还要进来一下，最后一轮进行生成， 取消了

        time0 = time.time()

        for class_i in range(1):
            # 特殊处理，选择只训练的模型
            if class_i in [0,1,2]:
                # 如果模型当前进度刚好是，才继续训练
                if starts[class_i] == global_step:
                    train_one_nerf_step(args, i, render_kwargs_train_s[class_i], render_kwargs_test_s[class_i], optimizer_s[class_i]
                                , i_train, i_test, images, poses, H, W, K, single_valid[class_i], class_i, hwf, clip_model=clip_model, target_embs=target_embs)
                    starts[class_i]+=1

        if global_step%args.i_cycle == 0 :
            except_i = 0
            occed_i = 1
            intermediate_except_path = render_with_intermediate(args, poses, hwf, K, class_i=except_i,render_kwargs_test=render_kwargs_test_s[except_i],cycle=global_step//args.i_cycle)
            images, single_valid[occed_i] = view_convert(intermediate_except_path, single_valid, images=images,cycle=global_step//args.i_cycle)
            #相当于更新了except nerf的有效训练像素位
            single_valid[except_i] = single_valid[except_i] | single_valid[occed_i]

        dt = time.time()-time0


        global_step += 1

        #return #测试时用，只进行一轮
    


    return
    except_i = 0
    intermediate_except_path = render_with_intermediate(args, poses, hwf, K, class_i=except_i,render_kwargs_test=render_kwargs_test_s[except_i])

    
    occed_i = 1

    complete_i = 2

    intermediate_complete_path = render_with_intermediate(args, poses, hwf, K, class_i=complete_i, render_kwargs_test=render_kwargs_test_s[complete_i])
    

    # 更新了images之后，再训10000次 部分模型
    N_iters_total = 200000 + 1
    for i in trange(global_step, N_iters_total):
        time0 = time.time()
        for class_i in range(0, args.class_num + 1):
            # 特殊处理，只训练被遮挡物
            #if class_i == 1:
            if class_i == 1:
                train_one_nerf_step(args, i, render_kwargs_train_s[class_i], render_kwargs_test_s[class_i], optimizer_s[class_i]
                                , i_train, i_test, images, poses, H, W, K, single_valid[class_i], class_i, hwf)


        dt = time.time()-time0


        global_step += 1
    




if __name__ == '__main__':

    args = initial()
    np.random.seed(args.random_seed)
    torch.manual_seed(args.random_seed)
    train(args)