Пример #1
0
def make_net(ims, samples, pr, reuse = True, train = True):
  if pr.net_type == 'i3d':
    import i3d_kinetics
    keep_prob = 0.5 if train else 1.
    if pr.use_i3d_logits:
      with tf.variable_scope('RGB', reuse = reuse):
        net = tfu.normalize_ims(ims)
        i3d_net = i3d_kinetics.InceptionI3d(pr.num_classes, spatial_squeeze = True, final_endpoint = 'Logits')
        logits, _ = i3d_net(net, is_training = train, dropout_keep_prob = keep_prob)
        return ut.Struct(logits = logits, prob = tf.nn.softmax(logits), last_conv = logits)
    else:
      with tf.variable_scope('RGB', reuse = reuse):
        i3d_net = i3d_kinetics.InceptionI3d(pr.num_classes, final_endpoint = 'Mixed_5c')
        net = tfu.normalize_ims(ims)
        net, _ = i3d_net(net, is_training = train, dropout_keep_prob = keep_prob)
      last_conv = net
      net = tf.reduce_mean(last_conv, [1, 2, 3], keep_dims = True)
      with slim.arg_scope(shift.arg_scope(pr, reuse = reuse, train = train)):
        logits = shift.conv3d(
          net, pr.num_classes, [1, 1, 1], scope = 'lb/logits', 
          activation_fn = None, normalizer_fn = None)[:, 0, 0, 0, :]
        return ut.Struct(logits = logits, 
                         prob = tf.nn.softmax(logits), 
                         last_conv = net)

  elif pr.net_type == 'shift':
    with slim.arg_scope(shift.arg_scope(pr, reuse = reuse, train = train)):
      # To train the network without audio, you can set samples to be an all-zero array, and
      # set pr.use_sound = False.
      shift_net = shift.make_net(ims, samples, pr, reuse = reuse, train = train)
      if pr.use_dropout:
        shift_net.last_conv = slim.dropout(shift_net.last_conv, is_training = train)

      net = shift_net.last_conv
      net = tf.reduce_mean(net, [1, 2, 3], keep_dims = True)
      logits = shift.conv3d(
        net, pr.num_classes, [1, 1, 1], scope = 'lb/logits', 
        activation_fn = None, normalizer_fn = None)[:, 0, 0, 0, :]
      return ut.Struct(logits = logits, prob = tf.nn.softmax(logits), last_conv = net)
  elif pr.net_type == 'c3d':
    import c3d
    with slim.arg_scope(shift.arg_scope(reuse = reuse, train = train)):
      net = c3d.make_net(ims, samples, pr, reuse = reuse, train = train)
      net = net.last_conv
      net = tf.reduce_mean(net, [1, 2, 3], keep_dims = True)
      logits = c3d.conv3d(
        net, pr.num_classes, [1, 1, 1], scope = 'lb/logits', 
        activation_fn = None, normalizer_fn = None)[:, 0, 0, 0, :]
      return ut.Struct(logits = logits, prob = tf.nn.softmax(logits), last_conv = net)
  else: 
    raise RuntimeError()
Пример #2
0
def make_discrim_spec(spec_in,
                      spec_out,
                      phase_in,
                      phase_out,
                      pr,
                      reuse=True,
                      train=True):
    with slim.arg_scope(unet_arg_scope(pr, reuse=reuse, train=train)):
        spec_in = normalize_spec(spec_in, pr)
        spec_out = normalize_spec(spec_out, pr)
        spec_in = ed(spec_in, 3)
        spec_out = ed(spec_out, 3)
        phase_in = ed(phase_in, 3)
        phase_out = ed(phase_out, 3)
        net = tf.concat([spec_in, phase_in, spec_out, phase_out], 3)

        net = conv2d_same(net, 32, 4, scope='discrim/spec/conv1', stride=2)
        net = conv2d_same(net, 64, 4, scope='discrim/spec/conv2', stride=2)
        net = conv2d(net, 128, 4, scope='discrim/spec/conv3', stride=2)
        net = conv2d(net, 128, 4, scope='discrim/spec/conv4', stride=2)
        # net = conv2d(net, 256, 4, scope = 'discrim/spec/conv5', stride = 2)
        # net = conv2d(net, 256, 4, scope = 'discrim/spec/conv6', stride = [1, 2])
        logits = conv2d(net,
                        1,
                        1,
                        scope='discrim/spec/logits',
                        stride=1,
                        normalizer_fn=None,
                        activation_fn=None)
        return ut.Struct(logits=logits)
