content: import os
  import numpy as np
  import torch
  import nvdiffrast.torch as dr
  import json
  import torch.nn.functional as F
  from PIL import Image
  import pymeshlab
  import cv2
  
  def back_to_texture(glctx, look_at, pos, tri, tex, uv, uv_idx, idx, vn):
      rast_out, rast_out_db = dr.rasterize(glctx, pos, tri, resolution=[tex.shape[0],tex.shape[1]])
      gb_normal, _ = dr.interpolate(vn[None], rast_out, tri)
      gb_normal = F.normalize(gb_normal, dim=-1)
      if idx == 2 or idx == 0:
          filter_camera = [torch.tensor([[[[1,0.,0.]]]]).cuda(), torch.tensor([[[[-1,0.,0.]]]]).cuda()]
      else:
          filter_camera = [torch.tensor([[[[0,-1.,0.]]]]).cuda(), torch.tensor([[[[0,1.,0.]]]]).cuda()]
      nmasks = []
      for fc in filter_camera:
          nmasks.append(((gb_normal * fc) > 0.75).int().sum(keepdim=True, dim=-1))
      gb_normal_mask = 1 - (nmasks[0] | nmasks[1])
     #Image.fromarray(np.clip(gb_normal_mask[0,...,0].detach().cpu().numpy() * 255, 0, 255).astype(np.uint8)).save(f"mask_normal_{idx}.png")
      gb_mask = rast_out[...,3:4] > 0
      tri_list = torch.unique(rast_out[...,3:4].reshape(-1))
      tri_list = (tri_list[1:] - 1).to(torch.int32)
      pos = pos[0]
  
      depth_map = rast_out[...,3:4].clone()
      depth_map[depth_map > 0] = 1
      depth_map = depth_map.to(torch.float32)
      dmax = (rast_out[...,2:3] * gb_mask).max()
      uv = torch.cat([uv * 2 - 1, torch.zeros(uv.shape[0], 1).cuda(), torch.ones(uv.shape[0], 1).cuda()], dim=1).unsqueeze(0)
      uv_idx = uv_idx[tri_list.to(torch.long)]
      rast_uv, rast_uv_db = dr.rasterize(glctx, uv, uv_idx, resolution=(1024, 1024))
      pos_clip = torch.cat([pos[...,:2], pos[...,3:]], -1)
      pos_2d, _ = dr.interpolate(pos_clip, rast_uv, tri[tri_list.to(torch.long)]) # pos (x, y, z, w)
      pos_coord = (pos_2d[...,:2] / (pos_2d[...,2:3] + 1e-6) + 1) / 2.
      texture_mask = (rast_uv[...,3:4] > 0).int()
      color = dr.texture(tex[None, ...] * gb_normal_mask, pos_coord, filter_mode='linear')
      color_mask = dr.texture(gb_normal_mask.to(torch.float32), pos_coord, filter_mode='linear')
      color_mask[color_mask > 0.82] = 1
      color_mask[color_mask <= 0.82] = 0
      color_mask = color_mask.to(torch.int32)
     #Image.fromarray(np.clip(color_mask[0].repeat(1,1,3).detach().cpu().numpy() * 255, 0, 255).astype(np.uint8)).save(f"depth_{idx}.png")
      texture_mask = texture_mask * color_mask
     #Image.fromarray(np.clip(color[0].detach().cpu().numpy() * 255, 0, 255).astype(np.uint8)).save(f"{idx}.png")
     #Image.fromarray(np.clip(texture_mask[0].repeat(1,1,3).detach().cpu().numpy() * 255, 0, 255).astype(np.uint8)).convert("RGB").save(f"mask-{idx}.png")
      return color, texture_mask, rast_uv
  
  def perspective(fovy=0.6913, aspect=1.0, n=0.1, f=1000.0, device=None):
      y = np.tan(fovy / 2)
      return torch.tensor([[1/(y*aspect),    0,            0,              0], 
                           [           0, 1/-y,            0,              0], 
                           [           0,    0, -(f+n)/(f-n), -(2*f*n)/(f-n)], 
                           [           0,    0,           -1,              0]]).to(torch.float32).cuda()
  
  def rec_mvp(trans, h, w):
      mv = trans
      fov = 40. / 180. * np.pi
      proj = perspective(fov, h / w, n=0.1, f=1000)
      mvp = proj @ mv
      return mvp
  
  def aggregate_texture(kd_map, textures, texture_masks, rast_uvs):
      texture = torch.zeros_like(textures[0])
      texture_mask = torch.zeros_like(texture_masks[0])
      ctex = []
      for idx in range(len(textures)):
          ctex.append(textures[idx] * texture_masks[idx] + 10 * (1 - texture_masks[idx]))
      cat_textures = torch.stack(ctex, dim=-2)
      dis_measure = (cat_textures - kd_map.unsqueeze(-2)).abs().sum(-1)
      _, choose_idx = dis_measure.min(-1)
  
      choose_idx = choose_idx.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, 1, 1, 3)
      final_texture_map = torch.gather(cat_textures, 3, choose_idx).squeeze(-2)
      #cv2.imwrite("final_texture_map.png", cv2.cvtColor((final_texture_map[0].cpu().numpy() * 255).astype(np.uint8), cv2.COLOR_BGR2RGB))
      #cv2.imwrite("final_texture_mask.png", (texture_mask[0].cpu().numpy() * 255).astype(np.uint8))
      zero_mask = (final_texture_map.max(dim=-1, keepdim=True)[0] > 0.1)
      close_mask = ((final_texture_map[0] - kd_map).abs().sum(dim=-1, keepdim=True) < 1.0).int()
      for idx in range(len(textures)):
          texture += textures[idx] * texture_masks[idx]
          texture_mask |= texture_masks[idx]
      texture_mask = texture_mask * zero_mask * close_mask[None]
      optimize_mask = (texture_mask == 0).int()
  
     #import pdb; pdb.set_trace()
     #mask = (texture_mask[0].cpu().numpy() * 255).astype(np.uint8)
     #cv2.imwrite("mask.png", mask)
     #kernel = np.ones((5,5), np.uint8)
     #dilated = cv2.dilate(mask, kernel, iterations=1)
     #cv2.imwrite("di_mask.png", dilated)
     #texture_mask[0] = torch.from_numpy(dilated).unsqueeze(-1).to(torch.float32) / 255.
  
      final_texture_map = final_texture_map[0] * texture_mask[0]
      Image.fromarray(np.rint(final_texture_map.cpu().numpy() * 255).astype(np.uint8)).save(f"final_texture.png")
  
     #cv2.imwrite("kd_map.png", cv2.cvtColor((kd_map.cpu().numpy() * 255).astype(np.uint8), cv2.COLOR_BGR2RGB))
     #cv2.imwrite("texture_map.png", cv2.cvtColor((final_texture_map.cpu().numpy() * 255).astype(np.uint8), cv2.COLOR_BGR2RGB))
     #result = cv2.seamlessClone((final_texture_map.cpu().numpy() * 255).astype(np.uint8), (kd_map.cpu().numpy() * 255).astype(np.uint8), mask, (mask.shape[1]//2, mask.shape[0]//2), cv2.NORMAL_CLONE)
     #cv2.imwrite("result.png", cv2.cvtColor(result * 255, cv2.COLOR_BGR2RGB))
  
      kd_map = kd_map * (1 - texture_mask[0]) + final_texture_map
      return kd_map, optimize_mask
  
  def refine(save_path, front_image, back_image, left_image, right_image):
      ms = pymeshlab.MeshSet()
      mesh_path = f"{save_path}/model-00.obj"
      ms.load_new_mesh(mesh_path)
      ms.apply_coord_laplacian_smoothing(stepsmoothnum=10)
      tl = open(mesh_path, "r").readlines()
      tex_uv = []
      uv_idx = []
      for line in tl:
          if line.startswith("vt"):
              uvs = line.split(" ")[1:3]
              tex_uv += [float(uvs[0]), 1.0-float(uvs[1])]
      tex_uv = torch.from_numpy(np.array(tex_uv)).to(torch.float32).cuda().reshape(-1, 2)
      m = ms.current_mesh()
      v_matrix = m.vertex_matrix()
      f_matrix = m.face_matrix()
      vn = m.vertex_normal_matrix()
      uv_idx = torch.arange(f_matrix.shape[0] * 3).reshape(-1, 3).to(torch.int32).cuda()
      vn = torch.tensor(vn).contiguous().cuda().to(torch.float32)
  
      frames = []
      front_camera = torch.tensor([[
          1,0,0,0,
          0,0,1,0,
          0,-1,0,-1.5,
          0,0,0,1,
      ]]).to(torch.float32).reshape(4,4).cuda()
      back_camera = torch.tensor([[
          1,0,0,0,
          0,0,1,0,
          0,1,0,-1.5,
          0,0,0,1,
      ]]).to(torch.float32).reshape(4,4).cuda()
      right_camera = torch.tensor([[
          0,-1,0,0,
          0,0,1,0,
          1,0,0,-1.5,
          0,0,0,1,
      ]]).to(torch.float32).reshape(4,4).cuda()
      left_camera = torch.tensor([[
          0,1,0,0,
          0,0,1,0,
          -1,0,0,-1.5,
          0,0,0,1,
      ]]).to(torch.float32).reshape(4,4).cuda()
      frames = [front_camera, left_camera, back_camera, right_camera]
  
      target_images = []
      for target_image in [front_image, left_image, back_image, right_image]:
          target_images.append(torch.from_numpy(np.asarray(target_image.convert("RGB"))).to(torch.float32).cuda() / 255.)
  
      pos = torch.tensor(v_matrix, dtype=torch.float32).contiguous().cuda()
      tri = torch.tensor(f_matrix, dtype=torch.int32).contiguous().cuda()
  
      kd_map = (torch.tensor(np.asarray(Image.open(f"{save_path}/texture_kd.jpg"))) / 255.).cuda()
      translate_tensor = torch.zeros((1,1,3)).cuda()
      pos = torch.cat([pos, torch.ones([pos.shape[0], 1]).cuda()],-1).unsqueeze(0)
      glctx = dr.RasterizeCudaContext()
      target_texture = []
      target_mask = []
      rast_uvs = []
      with torch.no_grad():
          for idx, trans in enumerate(frames):
              target_image = target_images[idx]
              look_at = -torch.linalg.inv(trans)[:3,2]
              mvp = rec_mvp(trans, h=target_images[0].shape[0], w=target_images[0].shape[1])
              trans_pos = pos.clone()
              trans_pos[...,:3] += translate_tensor
              view_pos = torch.matmul(mvp, trans_pos.unsqueeze(-1)).squeeze(-1) 
              texture, mask, rast_uv = back_to_texture(glctx, look_at, view_pos, tri, target_image, tex_uv, uv_idx, idx, vn)
              target_texture.append(texture)
              target_mask.append(mask)
              rast_uvs.append(rast_uv)
          kd_map, opt_mask = aggregate_texture(kd_map, target_texture, target_mask, rast_uvs)
          opt_mask = opt_mask[0]
      Image.fromarray((np.clip(kd_map.detach().cpu().numpy() * 255, 0, 255)).astype(np.uint8)).save(f"{save_path}/refined_texture_kd.jpg")
  
     #ms.save_current_mesh(f"{save_path}/model-00.obj")
      with open(f"{save_path}/model-00.mtl", "w") as f:
          f.write(f"newmtl default\nKa 0.0 0.0 0.0\nmap_Kd refined_texture_kd.jpg\nKs 0.0 0.0 0.0")