88 lines
2.8 KiB
C++
88 lines
2.8 KiB
C++
#pragma once
|
|
|
|
#include <limits>
|
|
#include <string>
|
|
#include <unordered_map>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
namespace rk3588 {
|
|
|
|
struct GallerySearchEntry {
|
|
int person_id = -1;
|
|
std::string name;
|
|
std::vector<float> emb;
|
|
};
|
|
|
|
struct GallerySearchResult {
|
|
int best_person_id = -1;
|
|
std::string best_name;
|
|
float best_sim = 0.0f;
|
|
float second_sim = 0.0f;
|
|
};
|
|
|
|
template <typename EntryT, typename PersonIdFn, typename NameFn, typename EmbFn>
|
|
GallerySearchResult SearchTop2ByPersonGeneric(const std::vector<EntryT>& entries,
|
|
const std::vector<float>& emb_normed,
|
|
PersonIdFn person_id_fn,
|
|
NameFn name_fn,
|
|
EmbFn emb_fn) {
|
|
GallerySearchResult r;
|
|
if (entries.empty() || emb_normed.empty()) return r;
|
|
|
|
std::unordered_map<int, std::pair<std::string, float>> per_person_best;
|
|
per_person_best.reserve(entries.size());
|
|
|
|
for (const auto& entry : entries) {
|
|
const auto& entry_emb = emb_fn(entry);
|
|
if (entry_emb.size() != emb_normed.size()) continue;
|
|
|
|
float sim = 0.0f;
|
|
for (size_t i = 0; i < emb_normed.size(); ++i) {
|
|
sim += emb_normed[i] * entry_emb[i];
|
|
}
|
|
|
|
const int person_id = person_id_fn(entry);
|
|
auto it = per_person_best.find(person_id);
|
|
if (it == per_person_best.end() || sim > it->second.second) {
|
|
per_person_best[person_id] = std::make_pair(name_fn(entry), sim);
|
|
}
|
|
}
|
|
|
|
float best = -std::numeric_limits<float>::infinity();
|
|
float second = -std::numeric_limits<float>::infinity();
|
|
int best_person_id = -1;
|
|
std::string best_name;
|
|
|
|
for (const auto& kv : per_person_best) {
|
|
const float sim = kv.second.second;
|
|
if (sim > best) {
|
|
second = best;
|
|
best = sim;
|
|
best_person_id = kv.first;
|
|
best_name = kv.second.first;
|
|
} else if (sim > second) {
|
|
second = sim;
|
|
}
|
|
}
|
|
|
|
if (best_person_id >= 0) {
|
|
r.best_person_id = best_person_id;
|
|
r.best_name = best_name;
|
|
r.best_sim = best;
|
|
r.second_sim = std::isfinite(second) ? second : 0.0f;
|
|
}
|
|
return r;
|
|
}
|
|
|
|
inline GallerySearchResult SearchTop2ByPerson(const std::vector<GallerySearchEntry>& entries,
|
|
const std::vector<float>& emb_normed) {
|
|
return SearchTop2ByPersonGeneric(
|
|
entries, emb_normed,
|
|
[](const GallerySearchEntry& entry) { return entry.person_id; },
|
|
[](const GallerySearchEntry& entry) -> const std::string& { return entry.name; },
|
|
[](const GallerySearchEntry& entry) -> const std::vector<float>& { return entry.emb; });
|
|
}
|
|
|
|
} // namespace rk3588
|