text: import numpy as np
  import torch
  from tqdm import tqdm
  import matplotlib.pyplot as plt
  from matplotlib import cm
  from PIL import Image, ImageFilter
  
  import regularization as reg
  import utils
  from metrics import get_pxacc_miou
  from utils import color_map
  import paths
  import models
  
  class Optimizer():
      ''' Optimization of Scribble based Segmentation.'''
  
  
      @staticmethod
      def train(model, data, criterion, optimizer, regularization, setting, save_name):
          ''' Training of scribble based Segmentation.
  
          Args:
              model:
                  Neural Network model used for the training.
              data:
                  Single image Dataset of the class "Single_Image".
              criterion:
                  Loss function.
              optimizer:
                  Optimizer.
              regularization:
                  Regularization of the energy function, e.g. total variation.
              setting:
                  Dictionary defining the setting for the training procedure.
              save_name:
                  Name under which the trainings results are saved.
          Returns:
              The final intersection over union value,
              the final pixel accuracy and the final loss value.
          '''
  
          model.train()
          loss_plot = []
  
          # Case1: Training on Pixelwise Networks
          if setting['name'] == "FC_Net":
  
              # Extract Data
              sample = data['2d']
              img2d = sample['rgb'].to(setting['dev'])
              xy2d = sample['xy'].to(setting['dev'])
              scribble2d = sample['scribble'].squeeze(1).to(setting['dev'])
              gt2d = sample['gt'].flatten()            
  
              # Train Loop
              for epoch in tqdm(range(setting['epochs'])):
                  reg_loss = 0
                  model.train()
  
                  # Extract only Scribbled Pixels
                  scribbled_pixels = scribble2d!=data.get_number_classes()
                  img = img2d[scribbled_pixels]
                  scribble = scribble2d[scribbled_pixels]
                  gt = gt2d[scribbled_pixels]
                  xy = xy2d[scribbled_pixels]
  
                  # Select randomly setting["bs"] many pixels
                  if setting["bs"] is not None:
                      no = min(setting["bs"], len(gt)-1)
                      indices = torch.randint(0, len(gt)-1, (no,))
                      img = img[indices]
                      scribble = scribble[indices]
                      gt = gt[indices]
                      xy = xy[indices]
              
                  # Zero Grad
                  optimizer.zero_grad()
  
                  # Predict
                  pred = model(img, xy)
                  reg_loss = setting['tau'] * regularization(pred)
                  loss = criterion(pred, scribble.long()) + reg_loss
  
                  # Optimization
                  loss.backward()
                  optimizer.step()
  
                  # Loss and Metrics
                  with torch.no_grad():
                      loss_plot.append(loss.item())
  
                      if epoch == (setting['epochs'])-1:
  
                          # Get Accuracies
                          model.eval()
                          argmax_pred = torch.argmax(model(img2d, xy2d), dim=1)
                          acc, miou = get_pxacc_miou(gt2d.cpu(), argmax_pred.cpu(), noneclass = data.get_number_classes())
  
                          # Plot
                          pred = pred.reshape(gt.shape[0],gt.shape[1],pred.shape[1]).permute(2,0,1).unsqueeze(0)
                          Optimizer.plot_state(pred.cpu(), img.cpu(), scribble.cpu(), gt.cpu(), loss_plot, save_name)
                          Optimizer.save_contoured_prediction(torch.argmax(pred[0],dim=0).cpu().numpy(), img[0].detach().cpu(), save_name)
  
          # Case2: Training on Convolutional Networks
          else:
  
              # Extract Data
              sample = data['3d']
              img = sample['rgb'].to(setting['dev'])
              xy = sample['xy'].to(setting['dev'])
              scribble = sample['scribble'].to(setting['dev'])
              gt = sample['gt']
  
              # Train Loop
              for epoch in tqdm(range(setting['epochs'])):
  
                  if setting['input'] != 'rgb':
                      xy.requires_grad = True
                      img.requires_grad = True
  
                  # Zero Grad
                  optimizer.zero_grad()
  
                  # Predict
                  pred = model(img, xy.float())
                  loss = criterion(pred, scribble.long())
                  
                  # Regularization
                  if setting['input'] != 'rgb':
                      
                      input_gradient = torch.autograd.grad(loss, xy, retain_graph=True, create_graph=True)
                      input_gradient_rgb = torch.autograd.grad(loss, img, retain_graph=True, create_graph=True)
  
                      loss += setting['xygrad'] * torch.mean(torch.abs(input_gradient[0])) * 1e6
                      loss += setting['rgbgrad'] * torch.sum(torch.abs(input_gradient_rgb[0]))
  
                  loss += setting['tau'] * regularization(pred)
  
                  # Optimization
                  loss.backward()
                  optimizer.step()
  
                  # Metrics
                  with torch.no_grad():
                      loss_plot.append(loss.item())
  
                      if epoch == (setting['epochs'])-1:
                          
                          argmax_pred = torch.argmax(pred.cpu(), dim=1)
                          acc, miou = get_pxacc_miou(gt.cpu(), argmax_pred[0].cpu(), noneclass = data.get_number_classes())
  
                          Optimizer.save_contoured_prediction(argmax_pred[0].numpy(), img[0].detach().cpu(), save_name)
                          Optimizer.plot_state(pred.cpu(), img.cpu(), scribble.cpu(), gt.cpu(), loss_plot, save_name)
                          
          return miou, acc, loss.item()
  
      @staticmethod
      def plot_state(prediction, image, scribble , ground_truth, loss_full, savename):
          vmin, vmax = 0, torch.max(scribble)
          blend = utils.blend_image_segmentation(image[0], torch.argmax(prediction, dim=1)[0])
  
          # Plot Loss
          plt.figure(figsize=(10,4))
          plt.plot(loss_full)
          plt.savefig(savename + '_loss.png')
          plt.close()
  
          # Plot Original Image
          _, axs = plt.subplots(1,4,figsize=(20,4))
          axs[0].imshow(image[0].permute(1,2,0))
          axs[0].set_title('Original Image')
          axs[1].imshow(blend, vmin=vmin, vmax=vmax)
          axs[1].set_title('Prediction')
          axs[2].imshow(scribble[0], vmin=vmin, vmax=vmax)
          axs[2].set_title('Scribble Data')
          axs[3].imshow(ground_truth, vmin=vmin, vmax=vmax)
          axs[3].set_title('Ground Truth')
          plt.savefig(savename + '_all.png')
          plt.close()
  
      @staticmethod
      def save_contoured_prediction(prediction, image, save_name):
          image = utils.blend_image_segmentation(image.detach().cpu(), prediction)
  
          # load segmented image as greyscale
          seg = Image.fromarray(np.uint8((prediction))).convert("L")
          image = np.uint8((np.array(image*255)))
  
          # draw contours
          colors = [(255,255,0),(0,255,0),(0,0,255),(0,255,255),(255,0,0),(255,0,255)]
          for idx, color in zip(range(len(colors)), colors):
              c_contour = seg.point(lambda p:p==(idx+1) and 255)
              edges = np.array(c_contour.filter(ImageFilter.FIND_EDGES))
              image[np.nonzero(edges)] = color
  
          # save result
          Image.fromarray(image).save(save_name+'_result.png')
  