small-image-descriptors/DescriptorManager.cpp

#include "DescriptorManager.h"

DescriptorManager::DescriptorManager() {} //don't use this
DescriptorManager::DescriptorManager(const fs::path & image_path) {
	img_path = image_path;
	if (!fs::exists(img_path)) {
		cout << "Error. Data folder not found." << endl;
		return; //No data: giving up.
	}
	numImages = std::count_if(
		fs::directory_iterator(img_path),
		fs::directory_iterator(),
		static_cast<bool(*)(const fs::path&)>(fs::is_regular_file));
	current_image = NULL;
	//Manual definitions.. SAD!
	descriptor_types[7] = new DCTMedian();
	descriptor_types[6] = new DCTPearson();
	descriptor_types[5] = new threshold();
	descriptor_types[4] = new pearson();
	descriptor_types[3] = new sobel();
	descriptor_types[2] = new median();
	descriptor_types[1] = new numColors();
	descriptor_types[0] = new averageColor();

	mutators[0] = new Grayscale();
	mutators[1] = new Crop();
	mutators[2] = new Flop();
	mutators[3] = new JPEG();
	mutators[4] = new LowQ();
	mutators[5] = new Hue();
	mutators[6] = new Brighten();
	mutators[7] = new Watermark();

	//Check if all folders are available, create if needed
	
	if (!fs::is_directory("meta"))
		fs::create_directory("meta");
	if (!fs::is_directory("results"))
		fs::create_directory("results");
	
	for (int i = 0; i < numDescriptors; i++) {
		fs::path descriptorBase("meta" / descriptor_types[i]->ToPath());
		fs::path resultFolder("results" / descriptor_types[i]->ToPath());
		if (!fs::is_directory(descriptorBase)) fs::create_directory(descriptorBase);
		if (!fs::is_directory(resultFolder)) fs::create_directory(resultFolder);
		for (int j = 0; j < numMutations; j++) {
			fs::path mutatorDescriptor(descriptorBase / mutators[j]->ToPath());
			if (!fs::is_directory(mutatorDescriptor)) fs::create_directory(mutatorDescriptor);
		}
	}
	for (int i = 0; i < numMutations; i++) {	
		fs::path mutationBase(img_path / mutators[i]->ToPath());
		if (!fs::is_directory(mutationBase)) fs::create_directory(mutationBase);
		
	}
}

//Returns a pointer to the image mutated by the given method.
//Either gets it from disk, or calculates it directly.
Image* DescriptorManager::GetMutated(Image* image, string image_name, int method) {
	fs::path mutPath(img_path / mutators[method]->ToPath() / image_name);
	Image* mutated;
	if (!fs::exists(mutPath)) { //only mutate if nothing on disk
		mutated = mutators[method]->getMutated(image);
		mutated->write(mutPath.generic_string());
	}
	else { 
		mutated = new Image;
		mutated->read(mutPath.generic_string());
	}
	return mutated;
}

//Fills the features array with all descriptors for the given image.
void DescriptorManager::LoadFeatures(string image_path, string image_name) {
	for (int i = 0; i < numDescriptors; i++) {
		//First check if we already have this information from the loaded cache
		if (features[i].count(image_name) > 0) {
			continue;
		}
		current_image = new Image(image_path);
		fs::path featurePath("meta" / descriptor_types[i]->ToPath() / image_name);

		uint64_t b1;
		if (fs::exists(featurePath)) { //feature on disk, load it in
			fs::ifstream ifs(featurePath, fs::fstream::binary);
			ifs.read(reinterpret_cast<char*>(&b1), sizeof(b1));
			ifs.close();
		}
		else { //not on disk, calculate it
			b1 = descriptor_types[i]->feature(current_image);
			fs::ofstream ofs(featurePath, fs::fstream::binary);
			ofs.write(reinterpret_cast<const char*>(&b1), sizeof(b1));
			ofs.close();
		}
		
		features[i][image_name] = b1;
	}
}

