<commit_before>"""
A searchable collection of images, listening for jobs on a Redis queue.
"""
import sys
import operator
import time
import bottleneck as bn
import pandas as pd
import numpy as np
import sklearn.metrics.pairwise as metrics
import vislab
import vislab.utils.redis_q
import vislab.datasets


feats_dir = vislab.config['paths']['feats']
experiment_dir = vislab.config['paths']['shared_data'] + '/results_mar23'
feat_filenames = {
    'flickr': {
        'caffe fc6': feats_dir + '/flickr/caffe_fc6.h5',
        'caffe fc7': feats_dir + '/flickr/caffe_fc7.h5',
    }
}
dataset_loaders = {
    'flickr': vislab.datasets.flickr.get_df
}


class SearchableCollection(object):
    """
    A searchable collection of images for a given dataset.
    """
    def __init__(self, dataset_name):
        assert(dataset_name in dataset_loaders)
        t = time.time()

        # Load image information in a dataframe.
        self.images = dataset_loaders[dataset_name]()

        # Downsample [optional for dev].
        self.images = self.images.iloc[:10000]

        # Load all features.
        self.features = {}
        for feature_name, filename in feat_filenames[dataset_name].iteritems():
            try:
                feats = pd.read_hdf(filename, 'df')
            except:
                feats = pd.read_pickle(filename)
            feats = feats.ix[self.images.index]
            self.features[feature_name] = feats

        # Append predictions to the images DataFrame.
        if 'style scores' in self.features:
            self.images = self.images.join(
                self.features['style scores'], rsuffix='preds')

        print('Initialized SearchableCollection in {:.3f} s'.format(
            time.time() - t))

    def nn_by_id_many_filters(
            self, image_id, feature, distance,
            page=1, filter_conditions_list=None, results_per_page=8):
        """
        Return several sets of results, each filtered by different
        filter_conditions.
        """
        print('nn_by_id_many_filters', feature)
        assert(feature in self.features)
        t = time.time()

        # Get all distances to query
        images = self.images
        feats = self.features[feature]
        feat = feats.ix[image_id].values
        feats = feats.ix[images.index].values

        nn_ind, nn_dist = nn(feat, feats, distance)
        nn_df = pd.DataFrame(
            {'distance': nn_dist[1:]},
            images.index[nn_ind[1:]])
        images_nn_df = nn_df.join(images)

        # TODO: figure out what to do for paging
        start_ind = 0
        end_ind = start_ind + results_per_page

        results_sets = []
        for conditions in filter_conditions_list:
            result_images = filter_df(images_nn_df, conditions)
            num_results = result_images.shape[0]
            result_images = result_images[start_ind:end_ind]
            result_images['image_id'] = result_images.index
            results = [row.to_dict() for _, row in result_images.iterrows()]
            results_data = {
                'results': results,
                'start_ind': start_ind,
                'page': page,
                'num_results': num_results,
                'time_elapsed': time.time() - t<commit_msg>significantly sped up SearchableCollection<commit_after>"""
A searchable collection of images, listening for jobs on a Redis queue.
"""
import sys
import operator
import time
import bottleneck as bn
import pandas as pd
import numpy as np
import sklearn.metrics.pairwise as metrics
import vislab
import vislab.utils.redis_q
import vislab.datasets


feats_dir = vislab.config['paths']['feats']
experiment_dir = vislab.config['paths']['shared_data'] + '/results_mar23'
feat_filenames = {
    'flickr': {
        'caffe fc6': feats_dir + '/flickr/caffe_fc6.h5',
        'caffe fc7': feats_dir + '/flickr/caffe_fc7.h5',
    }
}
dataset_loaders = {
    'flickr': vislab.datasets.flickr.get_df
}


class SearchableCollection(object):
    """
    A searchable collection of images for a given dataset.
    """
    def __init__(self, dataset_name):
        assert(dataset_name in dataset_loaders)
        t = time.time()

        # Load image information in a dataframe.
        self.images = dataset_loaders[dataset_name]()

        # Downsample [optional for dev].
        # self.images = self.images.iloc[:10000]

        # Cache the image index.
        self.index = self.images.index.tolist()

        # Load all features.
        self.features = {}
        self.features_index = {}
        for feature_name, filename in feat_filenames[dataset_name].iteritems():
            try:
                feats = pd.read_hdf(filename, 'df')
            except:
                feats = pd.read_pickle(filename)
            self.features[feature_name] = feats.ix[self.images.index].values

        # Append predictions to the images DataFrame.
        if 'style scores' in self.features:
            self.images = self.images.join(
                self.features['style scores'], rsuffix='preds')

        print('Initialized SearchableCollection in {:.3f} s'.format(
            time.time() - t))

    def nn_by_id_many_filters(
            self, image_id, feature, distance,
            page=1, filter_conditions_list=None, results_per_page=8):
        """
        Return several sets of results, each filtered by different
        filter_conditions.
        """
        print('nn_by_id_many_filters', feature)
        assert(feature in self.features)
        t = time.time()

        # Get all distances to query
        images = self.images
        feats = self.features[feature]
        feat = feats[self.index.index(image_id)]

        print feats.shape
        print feat.shape

        print('elapsed: {:.3f}'.format(time.time() - t))

        nn_ind, nn_dist = nn(feat, feats, distance)
        nn_df = pd.DataFrame(
            {'distance': nn_dist[1:]},
            images.index[nn_ind[1:]])
        images_nn_df = nn_df.join(images)

        # TODO: figure out what to do for paging
        start_ind = 0
        end_ind = start_ind + results_per_page

        results_sets = []
        for conditions in filter_conditions_list:
            result_images = filter_df(images_nn_df, conditions)
            num_results = result_images.shape[0]
            result_images = result_images[start_ind:end_ind]
            result_images['image_id'] = result_images.index
            results = [row.to_dict() for _, row in result_images.iterrows()]
            results_data = {
                'results': results,
                'start_ind': start_ind,
                'page': page,
                'num_results': num_results,
                'time_elapsed': time.time() - t
{'fraction_non_alphanumeric': 0.1020377807526402, 'fraction_numerical': 0.006098467945857504, 'mean_word_length': 2.300932744231713, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2624259', 'n_tokens_mistral': 2037, 'n_tokens_neox': 1971, 'n_words': 474}