Пример #3
0
def triangulate_search(Ps, projs, im_shapes):
    projs = [p[0] for p in projs]
    ts = np.array([center_from_P(P) for P in Ps])
    (max_angle, X_best, tri) = (None, None, ut.Struct(inliers=[]))
    for inds in ransac_sample(len(Ps), 2, 300):
        Ps_s, projs_s, shapes_s = ut.take_inds_each([Ps, projs, im_shapes],
                                                    inds)

        # # does not seem to help
        # F = F_from_Ps(Ps_s[0], Ps_s[1], center_from_P(Ps_s[1]))
        # if not F_test_matches(F, projs_s[0][na, :], projs_s[1][na, :])[0]:
        #   continue

        X = triangulate_nonlinear_pt(Ps_s, projs_s, shapes_s)

        angle = ut.angle_between(-X + ts[inds[0]], -X + ts[inds[1]])
        pt_ok = (in_front_of(Ps_s[0], X) and in_front_of(Ps_s[1], X)) \
                and (max(reproj_error(P, X, x) for P, x in zip(Ps_s, projs_s)) <= 100)

        #pt_ok = (in_front_of(Ps_s[0], X) and in_front_of(Ps_s[1], X))

        # F = F_from_Ps(Ps_s[0], Ps_s[1], center_from_P(Ps_s[1]))
        # if max(reproj_error(P, X, x) for P, x in zip(Ps_s, projs_s)) > 100 and F_test_matches(F, projs_s[0][na, :], projs_s[1][na, :])[0]:
        #   asdf

        if pt_ok:
            if max_angle <= angle:
                max_angle = angle
                X_best = X
                tri = ut.Struct(inliers=inds,
                                angle_deg=np.degrees(angle),
                                in_front=(in_front_of(Ps_s[0], X),
                                          in_front_of(Ps_s[1], X)),
                                cam_dist=pylab.dist(ts[inds[0]], ts[inds[1]]),
                                reproj_errors=[
                                    reproj_error(P, X, x)
                                    for P, x in zip(Ps_s, projs_s)
                                ])

    return X_best, tri
Пример #4
0
def mix_sounds(samples0, pr, quiet_thresh_db=40., samples1=None):
    if pr.normalize_rms:
        samples0 = mu.normalize_rms(samples0)
        if samples1 is not None:
            samples1 = mu.normalize_rms(samples1)

    if samples1 is None:
        n = shape(samples0, 0) / 2
        samples0 = samples0[:, :pr.sample_len]
        if pr.both_videos_in_batch:
            print 'Using both videos'
            samples1 = tf.concat(
                [samples0[n:, :pr.sample_len], samples0[:n, :pr.sample_len]],
                axis=0)
        else:
            print 'Only using first videos'
            samples1 = samples0[n:]
            samples0 = samples0[:n]
    else:
        samples0 = samples0[:, :pr.sample_len]
        samples1 = samples1[:, :pr.sample_len]

    if pr.augment_rms:
        print 'Augmenting rms'
        # scale0 = tf.random_uniform((shape(samples0, 0), 1, 1), 0.9, 1.1)
        # scale1 = tf.random_uniform((shape(samples1, 0), 1, 1), 0.9, 1.1)
        db = 0.25
        scale0 = 2.**tf.random_uniform((shape(samples0, 0), 1, 1), -db, db)
        scale1 = 2.**tf.random_uniform((shape(samples1, 0), 1, 1), -db, db)
        samples0 = scale0 * samples0
        samples1 = scale1 * samples1

    samples_mix = samples0 + samples1
    spec_mix, phase_mix = stft(make_mono(samples_mix), pr)
    spec0, phase0 = stft(make_mono(samples0), pr)
    spec1, phase1 = stft(make_mono(samples1), pr)

    print 'Before truncating specgram:', shape(spec_mix)
    spec_mix = spec_mix[:, :pr.spec_len]
    print 'After truncating specgram:', shape(spec_mix)
    phase_mix = phase_mix[:, :pr.spec_len]
    spec0 = spec0[:, :pr.spec_len]
    spec1 = spec1[:, :pr.spec_len]
    phase0 = phase0[:, :pr.spec_len]
    phase1 = phase1[:, :pr.spec_len]

    return ut.Struct(samples=samples_mix,
                     phase=phase_mix,
                     spec=spec_mix,
                     sample_parts=[samples0, samples1],
                     spec_parts=[spec0, spec1],
                     phase_parts=[phase0, phase1])