//Fills the mutatedFeatures array with all descriptors for the given mutation and image.
void DescriptorManager::LoadMutationFeatures(string image_path, string image_name, int mutation) {
	for (int i = 0; i < numDescriptors; i++) {
		//First check if we already have this information from the loaded cache
		if (mutatedFeatures[i][mutation].count(image_name) > 0) {
			continue;
		}
		if (!current_image) current_image->read(image_path);
		Image* image = GetMutated(current_image, image_name, i);
		fs::path featurePath("meta" / descriptor_types[i]->ToPath() / mutators[mutation]->ToPath() / image_name);

		uint64_t b1;
		if (fs::exists(featurePath)) { //feature on disk, load it in
			fs::ifstream ifs(featurePath, fs::fstream::binary);
			ifs.read(reinterpret_cast<char*>(&b1), sizeof(b1));
			ifs.close();
		}
		else { //not on disk, calculate it
			b1 = descriptor_types[i]->feature(image);
			fs::ofstream ofs(featurePath, fs::fstream::binary);
			ofs.write(reinterpret_cast<const char*>(&b1), sizeof(b1));
			ofs.close();
		}
		mutatedFeatures[i][mutation][image_name] = b1;	
		delete image;
	}
}

//Calculates all features, mutations and features of mutations.
//Do this before anything else.
void DescriptorManager::Init() {
	LoadFromDisk();
	int filecount = 0;
	for (fs::directory_iterator itr(img_path); itr != fs::directory_iterator(); ++itr) {
		if (!itr->path().has_extension()) continue; //skip directories and such
		string image_name = itr->path().filename().generic_string();
		string image_path = itr->path().string();
		{ //Progress bar
			filecount++; 
			cout << setw(60) << setfill('*') << "\r" << string(60, ' ') << flush;  
			cout << setfill('*') << setw(60) << right<< "\rProgress: " << filecount*100/numImages << "% - (" << filecount << "/" << numImages << "): " << image_name << flush;
		}
		//Store all base features in the features array
		LoadFeatures(image_path, image_name);
		//now calculate all the features for the mutated images as well
		for (int i = 0; i < numMutations; i++) {
			LoadMutationFeatures(image_path, image_name, i);
		}
		if (current_image) delete current_image;
		current_image = NULL;
		SaveToDisk();
	}
}


//Saves feature map to disk.
void DescriptorManager::SaveToDisk() {
	{ //Boost serialization
		ofstream f("features.map", ios::binary);
		ar::binary_oarchive oa(f);
		oa << features;

		ofstream fm("mutated_features.map", ios::binary);
		ar::binary_oarchive oam(fm);
		oam << mutatedFeatures;
	}
}

//Loads saved feature maps from disk.
void DescriptorManager::LoadFromDisk() {
	{ //Boost
		ifstream f("features.map", ios::binary);
		if (f.is_open()) {
			ar::binary_iarchive bi(f);
			bi >> features;
		}
		ifstream fm("mutated_features.map", ios::binary);
		if (fm.is_open()) {
			ar::binary_iarchive bim(fm);
			bim >> mutatedFeatures;
		}
	}
}

//Sorts all base (unmutated) images in order of distance from the given image
//After this function the distances array will be filled
void DescriptorManager::rankImages(string image_name, int method) {
	uint64_t target;
	if (features[method].count(image_name)) { //Check if this image was properly initialized
		target = features[method][image_name];
	}
	else {
		return;
	}
	distances.clear(); //start fresh
	for (map<string, uint64_t>::iterator it = features[method].begin(); it != features[method].end(); ++it) {
		uint64_t d = descriptor_types[method]->distance(target, it->second);
		pair<uint64_t, string> p(d, it->first);
		distances.insert(p); //pairs are compared by their first element..
	}
}

