def get_featurize_fn()

in conv_lstm_utils.py [0:0]

• 69 lines of code
• 13 McCabe index (conditional complexity)

def get_featurize_fn(args, pers):
    normalize = get_normalize_fn(args)
    featurizer = args.featurizer

    # TODO:
    # * check that it has enemy units that we see in input
    # * add resources for us (and maybe resources for the enemy in target)
    def featurize_reduced(fn):
        with np.load(fn) as data:
            info = data['info']
            race = [int(info[0]), int(info[3])]
            map_size = (info[6], info[7])
            # XXX and another copy...
            rdata = data['reduced'].tobytes()

            mfdata = io.BytesIO(data['map'])
        with gzip.GzipFile(fileobj=mfdata, mode='r') as f:
            map_features = np.load(f)

        vis = args.use_true_visibility
        feats = featurizer.featurize_reduced_all(
            rdata, args.skip_frames, args.combine_frames, map_size, pers,
            visibility=vis, ground_height=map_features[2])
        input = feats[0]
        targets = feats[1]
        if vis:
            visibility = feats[2]
        else:
            visibility = input.sum(-1, keepdims=True) > 0
        assert input.shape[0] > 1

        featurized = (
            # map
            th.from_numpy(map_features[np.newaxis, :]).type(th.FloatTensor),
            # race
            th.LongTensor([race if pers == 0 else list(reversed(race))]),
            # input
            th.from_numpy(normalize(input.transpose(0, 3, 1, 2))[:-1]).type(th.FloatTensor),
            # targets
            th.from_numpy(normalize(targets.transpose(0, 3, 1, 2))[1:]).type(th.FloatTensor),
            # game name
            fn,
            # visibility
            visibility.transpose(0, 3, 1, 2),
        )
        return featurized

    def featurize(fn):
        raise RuntimeError("This doesn't work anymore, add in game name, visibility")
        rep = replayer.load(fn)
        map = rep.getMap()
        map_size = map['walkability'].shape
        race = [x[0] for x in get_info(rep)]
        sl = make_start_loc(map['start_locations'], map_size)

        batch = []
        fromt = args.from_time if args.until_time != 0 else 0
        until = args.from_time if args.until_time != 0 else len(rep)
        for i in range(fromt, min(len(rep), until), args.skip_frames):
            frames = [rep.getFrame(k) for k in
                      range(i, min(i + args.combine_frames, len(rep)))]
            batch.append(frames)

        featurized = (
            # map
            th.from_numpy(
                np.stack([
                    map['walkability'],
                    map['buildability'],
                    map['ground_height'],
                    sl,
                ], axis=0)[np.newaxis, :]).type(th.FloatTensor),
            # race
            th.LongTensor([race if pers == 0 else list(reversed(race))]),
            # input
            th.from_numpy(
                normalize(featurizer.featurize(
                    batch, map_size, perspective=pers
                ).transpose(0, 3, 1, 2))[:-1]).type(th.FloatTensor),
            # targets
            th.from_numpy(
                normalize(featurizer.featurize(
                    batch, map_size, perspective=pers, full=True
                ).transpose(0, 3, 1, 2))[1:]).type(th.FloatTensor),
        )
        assert featurized[2].size(0) > 1

        return featurized

    if args.reduced:
        return featurize_reduced
    return featurize