def test(pr, gpu, test_on_train=False, center_crop=True):
    [gpu] = set_gpus([gpu])

    if pr.inputs == ['press']:
        net = PressClf(pr)
    else:
        #check_path = tf.train.latest_checkpoint(pr.train_dir)
        check_path = pj(pr.train_dir, 'net.tf-%d' % pr.model_iter)
        print 'Restoring from:', check_path
        net = NetClf(pr, check_path, gpu)

    if test_on_train:
        print 'Testing on train!'
        data = ut.load(pj(pr.dsdir, 'train.pk'))
    else:
        data = ut.load(pj(pr.dsdir, 'test.pk'))

    labels, probs, accs, vals = [], [], [], []
    for i in xrange(len(data)):
        ex = data[i]
        label = ex['is_gripping']

        def load_im(k, v):
            if k.startswith('gel') or k.startswith('im'):
                im = ig.uncompress(v)
            elif k.startswith('depth'):
                #v = np.tile(v, (1, 1, 3))
                im = v.astype('float32')
            else:
                raise RuntimeError()
            if center_crop:
                im = ut.crop_center(im, 224)
            return im

        inputs = {k: load_im(k, ex[k]) for k in im_names}
        inputs['initial_press_prob'] = ex['initial_press_prob']
        inputs['ee'] = ex['end_effector']

        pred, prob = net.predict(**inputs)
        #print prob, pred, label
        labels.append(label)
        probs.append(prob)
        accs.append(pred == label)
        if i % 50 == 0:
            print 'running average acc:', ut.f3(np.mean(accs))
        vals.append(
            ut.Struct(label=label,
                      prob=prob,
                      acc=accs[-1],
                      idx=i,
                      db_file=ex['db_file'],
                      object_name=ex['object_name']))

    labels = np.array(labels, 'bool')
    probs = np.array(probs, 'float32')
    accs = np.array(accs)

    acc = np.mean(accs)
    ap = sklearn.metrics.average_precision_score(labels, probs)
    print 'Accuracy:', acc
    print 'mAP:', ap
    print 'Base rate:', ut.f3(
        np.array(ut.mapattr(vals).label).astype('float32').mean())

    ut.save(pj(pr.resdir, 'eval_results.pk'),
            dict(acc=acc, ap=ap, results=(labels, probs)))
    ut.save(pj(pr.resdir, 'eval.pk'), vals)
    return acc
