Esempio n. 1
0
def train(train_run = (True,True), restore = False, epoch = 0):
  disc_train_run,segm_train_run = train_run
  train_run = disc_train_run or segm_train_run

  if not train_run:
    FLAGS.batch_size = 1
    FLAGS.num_epochs = 1

  train_run = disc_train_run or segm_train_run

  timestr        = time.strftime("TRAIN/%d_%b_%Y_%H_%M",time.localtime()) if train_run else time.strftime("TEST/%d_%b_%Y_%H_%M",time.localtime())
  timestr = timestr + "_EPOCH_%d"%epoch
  # Location to log to.
  split = 'TRAIN' if train_run else 'TEST'

  # We just run from the saved model directory for demoing.
  FLAGS.run_dir = FLAGS.mod_dir
  filestr        = FLAGS.run_dir + "tensorlogs/" + timestr + '/'
  ft.directoryFixer(filestr + 'patches/')

  print("Running from: " + filestr)

  tf.reset_default_graph()
  discriminator = DiscNet(disc_train_run,restore,timestr,split)

  tf.reset_default_graph()
  segmenter     = SegNet(segm_train_run,restore,timestr)

  # Starts the input generator
  print("\rSTARTING INPUT GENERATION THREADS...")
  coord          = tf.train.Coordinator()
  threads        = tf.train.start_queue_runners(sess = discriminator.sess, coord = coord)
  print("STARTING TRAINING...")

  step = 0
  full_pats = 0
  epoch = 0

  disc_losses = []
  segm_losses = []

  try:
    while not coord.should_stop():
      # Run the network and write summaries
      try:
        orig_img,orig_lab,img_pat,lab_pat,disc_log,count,disc_loss = discriminator.run()
        disc_losses.append(disc_loss)
      except IndexError:
        break
        # Epoch done.
      img_pat  = np.reshape(img_pat,(-1,FLAGS.patch_size,FLAGS.patch_size,3))
      lab_pat  = np.reshape(lab_pat,(-1,FLAGS.patch_size,FLAGS.patch_size,1))
      disc_log = np.squeeze(disc_log)


      # Do some processing, create new array with only patches we need.
      new_imgs = []
      new_labs = []

      for x in range(len(disc_log)):
        if(disc_log[x] == 1):
          new_imgs.append(img_pat[x])
          new_labs.append(lab_pat[x])

      new_imgs = np.array(new_imgs)
      new_labs = np.array(new_labs)

      imgs_labs_losses   = []

      if(np.sum(disc_log) > 0):
        seg_log,seg_loss,per_pat_loss = segmenter.run(new_imgs,new_labs,count)
        segm_losses.append(seg_loss)

        # Train on the worst 1/2 images twice.
        if train_run:
          im_loss = zip(new_imgs,new_labs,per_pat_loss)
          [imgs_labs_losses.append(im_lab_loss) for im_lab_loss in im_loss]


        # Do some more processing, weave the resultant patches into an array
        # of the resultant logit map
        y = 0



      if not train_run:
        full_pats = np.zeros(shape = lab_pat.shape)
        for x in range(len(disc_log)):
          if(disc_log[x] == 1):
            full_pats[x] = seg_log[y]
            y+=1

        orig_img = np.squeeze(orig_img).astype('uint8')
        orig_lab = util.cpu_colorize(np.squeeze(orig_lab))
        # Go from patches to full image logit map.
        result = patches_to_image(full_pats)
        img = Image.fromarray(util.cpu_colorize(result))
        img.save(filestr + '%d_log.png'%step)
        img = Image.fromarray(orig_img)
        img.save(filestr + '%d_img.png'%step)
        img = Image.fromarray(orig_lab)
        img.save(filestr + '%d_lab.png'%step)
        for x in range(new_imgs.shape[0]):
          img = Image.fromarray(np.squeeze(new_imgs[x]).astype(np.uint8))
          img.save(filestr + 'patches/' + '%d_%d_img_pat.png'%(step,x))
          img = Image.fromarray(np.squeeze(util.cpu_colorize(new_labs[x])))
          img.save(filestr + 'patches/' + '%d_%d_lab_pat.png'%(step,x))
          img = Image.fromarray(np.squeeze(util.cpu_colorize(seg_log[x])))
          img.save(filestr + 'patches/' + '%d_%d_lab_pat.png'%(step,x))
      step +=1

    # Train on bad patches over again.
    if segm_train_run:
      imgs_labs_losses.sort(key = lambda tup: tup[2])
      iterval = 0
      while len(imgs_labs_losses) > 1:
        iterval += 1
        imgs_labs_losses = imgs_labs_losses[len(imgs_labs_losses)//2:]
        generator.perturb(imgs_labs_losses)
        new_imgs = []
        new_labs = []
        [(new_imgs.append(new_img),new_labs.append(new_lab)) for new_img,new_lab,_ in imgs_labs_losses]
        # Run through 10 patches at a time as a batch.
        for x in range(len(new_imgs)//10 + 1):
          print('\rTRAINING ON HARD EXAMPLES %d/%d ITER %d'%(x,len(new_imgs)//10+1,iterval),end='')
          _ = segmenter.run(new_imgs[x*10:(x+1)*10],new_labs[x*10:(x+1)*10],0,False)
      print('\rDONE TRAINING HARD EXAMPLES')

  except KeyboardInterrupt:
    pass
  finally:
    if train_run:

      discriminator.save()
      segmenter.save()
    coord.request_stop()
  coord.join(threads)

  discriminator.close()
  segmenter.close()
  return np.mean(disc_losses),np.mean(segm_losses)
Esempio n. 2
0
def train(train_run=(True, True), restore=False, epoch=0):
    disc_train_run, segm_train_run = train_run
    train_run = disc_train_run or segm_train_run

    if not train_run:
        FLAGS.batch_size = 1
        FLAGS.num_epochs = 1

    timestr = time.strftime("TRAIN/%d_%b_%Y_%H_%M",
                            time.localtime()) if train_run else time.strftime(
                                "TEST/%d_%b_%Y_%H_%M", time.localtime())
    timestr = timestr + "_EPOCH_%d" % epoch
    # Location to log to.
    split = 'TRAIN' if train_run else 'TEST'

    filestr = FLAGS.run_dir + "tensorlogs/" + timestr + '/'
    ft.directoryFixer(filestr + 'patches/')

    print("Running from: " + filestr)

    tf.reset_default_graph()
    discriminator = DiscNet(disc_train_run, restore, timestr, split)

    tf.reset_default_graph()
    segmenter = SegNet(segm_train_run, restore, timestr)

    # Starts the input generator
    print("\rSTARTING INPUT GENERATION THREADS...")
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=discriminator.sess,
                                           coord=coord)
    print("STARTING TRAINING...")

    step = 0
    full_pats = 0
    epoch = 0

    disc_losses = []
    segm_losses = []
    sum_cmat = np.zeros((FLAGS.num_classes, FLAGS.num_classes))

    game_1 = []
    game_2 = []
    game_3 = []
    game_4 = []
    ssim = []
    psnr = []
    rers = []
    maes = []
    mses = []

    try:
        while not coord.should_stop():
            # Run the network and write summaries
            start = timeit.default_timer()
            try:
                orig_img, orig_lab, ids, img_pat, lab_pat, disc_log, count, disc_loss = discriminator.run(
                )
                disc_losses.append(disc_loss)
            except IndexError:
                break
                # Epoch done.
            if FLAGS.cell_selection or (FLAGS.positive_training and train_run):
                img_pat = np.reshape(
                    img_pat, (-1, FLAGS.patch_size, FLAGS.patch_size, 3))
                lab_pat = np.reshape(
                    lab_pat, (-1, FLAGS.patch_size, FLAGS.patch_size, 1))
                disc_log = np.squeeze(disc_log)
            else:
                img_pat = np.reshape(orig_img, (-1, FLAGS.imgH, FLAGS.imgW, 3))
                lab_pat = np.reshape(orig_lab, (-1, FLAGS.imgH, FLAGS.imgW, 1))
                disc_log = np.ones((FLAGS.batch_size))
            # Do some processing, create new array with only patches we need.
            new_imgs = []
            new_labs = []

            for x in range(len(disc_log)):
                if (disc_log[x] == 1):
                    new_imgs.append(img_pat[x])
                    new_labs.append(lab_pat[x])

            new_imgs = np.array(new_imgs)
            new_labs = np.array(new_labs)

            imgs_labs_losses = []

            if (np.sum(disc_log) > 0):
                seg_log, seg_loss, per_pat_loss, cmat = segmenter.run(
                    new_imgs, new_labs, count)
                segm_losses.append(seg_loss)
                sum_cmat += cmat

                # Train on the worst 1/2 images twice.
                if train_run:
                    im_loss = zip(new_imgs, new_labs, per_pat_loss)
                    [
                        imgs_labs_losses.append(im_lab_loss)
                        for im_lab_loss in im_loss
                    ]

                # Do some more processing, weave the resultant patches into an array
                # of the resultant logit map
                y = 0
            stop = timeit.default_timer()
            print('\rTime: ', stop - start, end='')

            if not train_run:
                full_pats = np.zeros(shape=lab_pat.shape, dtype=np.float32)
                for x in range(len(disc_log)):
                    if (disc_log[x] == 1):
                        full_pats[x] = seg_log[y]
                        y += 1

                orig_img = np.squeeze(orig_img).astype('uint8')
                orig_lab = np.squeeze(orig_lab)
                result = np.squeeze(patches_to_image(full_pats))

                with h5py.File(filestr + '%s_log.png' % ids, 'w') as hf:
                    hf['density'] = np.squeeze(orig_lab)
                with h5py.File(filestr + '%s_ann.png' % ids, 'w') as hf:
                    hf['density'] = np.squeeze(result)

                # Before we colorize the result, we want to run the GAME metrics
                game_1.append(util.cpu_GAME(orig_lab, result, 1))
                game_2.append(util.cpu_GAME(orig_lab, result, 2))
                game_3.append(util.cpu_GAME(orig_lab, result, 3))
                game_4.append(util.cpu_GAME(orig_lab, result, 4))
                ssim.append(util.cpu_psnr(orig_lab, result))
                psnr.append(util.cpu_ssim(orig_lab, result))
                rers.append(
                    np.abs(np.sum(orig_lab) - np.sum(result)) /
                    np.sum(orig_lab))
                maes.append(np.abs(np.sum(orig_lab) - np.sum(result)))
                mses.append((np.sum(orig_lab) - np.sum(result))**2)

                result = util.cpu_colorize(result)
                orig_lab = util.cpu_colorize(np.squeeze(orig_lab))

                img = Image.fromarray(result)
                img.save(filestr + '%d_log.png' % step)
                img = Image.fromarray(orig_img)
                img.save(filestr + '%d_img.png' % step)
                img = Image.fromarray(orig_lab)
                img.save(filestr + '%d_lab.png' % step)
                for x in range(new_imgs.shape[0]):
                    img = Image.fromarray(
                        np.squeeze(new_imgs[x]).astype(np.uint8))
                    img.save(filestr + 'patches/' + '%d_%d_img_pat.png' %
                             (step, x))
                    img = Image.fromarray(
                        np.squeeze(util.cpu_colorize(new_labs[x])))
                    img.save(filestr + 'patches/' + '%d_%d_lab_pat.png' %
                             (step, x))
                    img = Image.fromarray(
                        np.squeeze(util.cpu_colorize(seg_log[x])))
                    img.save(filestr + 'patches/' + '%d_%d_log_pat.png' %
                             (step, x))
            step += 1

        # Train on bad patches over again.
        if segm_train_run:
            imgs_labs_losses.sort(key=lambda tup: tup[2])
            iterval = 0
            while len(imgs_labs_losses) > 1:
                iterval += 1
                imgs_labs_losses = imgs_labs_losses[len(imgs_labs_losses) //
                                                    2:]
                generator.perturb(imgs_labs_losses)
                new_imgs = []
                new_labs = []
                [(new_imgs.append(new_img), new_labs.append(new_lab))
                 for new_img, new_lab, _ in imgs_labs_losses]
                # Run through 10 patches at a time as a batch.
                for x in range(len(new_imgs) // 10 + 1):
                    print('\rTRAINING ON HARD EXAMPLES %d/%d ITER %d' %
                          (x, len(new_imgs) // 10 + 1, iterval),
                          end='')
                    _ = segmenter.run(new_imgs[x * 10:(x + 1) * 10],
                                      new_labs[x * 10:(x + 1) * 10], 0, False)
            print('\rDONE TRAINING HARD EXAMPLES')

    except KeyboardInterrupt:
        pass
    finally:
        if train_run:
            discriminator.save()
            segmenter.save()
        else:
            print("GAME 1,2,3,4")
            print("Game 1, ", np.mean(game_1))
            print("Game 2, ", np.mean(game_2))
            print("Game 3, ", np.mean(game_3))
            print("Game 4, ", np.mean(game_4))
            print("SSIM, ", np.mean(ssim))
            print("PSNR, ", np.mean(psnr))
            print("Rel_Err, ", np.mean(rers))
            print("MAE, ", np.mean(maes))
            print("MSE, ", np.mean(mses))
        coord.request_stop()
    coord.join(threads)

    np.save(filestr + "data.dat", sum_cmat)

    discriminator.close()
    segmenter.close()
    return np.mean(disc_losses), np.mean(segm_losses)