text: import os
  from tqdm import trange
  import torch
  from torch.nn import functional as F
  from torch import distributions as dist
  from im2mesh.common import (
      compute_iou, make_3d_grid, transform_points
  )
  from im2mesh.utils import visualize as vis
  from im2mesh.training import BaseTrainer
  import numpy as np
  
  
  class Trainer(BaseTrainer):
      ''' Trainer object for the Occupancy Network.
  
      Args:
          model (nn.Module): Occupancy Network model
          optimizer (optimizer): pytorch optimizer object
          device (device): pytorch device
          input_type (str): input type
          vis_dir (str): visualization directory
          threshold (float): threshold value
          eval_sample (bool): whether to evaluate samples
  
      '''
  
      def __init__(self, model, optimizer, device=None, input_type='img',
                   vis_dir=None, threshold=0.5, eval_sample=False, camera_space=False):
          self.model = model
          self.optimizer = optimizer
          self.device = device
          self.input_type = input_type
          self.vis_dir = vis_dir
          self.threshold = threshold
          self.eval_sample = eval_sample
          self.camera_space = camera_space
  
          if vis_dir is not None and not os.path.exists(vis_dir):
              os.makedirs(vis_dir)
  
      def train_step(self, data):
          ''' Performs a training step.
  
          Args:
              data (dict): data dictionary
          '''
          self.model.train()
          self.optimizer.zero_grad()
          loss = self.compute_loss(data)
          loss.backward()
          self.optimizer.step()
          return loss.item()
  
      def eval_step(self, data):
          ''' Performs an evaluation step.
  
          Args:
              data (dict): data dictionary
          '''
          self.model.eval()
  
          device = self.device
          threshold = self.threshold
          eval_dict = {}
  
          # Compute elbo
          points = data.get('points').to(device)
          occ = data.get('points.occ').to(device)
  
          inputs = data.get('inputs', torch.empty(points.size(0), 0)).to(device)
          voxels_occ = data.get('voxels')
  
          points_iou = data.get('points_iou').to(device)
          occ_iou = data.get('points_iou.occ').to(device)
  
  
          '''TRANSFORMATION TO CAMERA SPACE'''
          if self.camera_space:
              transform = data.get('inputs.world_mat').to(device)
              R = transform[:, :, :3]
              points = transform_points(points, transform)
              points_iou = transform_points(points_iou, R)
          '''END'''
  
          kwargs = {}
  
          with torch.no_grad():
              elbo, rec_error, kl = self.model.compute_elbo(
                  points, occ, inputs, **kwargs)
  
          eval_dict['loss'] = -elbo.mean().item()
          eval_dict['rec_error'] = rec_error.mean().item()
          eval_dict['kl'] = kl.mean().item()
  
          # Compute iou
          batch_size = points.size(0)
  
          with torch.no_grad():
              p_out = self.model(points_iou, inputs,
                                 sample=self.eval_sample, **kwargs)
  
          occ_iou_np = (occ_iou >= 0.5).cpu().numpy()
          occ_iou_hat_np = (p_out.probs >= threshold).cpu().numpy()
          iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean()
          eval_dict['iou'] = iou
  
          '''
          with torch.no_grad():
              p_out_r = self.model(points_iou_r, inputs,
                                 sample=self.eval_sample, **kwargs)
  
          occ_iou_r_hat_np = (p_out_r.probs >= threshold).cpu().numpy()
          iou_r = compute_iou(occ_iou_np, occ_iou_r_hat_np).mean()
          eval_dict['iou_r'] = iou_r
  
          data['iou_r'] = iou_r
          data['iou'] = iou
          data['occ_iou_r_hat_np'] = occ_iou_r_hat_np
          data['occ_iou_hat_np'] = occ_iou_hat_np
          
          SAVE IOU DATA
          # Create npz and save
          im_path = str(data.get('inputs.image_path')[0]).split('/')
          model=im_path[3]
          im_nr=im_path[5].split('.')[0]
          out_dir = 'IOU'
          np.savez(os.path.join(out_dir, model+'_'+im_nr), iou_r=iou_r,
                   occ_iou_hat_np=occ_iou_hat_np,
                   occ_iou_r_hat_np=occ_iou_r_hat_np,
                   iou=iou,
                   occ_iou=occ_iou.cpu().numpy(),
                   transform=transform.cpu().numpy(),
                   points_iou=points_iou.cpu().numpy(),
                   points_iou_r=points_iou_r.cpu().numpy(),
                   model=model,
                   im_nr=im_nr,
                   image=inputs.cpu().numpy()
                   )
          END'''
  
          # Estimate voxel iou
          if voxels_occ is not None:
              voxels_occ = voxels_occ.to(device)
              points_voxels = make_3d_grid(
                  (-0.5 + 1/64,) * 3, (0.5 - 1/64,) * 3, (32,) * 3)
              points_voxels = points_voxels.expand(
                  batch_size, *points_voxels.size())
              points_voxels = points_voxels.to(device)
              with torch.no_grad():
                  p_out = self.model(points_voxels, inputs,
                                     sample=self.eval_sample, **kwargs)
  
              voxels_occ_np = (voxels_occ >= 0.5).cpu().numpy()
              occ_hat_np = (p_out.probs >= threshold).cpu().numpy()
              iou_voxels = compute_iou(voxels_occ_np, occ_hat_np).mean()
  
              eval_dict['iou_voxels'] = iou_voxels
  
          return eval_dict
  
      def visualize(self, data):
          ''' Performs a visualization step for the data.
  
          Args:
              data (dict): data dictionary
          '''
          device = self.device
  
          batch_size = data['points'].size(0)
          inputs = data.get('inputs', torch.empty(batch_size, 0)).to(device)
  
          shape = (32, 32, 32)
          p = make_3d_grid([-0.5] * 3, [0.5] * 3, shape).to(device)
          p = p.expand(batch_size, *p.size())
  
          kwargs = {}
          with torch.no_grad():
              p_r = self.model(p, inputs, sample=self.eval_sample, **kwargs)
  
          occ_hat = p_r.probs.view(batch_size, *shape)
          voxels_out = (occ_hat >= self.threshold).cpu().numpy()
  
          for i in trange(batch_size):
              input_img_path = os.path.join(self.vis_dir, '%03d_in.png' % i)
              vis.visualize_data(
                  inputs[i].cpu(), self.input_type, input_img_path)
              vis.visualize_voxels(
                  voxels_out[i], os.path.join(self.vis_dir, '%03d.png' % i))
  
      def compute_loss(self, data):
          ''' Computes the loss.
  
          Args:
              data (dict): data dictionary
          '''
          device = self.device
          p = data.get('points').to(device)
          occ = data.get('points.occ').to(device)
          inputs = data.get('inputs', torch.empty(p.size(0), 0)).to(device)
  
          '''TRANSFORMATION TO CAMERA SPACE'''
          if self.camera_space:
              transform = data.get('inputs.world_mat').to(device)
              R = transform[:, :, :3]
              p = transform_points(p, transform)
          '''END'''
  
          kwargs = {}
  
          c = self.model.encode_inputs(inputs)
          q_z = self.model.infer_z(p, occ, c, **kwargs)
          z = q_z.rsample()
          # KL-divergence
          kl = dist.kl_divergence(q_z, self.model.p0_z).sum(dim=-1)
          loss = kl.mean()
  
          # General points
          logits = self.model.decode(p, z, c, **kwargs).logits
          loss_i = F.binary_cross_entropy_with_logits(
              logits, occ, reduction='none')
          loss = loss + loss_i.sum(-1).mean()
  
          return loss
  