Пример #6
0
def make_net(ims,
             sfs,
             spec,
             phase,
             pr,
             reuse=True,
             train=True,
             vid_net_full=None):
    if pr.mono:
        print 'Using mono!'
        sfs = make_mono(sfs, tile=True)

    if vid_net_full is None:
        if pr.net_style == 'static':
            n = shape(ims, 1)
            if 0:
                ims_tile = tf.tile(ims[:, n / 2:n / 2 + 1], (1, n, 1, 1, 1))
            else:
                ims = tf.cast(ims, tf.float32)
                ims_tile = tf.tile(ims[:, n / 2:n / 2 + 1], (1, n, 1, 1, 1))
            vid_net_full = shift_net.make_net(ims_tile, sfs, pr, None, reuse,
                                              train)
        elif pr.net_style == 'no-im':
            vid_net_full = None
        elif pr.net_style == 'full':
            vid_net_full = shift_net.make_net(ims, sfs, pr, None, reuse, train)
        elif pr.net_style == 'i3d':
            with tf.variable_scope('RGB', reuse=reuse):
                import sep_i3d
                i3d_net = sep_i3d.InceptionI3d(1)
                vid_net_full = ut.Struct(
                    scales=i3d_net(ims, is_training=train))

    with slim.arg_scope(unet_arg_scope(pr, reuse=reuse, train=train)):  # test
        acts = []

        def conv(*args, **kwargs):
            out = conv2d(*args, activation_fn=None, **kwargs)
            acts.append(out)
            out = mu.lrelu(out, 0.2)
            return out

        def deconv(*args, **kwargs):
            args = list(args)
            if kwargs.get('do_pop', True):
                skip_layer = acts.pop()
            else:
                skip_layer = acts[-1]
            if 'do_pop' in kwargs:
                del kwargs['do_pop']
            x = args[0]
            if kwargs.get('concat', True):
                x = tf.concat([x, skip_layer], 3)
            if 'concat' in kwargs:
                del kwargs['concat']
            args[0] = tf.nn.relu(x)
            return deconv2d(*args, activation_fn=None, **kwargs)

        def merge_level(net, n):
            if vid_net_full is None:
                return net
            vid_net = tf.reduce_mean(vid_net_full.scales[n], [2, 3],
                                     keep_dims=True)
            vid_net = vid_net[:, :, 0, :, :]
            s = shape(vid_net)
            if shape(net, 1) != s[1]:
                vid_net = tf.image.resize_images(vid_net, [shape(net, 1), 1])
                print 'Video net before merge:', s, 'After:', shape(vid_net)
            else:
                print 'No need to resize:', s, shape(net)
            vid_net = tf.tile(vid_net, (1, 1, shape(net, 2), 1))
            net = tf.concat([net, vid_net], 3)
            acts[-1] = net
            return net

        num_freq = shape(spec, 2)
        net = tf.concat([
            ed(normalize_spec(spec, pr), 3),
            ed(normalize_phase(phase, pr), 3)
        ], 3)

        net = net[:, :, :pr.freq_len, :]
        net = conv(net, 64, 4, scope='gen/conv1', stride=[1, 2])
        net = conv(net, 128, 4, scope='gen/conv2', stride=[1, 2])
        net = conv(net, 256, 4, scope='gen/conv3', stride=2)
        net = merge_level(net, 0)
        net = conv(net, 512, 4, scope='gen/conv4', stride=2)
        net = merge_level(net, 1)
        net = conv(net, 512, 4, scope='gen/conv5', stride=2)
        net = merge_level(net, 2)
        net = conv(net, 512, 4, scope='gen/conv6', stride=2)
        net = conv(net, 512, 4, scope='gen/conv7', stride=2)
        net = conv(net, 512, 4, scope='gen/conv8', stride=2)
        net = conv(net, 512, 4, scope='gen/conv9', stride=2)

        net = deconv(net, 512, 4, scope='gen/deconv1', stride=2, concat=False)
        net = deconv(net, 512, 4, scope='gen/deconv2', stride=2)
        net = deconv(net, 512, 4, scope='gen/deconv3', stride=2)
        net = deconv(net, 512, 4, scope='gen/deconv4', stride=2)
        net = deconv(net, 512, 4, scope='gen/deconv5', stride=2)
        net = deconv(net, 256, 4, scope='gen/deconv6', stride=2)
        net = deconv(net, 128, 4, scope='gen/deconv7', stride=2)
        net = deconv(net, 64, 4, scope='gen/deconv8', stride=[1, 2])

        out_fg = deconv(net,
                        2,
                        4,
                        scope='gen/fg',
                        stride=[1, 2],
                        normalizer_fn=None,
                        do_pop=False)
        out_bg = deconv(net,
                        2,
                        4,
                        scope='gen/bg',
                        stride=[1, 2],
                        normalizer_fn=None,
                        do_pop=False)

        def process(out):
            pred_spec = out[..., 0]
            pred_spec = tf.tanh(pred_spec)
            pred_spec = unnormalize_spec(pred_spec, pr)

            pred_phase = out[..., 1]
            pred_phase = tf.tanh(pred_phase)
            pred_phase = unnormalize_phase(pred_phase, pr)

            val = soundrep.db_from_amp(0.) if pr.log_spec else 0.
            pred_spec = tf.pad(pred_spec,
                               [(0, 0), (0, 0),
                                (0, num_freq - shape(pred_spec, 2))],
                               constant_values=val)

            if pr.phase_type == 'pred':
                pred_phase = tf.concat([pred_phase, phase[..., -1:]], 2)
            elif pr.phase_type == 'orig':
                pred_phase = phase
            else:
                raise RuntimeError()

            # if ut.hastrue(pr, 'griffin_lim'):
            #   print 'using griffin-lim'
            #   pred_wav = griffin_lim(pred_spec, pred_phase, pr)
            # else:
            pred_wav = istft(pred_spec, pred_phase, pr)
            return pred_spec, pred_phase, pred_wav

        pred_spec_fg, pred_phase_fg, pred_wav_fg = process(out_fg)
        pred_spec_bg, pred_phase_bg, pred_wav_bg = process(out_bg)

        return ut.Struct(
            pred_spec_fg=pred_spec_fg,
            pred_wav_fg=pred_wav_fg,
            pred_phase_fg=pred_phase_fg,
            pred_spec_bg=pred_spec_bg,
            pred_phase_bg=pred_phase_bg,
            pred_wav_bg=pred_wav_bg,
            vid_net=vid_net_full,
        )
