121 lines
4.5 KiB
JavaScript
121 lines
4.5 KiB
JavaScript
const sequelize = require('../database/database.js');
|
|
const { Project, Img, Ann } = require('../models');
|
|
const { fetchLableStudioProject, fetchProjectIdsAndTitles } = require('./fetch-labelstudio.js');
|
|
|
|
const updateStatus = { running: false };
|
|
|
|
async function seedLabelStudio() {
|
|
updateStatus.running = true;
|
|
console.log('Seeding started');
|
|
try {
|
|
await sequelize.sync();
|
|
const projects = await fetchProjectIdsAndTitles();
|
|
|
|
for (const project of projects) {
|
|
console.log(`Processing project ${project.id} (${project.title})`);
|
|
|
|
// Upsert project in DB
|
|
await Project.upsert({ project_id: project.id, title: project.title });
|
|
|
|
// Fetch project data (annotations array)
|
|
const data = await fetchLableStudioProject(project.id);
|
|
if (!Array.isArray(data) || data.length === 0) {
|
|
console.log(`No annotation data for project ${project.id}`);
|
|
continue;
|
|
}
|
|
|
|
// Remove old images and annotations for this project
|
|
const oldImages = await Img.findAll({ where: { project_id: project.id } });
|
|
const oldImageIds = oldImages.map(img => img.image_id);
|
|
if (oldImageIds.length > 0) {
|
|
await Ann.destroy({ where: { image_id: oldImageIds } });
|
|
await Img.destroy({ where: { project_id: project.id } });
|
|
console.log(`Deleted ${oldImageIds.length} old images and their annotations for project ${project.id}`);
|
|
}
|
|
|
|
// Prepare arrays
|
|
const imagesBulk = [];
|
|
const annsBulk = [];
|
|
|
|
for (const ann of data) {
|
|
// Extract width/height
|
|
let width = null;
|
|
let height = null;
|
|
if (Array.isArray(ann.label_rectangles) && ann.label_rectangles.length > 0) {
|
|
width = ann.label_rectangles[0].original_width;
|
|
height = ann.label_rectangles[0].original_height;
|
|
} else if (Array.isArray(ann.label) && ann.label.length > 0 && ann.label[0].original_width && ann.label[0].original_height) {
|
|
width = ann.label[0].original_width;
|
|
height = ann.label[0].original_height;
|
|
}
|
|
|
|
// Only push image and annotations if width and height are valid
|
|
if (width && height) {
|
|
imagesBulk.push({
|
|
project_id: project.id,
|
|
image_path: ann.image,
|
|
width,
|
|
height
|
|
});
|
|
|
|
// Handle multiple annotations per image
|
|
if (Array.isArray(ann.label_rectangles)) {
|
|
for (const ann_detail of ann.label_rectangles) {
|
|
annsBulk.push({
|
|
image_path: ann.image,
|
|
x: (ann_detail.x * width) / 100,
|
|
y: (ann_detail.y * height) / 100,
|
|
width: (ann_detail.width * width) / 100,
|
|
height: (ann_detail.height * height) / 100,
|
|
Label: Array.isArray(ann_detail.rectanglelabels) ? (ann_detail.rectanglelabels[0] || 'unknown') : (ann_detail.rectanglelabels || 'unknown')
|
|
});
|
|
}
|
|
} else if (Array.isArray(ann.label)) {
|
|
for (const ann_detail of ann.label) {
|
|
annsBulk.push({
|
|
image_path: ann.image,
|
|
x: (ann_detail.x * width) / 100,
|
|
y: (ann_detail.y * height) / 100,
|
|
width: (ann_detail.width * width) / 100,
|
|
height: (ann_detail.height * height) / 100,
|
|
Label: Array.isArray(ann_detail.rectanglelabels) ? (ann_detail.rectanglelabels[0] || 'unknown') : (ann_detail.rectanglelabels || 'unknown')
|
|
});
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// 1) Insert images and get generated IDs
|
|
const insertedImages = await Img.bulkCreate(imagesBulk, { returning: true });
|
|
|
|
// 2) Map image_path -> image_id
|
|
const imageMap = {};
|
|
for (const img of insertedImages) {
|
|
imageMap[img.image_path] = img.image_id;
|
|
}
|
|
|
|
// 3) Assign correct image_id to each annotation
|
|
for (const ann of annsBulk) {
|
|
ann.image_id = imageMap[ann.image_path];
|
|
delete ann.image_path; // cleanup
|
|
}
|
|
|
|
// 4) Insert annotations
|
|
await Ann.bulkCreate(annsBulk);
|
|
|
|
console.log(`Inserted ${imagesBulk.length} images and ${annsBulk.length} annotations for project ${project.id}`);
|
|
}
|
|
|
|
console.log('Seeding done');
|
|
return { success: true, message: 'Data inserted successfully!' };
|
|
} catch (error) {
|
|
console.error('Error inserting data:', error);
|
|
return { success: false, message: error.message };
|
|
} finally {
|
|
updateStatus.running = false;
|
|
console.log('updateStatus.running set to false');
|
|
}
|
|
}
|
|
|
|
module.exports = { seedLabelStudio, updateStatus };
|