text: from __future__ import print_function
  
  import sys
  from datetime import datetime
  
  import torch
  
  from src.utils.cluster.cluster_eval import cluster_subheads_eval
  from src.utils.cluster.transforms import sobel_process
  
  
  def segmentation_eval(config, net,
                        mapping_assignment_dataloader,
                        mapping_test_dataloader,
                        sobel, using_IR=False, verbose=0, return_only=False):
    torch.cuda.empty_cache()
    net.eval()
  
    stats_dict = cluster_subheads_eval(config, net,
                                       mapping_assignment_dataloader=mapping_assignment_dataloader,
                                       mapping_test_dataloader=mapping_test_dataloader,
                                       sobel=sobel,
                                       using_IR=using_IR,
                                       get_data_fn=_segmentation_get_data,
                                       verbose=verbose)
  
    net.train()
  
    acc = stats_dict["best"]
    is_best = (len(config.epoch_acc) > 0) and (acc > max(config.epoch_acc))
  
    torch.cuda.empty_cache()
  
    if not return_only:
      config.epoch_stats.append(stats_dict)
      config.epoch_acc.append(acc)
      config.epoch_avg_subhead_acc.append(stats_dict["avg"])
  
      return is_best
    else:
      return stats_dict
  
  
  def _segmentation_get_data(config, net, dataloader, sobel=False,
                             using_IR=False, verbose=0):
    # returns (vectorised) cuda tensors for flat preds and targets
    # sister of _clustering_get_data
  
    assert (config.output_k <= 255)
  
    num_batches = len(dataloader)
    num_samples = 0
  
    # upper bound, will be less for last batch
    samples_per_batch = config.batch_sz * config.input_sz * config.input_sz
  
    if verbose > 0:
      print("started _segmentation_get_data %s" % datetime.now())
      sys.stdout.flush()
  
    # vectorised
    flat_predss_all = [torch.zeros((num_batches * samples_per_batch),
                                   dtype=torch.uint8).cuda() for _ in xrange(
      config.num_subheads)]
    flat_targets_all = torch.zeros((num_batches * samples_per_batch),
                                   dtype=torch.uint8).cuda()
    mask_all = torch.zeros((num_batches * samples_per_batch),
                           dtype=torch.uint8).cuda()
  
    if verbose > 0:
      batch_start = datetime.now()
      all_start = batch_start
      print("starting batches %s" % batch_start)
  
    for b_i, batch in enumerate(dataloader):
  
      imgs, flat_targets, mask = batch
      imgs = imgs.cuda()
  
      if sobel:
        imgs = sobel_process(imgs, config.include_rgb, using_IR=using_IR)
  
      with torch.no_grad():
        x_outs = net(imgs)
  
      assert (x_outs[0].shape[1] == config.output_k)
      assert (x_outs[0].shape[2] == config.input_sz and x_outs[0].shape[
        3] == config.input_sz)
  
      # actual batch size
      actual_samples_curr = (
        flat_targets.shape[0] * config.input_sz * config.input_sz)
      num_samples += actual_samples_curr
  
      # vectorise: collapse from 2D to 1D
      start_i = b_i * samples_per_batch
      for i in xrange(config.num_subheads):
        x_outs_curr = x_outs[i]
        assert (not x_outs_curr.requires_grad)
        flat_preds_curr = torch.argmax(x_outs_curr, dim=1)
        flat_predss_all[i][
        start_i:(start_i + actual_samples_curr)] = flat_preds_curr.view(-1)
  
      flat_targets_all[
      start_i:(start_i + actual_samples_curr)] = flat_targets.view(-1)
      mask_all[start_i:(start_i + actual_samples_curr)] = mask.view(-1)
  
      if verbose > 0 and b_i < 3:
        batch_finish = datetime.now()
        print("finished batch %d, %s, took %s, of %d" %
              (b_i, batch_finish, batch_finish - batch_start, num_batches))
        batch_start = batch_finish
        sys.stdout.flush()
  
    if verbose > 0:
      all_finish = datetime.now()
      print(
        "finished all batches %s, took %s" % (all_finish, all_finish - all_start))
      sys.stdout.flush()
  
    flat_predss_all = [flat_predss_all[i][:num_samples] for i in
                       xrange(config.num_subheads)]
    flat_targets_all = flat_targets_all[:num_samples]
    mask_all = mask_all[:num_samples]
  
    flat_predss_all = [flat_predss_all[i].masked_select(mask=mask_all) for i in
                       xrange(config.num_subheads)]
    flat_targets_all = flat_targets_all.masked_select(mask=mask_all)
  
    if verbose > 0:
      print("ended _segmentation_get_data %s" % datetime.now())
      sys.stdout.flush()
  
    selected_samples = mask_all.sum()
    assert (len(flat_predss_all[0].shape) == 1 and
            len(flat_targets_all.shape) == 1)
    assert (flat_predss_all[0].shape[0] == selected_samples)
    assert (flat_targets_all.shape[0] == selected_samples)
  
    return flat_predss_all, flat_targets_all
  