Пример #7
0
def make_net(ims, sfs, pr, im_net=None, reuse=True, train=True):
    if pr.subsample_frames and ims is not None:
        ims = ims[:, ::pr.subsample_frames]

    im_scales = []
    scales = []
    with slim.arg_scope(arg_scope(pr=pr, reuse=reuse, train=train)):
        sf_net = normalize_sfs(sfs)
        sf_net = ed(sf_net, 2)
        sf_net = conv2d(sf_net,
                        64, [65, 1],
                        scope='sf/conv1_1',
                        stride=4,
                        reuse=reuse)
        sf_net = slim.max_pool2d(sf_net, [4, 1], [4, 1])

        sf_net = block2(sf_net,
                        128, [15, 1],
                        'sf/conv2_1',
                        stride=[4, 1],
                        reuse=reuse)
        sf_net = block2(sf_net,
                        128, [15, 1],
                        'sf/conv3_1',
                        stride=[4, 1],
                        reuse=reuse)
        sf_net = block2(sf_net,
                        256, [15, 1],
                        'sf/conv4_1',
                        stride=[4, 1],
                        reuse=reuse)

        if im_net is None:
            im_net = mu.normalize_ims(ims)
            im_net = conv3d(im_net, 64, [5, 7, 7], scope='im/conv1', stride=2)
            im_net = pool3d(im_net, [1, 3, 3], [1, 2, 2])

            im_net = block3(im_net,
                            64, [3, 3, 3],
                            'im/conv2_1',
                            stride=1,
                            reuse=reuse)
            im_net = block3(im_net,
                            64, [3, 3, 3],
                            'im/conv2_2',
                            stride=2,
                            reuse=reuse)
            scales.append(im_net)
            im_scales.append(im_net)

        net = merge(sf_net, im_net, pr, reuse, train=train)
        net = block3(net, 128, [3, 3, 3], 'im/conv3_1', stride=1, reuse=reuse)
        net = block3(net, 128, [3, 3, 3], 'im/conv3_2', stride=1, reuse=reuse)
        scales.append(net)
        im_scales.append(net)

        #s = (2 if ut.hastrue(pr, 'cam2') else 1)
        s = 2
        net = block3(net,
                     256, [3, 3, 3],
                     'im/conv4_1',
                     stride=[2, s, s],
                     reuse=reuse)
        net = block3(net, 256, [3, 3, 3], 'im/conv4_2', stride=1, reuse=reuse)
        im_scales.append(net)

        time_stride = (2 if ut.hastrue(pr, 'extra_stride') else 1)
        print 'time_stride =', time_stride
        s = (1 if pr.cam else 2)
        net = block3(net,
                     512, [3, 3, 3],
                     'im/conv5_1',
                     stride=[time_stride, s, s],
                     reuse=reuse)
        net = block3(net, 512, [3, 3, 3], 'im/conv5_2', stride=1, reuse=reuse)
        scales.append(net)
        im_scales.append(net)

        last_conv = net
        net = tf.reduce_mean(net, [1, 2, 3], keep_dims=True)

        logits_name = 'joint/logits'
        logits = conv3d(net,
                        1, [1, 1, 1],
                        scope=logits_name,
                        activation_fn=None,
                        normalizer_fn=None)[:, 0, 0, 0, :]
        cam = conv3d(last_conv,
                     1, [1, 1, 1],
                     scope=logits_name,
                     activation_fn=None,
                     normalizer_fn=None,
                     reuse=True)

        return ut.Struct(logits=logits,
                         cam=cam,
                         last_conv=last_conv,
                         im_net=im_net,
                         scales=scales,
                         im_scales=im_scales)
