github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Added ability to set the caption extention in dataset viewer, captioner, and trainer so one dataset can have multiple caption styles in different files with different extensions. Added dataset caption template for a blank ideogram 4 formatted template.

Browse Source
This commit is contained in:
Jaret Burkett
2026-06-06 08:32:24 -06:00
parent 10cdeb394e
commit 41157b460c
19 changed files with 219 additions and 270 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
ui/src/app/api/caption/get/route.ts
View File

@@ -22,15 +22,16 @@ export async function POST(request: NextRequest) {
return new NextResponse(null, { status: 499 });
}
const { imgPath } = body;
const { imgPath, ext } = body;
console.log('Received POST request for caption:', imgPath);
try {
// Decode the path
const filepath = imgPath;
console.log('Decoded image path:', filepath);
// caption name is the filepath without extension but with .txt
const captionPath = filepath.replace(/\.[^/.]+$/, '') + '.txt';
// caption name is the filepath without extension but with the caption extension (default txt)
const captionExt = ((ext || 'txt') as string).replace(/^\.+/, '').trim() || 'txt';
const captionPath = filepath.replace(/\.[^/.]+$/, '') + '.' + captionExt;
// Get allowed directories
const allowedDir = await getDatasetsRoot();

5
ui/src/app/api/caption/getBatch/route.ts
View File

@@ -21,11 +21,12 @@ export async function POST(request: NextRequest) {
return new NextResponse(null, { status: 499 });
}
const { imgPaths } = body as { imgPaths?: string[] };
const { imgPaths, ext } = body as { imgPaths?: string[]; ext?: string };
if (!Array.isArray(imgPaths)) {
return NextResponse.json({ error: 'imgPaths must be an array' }, { status: 400 });
}
const captionExt = ((ext || 'txt') as string).replace(/^\.+/, '').trim() || 'txt';
const allowedDir = await getDatasetsRoot();
const captions: Record<string, string> = {};
@@ -33,7 +34,7 @@ export async function POST(request: NextRequest) {
if (typeof imgPath !== 'string') continue;
if (!isUnderRoot(imgPath, allowedDir)) continue;
const captionPath = imgPath.replace(/\.[^/.]+$/, '') + '.txt';
const captionPath = imgPath.replace(/\.[^/.]+$/, '') + '.' + captionExt;
try {
captions[imgPath] = fs.existsSync(captionPath) ? fs.readFileSync(captionPath, 'utf-8') : '';
} catch {

7
ui/src/app/api/img/caption/route.ts
View File

@@ -5,7 +5,7 @@ import { getDatasetsRoot } from '@/server/settings';
export async function POST(request: Request) {
try {
const body = await request.json();
const { imgPath, caption } = body;
const { imgPath, caption, ext } = body;
let datasetsPath = await getDatasetsRoot();
// make sure the dataset path is in the image path
if (!imgPath.startsWith(datasetsPath)) {
@@ -17,8 +17,9 @@ export async function POST(request: Request) {
return NextResponse.json({ error: 'Image does not exist' }, { status: 404 });
}
// check for caption
const captionPath = imgPath.replace(/\.[^/.]+$/, '') + '.txt';
// check for caption (default extension txt)
const captionExt = ((ext || 'txt') as string).replace(/^\.+/, '').trim() || 'txt';
const captionPath = imgPath.replace(/\.[^/.]+$/, '') + '.' + captionExt;
// save caption to file
fs.writeFileSync(captionPath, caption);