You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
138 lines
4.4 KiB
138 lines
4.4 KiB
const express = require('express');
|
|
const bodyParser = require('body-parser');
|
|
const vptree = require('vptree');
|
|
const similarity = require('compute-cosine-similarity');
|
|
const fs = require('fs');
|
|
const path = require('path');
|
|
const app = express();
|
|
const port = 2000;
|
|
|
|
app.use(bodyParser.json());
|
|
|
|
// Function to load JSON data
|
|
function loadJSON(filePath) {
|
|
return new Promise((resolve, reject) => {
|
|
fs.readFile(filePath, 'utf8', (err, data) => {
|
|
if (err) {
|
|
return reject(err);
|
|
}
|
|
try {
|
|
const jsonData = JSON.parse(data);
|
|
resolve(jsonData);
|
|
} catch (parseError) {
|
|
reject(parseError);
|
|
}
|
|
});
|
|
});
|
|
}
|
|
|
|
function cosineDistanceMatching(vector1, vector2) {
|
|
let cosineSimilarity = similarity(vector1, vector2);
|
|
let distance = 2 * (1 - cosineSimilarity);
|
|
return Math.sqrt(distance);
|
|
}
|
|
|
|
let tree, vectors, jsonEmbeddings, filteredEmbeddings;
|
|
|
|
async function buildVPTree() {
|
|
// Load JSON data from embeddings.json
|
|
const jsonFilePath = path.join(__dirname, 'embeddings.json');
|
|
try {
|
|
jsonEmbeddings = await loadJSON(jsonFilePath);
|
|
|
|
// Filter embeddings
|
|
filteredEmbeddings = jsonEmbeddings.filter(item => item.lost !== 1);
|
|
vectors = filteredEmbeddings.map(item => item.vector);
|
|
|
|
tree = vptree.build(vectors, cosineDistanceMatching);
|
|
const treeString = tree.stringify();
|
|
const fileName = "sv-tree";
|
|
const filePath = path.join(__dirname, `${fileName}.txt`);
|
|
|
|
await fs.promises.writeFile(filePath, treeString, 'utf8');
|
|
console.log("Tree saved successfully.");
|
|
} catch (err) {
|
|
console.error('Error building VP-tree:', err);
|
|
}
|
|
}
|
|
|
|
async function loadVPTree() {
|
|
try {
|
|
const treeData = await loadTreeFromDisk("sv-tree");
|
|
// Rebuild the VP-tree using the saved structure
|
|
const jsonFilePath = path.join(__dirname, 'embeddings.json');
|
|
try {
|
|
jsonEmbeddings = await loadJSON(jsonFilePath);
|
|
filteredEmbeddings = jsonEmbeddings.filter(item => item.lost !== 1); // Filter again after loading
|
|
vectors = filteredEmbeddings.map(item => item.vector);
|
|
|
|
console.log(vectors.length);
|
|
tree = vptree.load(vectors, cosineDistanceMatching, treeData);
|
|
console.log('Tree loaded successfully.');
|
|
} catch (loadError) {
|
|
console.error('Error loading the VP-tree:', loadError);
|
|
}
|
|
} catch (err) {
|
|
console.error('Failed to load tree from disk:', err);
|
|
}
|
|
}
|
|
|
|
function loadTreeFromDisk(fileName) {
|
|
const filePath = path.join(__dirname, `${fileName}.txt`);
|
|
return new Promise((resolve, reject) => {
|
|
fs.readFile(filePath, 'utf8', (err, data) => {
|
|
if (err) {
|
|
console.error('Error reading the file:', err);
|
|
return reject(err);
|
|
}
|
|
|
|
// Parse the data as a JavaScript object
|
|
try {
|
|
const treeData = eval(`(${data})`); // Using eval to convert to object
|
|
resolve(treeData);
|
|
} catch (parseError) {
|
|
console.error('Error parsing the tree data:', parseError);
|
|
reject(parseError);
|
|
}
|
|
});
|
|
});
|
|
}
|
|
|
|
async function initialize() {
|
|
await buildVPTree();
|
|
await loadVPTree();
|
|
}
|
|
|
|
// Initialize the VP-tree building and loading process
|
|
initialize();
|
|
|
|
app.post('/search/', (req, res) => {
|
|
const { vector } = req.body;
|
|
|
|
if (!Array.isArray(vector) || vector.length !== 7 || !vector.every(num => typeof num === 'number')) {
|
|
return res.status(400).send({ success: false, message: 'Invalid vector provided. It must be an array of 7 floating-point numbers.' });
|
|
}
|
|
|
|
let idx = queryVPTree(vector);
|
|
if (idx !== undefined && idx < filteredEmbeddings.length) {
|
|
res.send({
|
|
folder: filteredEmbeddings[idx].folder,
|
|
image: filteredEmbeddings[idx].image,
|
|
video: filteredEmbeddings[idx].video,
|
|
frame: filteredEmbeddings[idx].frame,
|
|
lost: filteredEmbeddings[idx].lost
|
|
});
|
|
} else {
|
|
res.status(404).send({ success: false, message: 'No matching entry found.' });
|
|
}
|
|
});
|
|
|
|
function queryVPTree(value) {
|
|
let nearest = tree.search(value);
|
|
let index = nearest[0].i;
|
|
return index;
|
|
}
|
|
|
|
app.listen(port, () => {
|
|
console.log(`Server listening at http://localhost:${port}`);
|
|
});
|
|
|