//Adds mutated descriptors of a given image to a copy of the distances array.
//It inserts these mutated descriptors in a sorted way
//Assumes the distances array is filled with the normal distances
multiset<pair<uint64_t, string>> DescriptorManager::rankMutations(string image_name, int method) {
	uint64_t target;
	if (features[method].count(image_name)) { //Check if this image was properly initialized
		target = features[method][image_name];
	}

	//cout << "Finding hits for " << image_name << " which has feature-number " << target.to_ullong() << endl;
	multiset<pair<uint64_t, string>> local_distance = distances; //get our local copy of the distances.
	for (int i = 0; i < numMutations; i++) {
		uint64_t b1 = mutatedFeatures[method][i][image_name];
		uint64_t d = descriptor_types[method]->distance(target, b1);
		pair<uint64_t, string> p(d, mutators[i]->ToString());
		local_distance.insert(p);
		//upper_bound
	}
	return local_distance;
}

//Takes all the calculated features and rankes 
void DescriptorManager::runExperiments() {
	fs::ofstream ofs;
	for (int i = 0; i < numDescriptors; i++) {
		for (map<string, uint64_t>::iterator it = features[i].begin(); it != features[i].end(); ++it) {
			//cout << "=== " << descriptor_types[i]->ToString() << ": " << it->first << " ===" << endl;
			rankImages(it->first, i);
			multiset<pair<uint64_t, string>> local_distance = rankMutations(it->first, i);

			//Now we can draw some conclusions from the received data:
			//Put all results in some useful files in the results directory:
			fs::path resultFolder("results" / descriptor_types[i]->ToPath());
			for (int i = 0; i < numMutations; i++) {
				int j = 0;
				fs::path resultPath(resultFolder / mutators[i]->ToPath());
				ofs.open(resultPath, ofstream::out | ofstream::app);
				for (std::multiset<pair<uint64_t, string>>::iterator it = local_distance.begin(); it != local_distance.end(); ++it) {
					if (it->second == mutators[i]->ToString()) {
						double percentile = (j + 1) / (double)(numImages+numMutations);

						ofs << (int)(percentile*100) << endl;
						//cout << "(" << it->first << ", " << it->second << ")" << endl;
					}
					j++;
				}
				ofs.close();
			}

		}
	}
}

//outputs the gathered data to iostream
//Tries to do some nice formatting and grouping of bits.
void DescriptorManager::outputMap(int a) {
	cout << "Outputting map " << descriptor_types[a]->ToString() << endl;
	map<string, uint64_t> a_map = features[a];
	if (a_map.begin() == a_map.end()) {
		cout << "------------------------" << endl << "No data found" << endl << endl;
	}

	for (map<string, uint64_t>::iterator it = a_map.begin(); it != a_map.end(); ++it) {
		string bits = bitset<64>(it->second).to_string();
		int groupsize = 8;
		
		cout << setfill(' ') << setw(14) << it->first << " => " << setw(20) << it->second << " (";

		for (int i=0; i < bits.length(); ++i) {
    		cout << bits[i];
			if(i % groupsize == groupsize-1 && bits.length()-i > groupsize) cout << " ";
		}
		 
		cout << ")" << endl;
	}
} 

//calls its buddy function to output all data to iostream
void DescriptorManager::outputMap() {
	for (int i = 0; i < numDescriptors; i++) {
		outputMap(i);
	}
}

//Outputs all generated features ordered on distance from the query image.
void DescriptorManager::outputRank(string imagename, int method) {
	cout << "Outputting rank.." << endl;
	for (std::multiset<pair<uint64_t, string>>::iterator it = distances.begin(); it != distances.end(); ++it) {
			cout << "(" << it->first << ", " << it->second << ")" << endl;
	}
}

//Just runs the preview functions of all the descriptors on given image.
//No need to init before this.
void DescriptorManager::generatePreviews(string imagename) {
	fs::path imagePath(img_path / imagename);
	Image* img = new Image(imagePath.generic_string());
	for (int i = 0; i < numDescriptors; i++) {
		descriptor_types[i]->preview(img);
	}
}