Added a Download Full Dataset and Download Captions option to the datasets page so they captions or dataset can be easily downloaded.

This commit is contained in:
Jaret Burkett
2026-08-12 21:37:35 -06:00
parent 81adcc2176
commit 742a4c8cef
3 changed files with 160 additions and 3 deletions

View File

@@ -4,24 +4,103 @@ import fs from 'fs';
import fsp from 'fs/promises';
import path from 'path';
import archiver from 'archiver';
import { getTrainingFolder } from '@/server/settings';
import { getDatasetsRoot, getTrainingFolder } from '@/server/settings';
export const runtime = 'nodejs'; // ensure Node APIs are available
export const dynamic = 'force-dynamic'; // long-running, non-cached
type PostBody = {
zipTarget: 'samples'; //only samples for now
jobName: string;
zipTarget: 'samples' | 'dataset' | 'dataset_captions';
jobName?: string; // required for 'samples'
datasetName?: string; // required for dataset targets
};
const CAPTION_EXTS = ['.txt', '.json', '.caption'];
async function resolveSafe(p: string) {
// resolve symlinks + normalize
return await fsp.realpath(p);
}
/**
* Recursively collects files to zip, skipping hidden entries (.thumbs etc).
* When captionsOnly is set, only caption files (.txt/.json/.caption) are kept.
*/
async function collectFiles(dir: string, captionsOnly: boolean): Promise<string[]> {
const entries = await fsp.readdir(dir, { withFileTypes: true });
let out: string[] = [];
for (const entry of entries) {
if (entry.name.startsWith('.')) continue;
const itemPath = path.join(dir, entry.name);
if (entry.isDirectory()) {
out = out.concat(await collectFiles(itemPath, captionsOnly));
} else if (entry.isFile()) {
if (!captionsOnly || CAPTION_EXTS.includes(path.extname(entry.name).toLowerCase())) {
out.push(itemPath);
}
}
}
return out;
}
async function zipDataset(datasetName: string, captionsOnly: boolean) {
const datasetsRoot = await resolveSafe(await getDatasetsRoot());
// basename strips any path segments so the name can't escape the datasets root
const safeName = path.basename(datasetName);
let folderPath: string;
try {
folderPath = await resolveSafe(path.join(datasetsRoot, safeName));
} catch {
return NextResponse.json({ error: `Dataset '${datasetName}' not found` }, { status: 404 });
}
const files = await collectFiles(folderPath, captionsOnly);
if (files.length === 0) {
return NextResponse.json({ error: captionsOnly ? 'No captions found' : 'Dataset is empty' }, { status: 404 });
}
// Hidden dir at the datasets root: skipped by the dataset list/image walkers,
// but still under datasetsRoot so /api/files/[...filePath] is allowed to serve it.
const zipDir = path.join(datasetsRoot, '.zips');
await fsp.mkdir(zipDir, { recursive: true });
const outputPath = path.join(zipDir, `${safeName}_${captionsOnly ? 'captions' : 'dataset'}.zip`);
await fsp.rm(outputPath, { force: true });
await new Promise<void>((resolve, reject) => {
const output = fs.createWriteStream(outputPath);
const archive = archiver('zip', { zlib: { level: 9 } });
output.on('close', () => resolve());
output.on('error', reject);
archive.on('error', reject);
archive.pipe(output);
for (const file of files) {
archive.file(file, { name: path.join(safeName, path.relative(folderPath, file)) });
}
archive.finalize().catch(reject);
});
return NextResponse.json({
ok: true,
zipPath: outputPath,
fileName: path.basename(outputPath),
});
}
export async function POST(request: NextRequest) {
try {
const body = (await request.json()) as PostBody;
if (body?.zipTarget === 'dataset' || body?.zipTarget === 'dataset_captions') {
if (!body.datasetName) {
return NextResponse.json({ error: 'datasetName is required' }, { status: 400 });
}
return await zipDataset(body.datasetName, body.zipTarget === 'dataset_captions');
}
if (!body || !body.jobName) {
return NextResponse.json({ error: 'jobName is required' }, { status: 400 });
}

View File

@@ -13,6 +13,7 @@ import { apiClient } from '@/utils/api';
import useSettings from '@/hooks/useSettings';
import { pathJoin } from '@/utils/basic';
import AutoCaptionButton from '@/components/AutoCaptionButton';
import DatasetActionBar from '@/components/DatasetActionBar';
import CaptionMonitor from '@/components/CaptionMonitor';
import { CreatableSelectInput } from '@/components/formInputs';
@@ -157,6 +158,7 @@ export default function DatasetPage({ params }: { params: { datasetName: string
<span className="sm:hidden">+ Add</span>
<span className="hidden sm:inline">Add Images</span>
</Button>
<DatasetActionBar datasetName={datasetName} />
</div>
</TopBar>
<MainContent

View File

@@ -0,0 +1,76 @@
'use client';
import { useState } from 'react';
import { Menu, MenuButton, MenuItem, MenuItems } from '@headlessui/react';
import { Cog, Download, Captions } from 'lucide-react';
import { LuLoader } from 'react-icons/lu';
import { apiClient } from '@/utils/api';
type DatasetZipTarget = 'dataset' | 'dataset_captions';
interface DatasetActionBarProps {
datasetName: string;
className?: string;
}
export default function DatasetActionBar({ datasetName, className }: DatasetActionBarProps) {
const [zipping, setZipping] = useState<DatasetZipTarget | null>(null);
const downloadZip = async (zipTarget: DatasetZipTarget) => {
if (zipping) return;
setZipping(zipTarget);
try {
const res = await apiClient.post('/api/zip', { zipTarget, datasetName });
const zipPath = res.data.zipPath;
if (!zipPath) throw new Error('No zipPath in response');
// Cache-buster: /api/files serves with a long max-age, and the zip is rebuilt
// at the same path every time, so a bare URL could hand back a stale download.
const downloadPath = `/api/files/${encodeURIComponent(zipPath)}?v=${Date.now()}`;
const a = document.createElement('a');
a.href = downloadPath;
a.download = res.data.fileName || `${datasetName}.zip`;
document.body.appendChild(a);
a.click();
a.remove();
} catch (err) {
console.error('Error downloading zip:', err);
} finally {
setZipping(null);
}
};
const iconSizeClass = 'w-5 h-5 sm:w-6 sm:h-6';
return (
<div className={`flex items-center flex-shrink-0 ${className ?? ''}`}>
<Menu>
<MenuButton className={'ml-1 sm:ml-2'}>
{zipping ? <LuLoader className={`${iconSizeClass} animate-spin`} /> : <Cog className={iconSizeClass} />}
</MenuButton>
<MenuItems
anchor={{ to: 'bottom end', gap: 16 }}
className="bg-gray-900 border border-gray-700 rounded shadow-lg w-64 px-2 py-2 z-50"
>
<MenuItem>
<div
className="cursor-pointer px-4 py-1 hover:bg-gray-800 rounded flex items-center gap-2"
onClick={() => downloadZip('dataset')}
>
<Download className="w-4 h-4" />
Download Full Dataset
</div>
</MenuItem>
<MenuItem>
<div
className="cursor-pointer px-4 py-1 hover:bg-gray-800 rounded flex items-center gap-2"
onClick={() => downloadZip('dataset_captions')}
>
<Captions className="w-4 h-4" />
Download Captions
</div>
</MenuItem>
</MenuItems>
</Menu>
</div>
);
}