Пример #8
0
    def make_model(self):
        with tf.device(self.default_gpu):
            pr = self.pr
            # steps
            self.step = tf.get_variable('global_step', [],
                                        trainable=False,
                                        initializer=tf.constant_initializer(0),
                                        dtype=tf.int64)
            self.lr = tf.constant(pr.base_lr)

            # model
            opt = make_opt(pr.opt_method, pr.base_lr, pr)
            self.inputs = read_data(pr, self.gpus)

            gpu_grads, gpu_losses = {}, {}
            for i, gpu in enumerate(self.gpus):
                with tf.device(gpu):
                    reuse = (i > 0)
                    with tf.device('/cpu:0'):
                        ims = self.inputs[i]['ims']
                        samples_ex = self.inputs[i]['samples']
                        assert pr.both_examples
                        assert not pr.small_augment
                        labels = tf.random_uniform([shape(ims, 0)],
                                                   0,
                                                   2,
                                                   dtype=tf.int64,
                                                   name='labels_sample')
                        samples0 = tf.where(tf.equal(labels, 1),
                                            samples_ex[:, 1], samples_ex[:, 0])
                        samples1 = tf.where(tf.equal(labels, 0),
                                            samples_ex[:, 1], samples_ex[:, 0])
                        labels1 = 1 - labels

                    net0 = make_net(ims,
                                    samples0,
                                    pr,
                                    reuse=reuse,
                                    train=self.is_training)
                    net1 = make_net(None,
                                    samples1,
                                    pr,
                                    im_net=net0.im_net,
                                    reuse=True,
                                    train=self.is_training)
                    labels = tf.concat([labels, labels1], 0)
                    net = ut.Struct(
                        logits=tf.concat([net0.logits, net1.logits], 0),
                        cam=tf.concat([net0.cam, net1.cam], 0),
                        last_conv=tf.concat([net0.last_conv, net1.last_conv],
                                            0))

                    loss = mu.Loss('loss')
                    loss.add_loss(slim_losses_with_prefix(None), 'reg')
                    loss.add_loss_acc(sigmoid_loss(net.logits, labels),
                                      'label')
                    grads = opt.compute_gradients(loss.total_loss())

                    ut.add_dict_list(gpu_grads, loss.name, grads)
                    ut.add_dict_list(gpu_losses, loss.name, loss)
                    #self.loss = loss

                    if i == 0:
                        self.net = net

            self.loss = mu.merge_losses(gpu_losses['loss'])
            for name, val in zip(self.loss.get_loss_names(),
                                 self.loss.get_losses()):
                tf.summary.scalar(name, val)

            if not self.is_training:
                #pr_test = pr.copy()
                pr_test = self.pr_test.copy()
                pr_test.augment_ims = False
                print 'pr_test ='
                print pr_test

                self.test_ims, self.test_samples, self.test_ytids = mu.on_cpu(
                    lambda: shift_dset.make_db_reader(
                        pr_test.test_list,
                        pr_test,
                        pr.test_batch, ['im', 'samples', 'ytid'],
                        one_pass=True))

                if pr_test.do_shift:
                    self.test_labels = tf.random_uniform(
                        [shape(self.test_ims, 0)], 0, 2, dtype=tf.int64)
                    self.test_samples = tf.where(tf.equal(self.test_labels, 1),
                                                 self.test_samples[:, 1],
                                                 self.test_samples[:, 0])
                else:
                    self.test_labels = tf.ones(shape(self.test_ims, 0),
                                               dtype=tf.int64)
                    #self.test_samples = tf.where(tf.equal(self.test_labels, 1), self.test_samples[:, 1], self.test_samples[:, 0])
                    print 'sample shape:', shape(self.test_samples)

                self.test_net = make_net(self.test_ims,
                                         self.test_samples,
                                         pr_test,
                                         reuse=True,
                                         train=self.is_training)

            (gs, vs) = zip(*mu.average_grads(gpu_grads['loss']))
            if pr.grad_clip is not None:
                gs, _ = tf.clip_by_global_norm(gs, pr.grad_clip)
            gs = [
                mu.print_every(gs[0], 100,
                               ['grad norm:', tf.global_norm(gs)])
            ] + list(gs[1:])
            gvs = zip(gs, vs)

            bn_ups = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
            if pr.multipass:
                ops = [
                    opt.apply_gradients(gvs, global_step=self.step)
                    for i in xrange(pr.multipass_count)
                ]

                def op_helper(count=[0]):
                    op = ops[count[0] % len(ops)]
                    count[0] += 1
                    return op

                self.train_op = op_helper
            else:
                op = tf.group(opt.apply_gradients(gvs, global_step=self.step),
                              *bn_ups)
                self.train_op = lambda: op

            self.coord = tf.train.Coordinator()
            self.saver = tf.train.Saver()

            self.init_op = tf.group(tf.global_variables_initializer(),
                                    tf.local_variables_initializer())
            self.sess.run(self.init_op)
            tf.train.start_queue_runners(sess=self.sess, coord=self.coord)

            self.merged_summary = tf.summary.merge_all()
            print 'Tensorboard command:'
            summary_dir = ut.mkdir(pj(pr.summary_dir, ut.simple_timestamp()))
            print 'tensorboard --logdir=%s' % summary_dir
            self.sum_writer = tf.summary.FileWriter(summary_dir,
                                                    self.sess.graph)
