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:
@@ -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 });
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
76
ui/src/components/DatasetActionBar.tsx
Normal file
76
ui/src/components/DatasetActionBar.tsx
Normal 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>
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user