def test(pr, gpu, test_on_train=False, crop_type='center'):
    [gpu] = set_gpus([gpu])

    if ut.hastrue(pr, 'use_clf'):
        net = SVMClf(pr)
    else:
        #check_path = tf.train.latest_checkpoint(pr.train_dir)
        check_path = pj(pr.train_dir, 'net.tf-%d' % pr.model_iter)
        print 'Restoring from:', check_path
        net = NetClf(pr, check_path, gpu)

    if test_on_train:
        print 'Testing on train!'
        data = ut.load(pj(pr.dsdir, 'train.pk'))
    else:
        data = ut.load(pj(pr.dsdir, 'test.pk'))

    labels, probs, accs, vals = [], [], [], []
    for i in xrange(len(data)):
        ex = data[i]
        label = ex['is_gripping']

        def load_im(k, v):
            if k.startswith('gel') or k.startswith('im'):
                im = ig.uncompress(v)
            elif k.startswith('depth'):
                #v = np.tile(v, (1, 1, 3))
                im = v.astype('float32')
            else:
                raise RuntimeError()

            if crop_type == 'center':
                crops = [ut.crop_center(im, 224)]
            elif crop_type == 'multi':
                crops = []
                dh = (im.shape[0] - crop_dim)
                num_dim_samples = 3
                for y in np.linspace(0, dh, num_dim_samples).astype('l'):
                    dw = (im.shape[1] - crop_dim)
                    for x in np.linspace(0, dw, num_dim_samples).astype('l'):
                        crops.append(im[y:y + crop_dim, x:x + crop_dim])
            return ut.shuffled_with_seed(crops, k.split('_')[0] + str(i))

        all_inputs = {k: load_im(k, ex[k]) for k in im_names}
        ps = []
        for j in xrange(len(all_inputs['gel0_pre'])):
            inputs = {k: all_inputs[k][j] for k in im_names}
            inputs['initial_press_prob'] = ex['initial_press_prob']
            inputs['ee'] = ex['end_effector']
            _, prob = net.predict(**inputs)
            ps.append(prob)
        prob = np.mean(ps)
        pred = int(prob >= net.thresh)

        print prob, pred, label
        labels.append(label)
        probs.append(prob)
        accs.append(pred == label)
        print 'running average acc:', np.mean(accs)
        vals.append(
            ut.Struct(label=label,
                      prob=prob,
                      acc=accs[-1],
                      idx=i,
                      db_file=ex['db_file'],
                      object_name=ex['object_name']))

    labels = np.array(labels, 'bool')
    probs = np.array(probs, 'float32')
    accs = np.array(accs)

    acc = np.mean(accs)
    ap = sklearn.metrics.average_precision_score(labels, probs)
    print 'Accuracy:', acc
    print 'mAP:', ap
    print 'Base rate:', ut.f3(
        np.array(ut.mapattr(vals).label).astype('float32').mean())

    ut.save(pj(pr.resdir, 'eval_results.pk'),
            dict(acc=acc, ap=ap, results=(labels, probs)))
    ut.save(pj(pr.resdir, 'eval.pk'), vals)