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 fsp from 'fs/promises';
|
||||||
import path from 'path';
|
import path from 'path';
|
||||||
import archiver from 'archiver';
|
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 runtime = 'nodejs'; // ensure Node APIs are available
|
||||||
export const dynamic = 'force-dynamic'; // long-running, non-cached
|
export const dynamic = 'force-dynamic'; // long-running, non-cached
|
||||||
|
|
||||||
type PostBody = {
|
type PostBody = {
|
||||||
zipTarget: 'samples'; //only samples for now
|
zipTarget: 'samples' | 'dataset' | 'dataset_captions';
|
||||||
jobName: string;
|
jobName?: string; // required for 'samples'
|
||||||
|
datasetName?: string; // required for dataset targets
|
||||||
};
|
};
|
||||||
|
|
||||||
|
const CAPTION_EXTS = ['.txt', '.json', '.caption'];
|
||||||
|
|
||||||
async function resolveSafe(p: string) {
|
async function resolveSafe(p: string) {
|
||||||
// resolve symlinks + normalize
|
// resolve symlinks + normalize
|
||||||
return await fsp.realpath(p);
|
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) {
|
export async function POST(request: NextRequest) {
|
||||||
try {
|
try {
|
||||||
const body = (await request.json()) as PostBody;
|
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) {
|
if (!body || !body.jobName) {
|
||||||
return NextResponse.json({ error: 'jobName is required' }, { status: 400 });
|
return NextResponse.json({ error: 'jobName is required' }, { status: 400 });
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -13,6 +13,7 @@ import { apiClient } from '@/utils/api';
|
|||||||
import useSettings from '@/hooks/useSettings';
|
import useSettings from '@/hooks/useSettings';
|
||||||
import { pathJoin } from '@/utils/basic';
|
import { pathJoin } from '@/utils/basic';
|
||||||
import AutoCaptionButton from '@/components/AutoCaptionButton';
|
import AutoCaptionButton from '@/components/AutoCaptionButton';
|
||||||
|
import DatasetActionBar from '@/components/DatasetActionBar';
|
||||||
import CaptionMonitor from '@/components/CaptionMonitor';
|
import CaptionMonitor from '@/components/CaptionMonitor';
|
||||||
import { CreatableSelectInput } from '@/components/formInputs';
|
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="sm:hidden">+ Add</span>
|
||||||
<span className="hidden sm:inline">Add Images</span>
|
<span className="hidden sm:inline">Add Images</span>
|
||||||
</Button>
|
</Button>
|
||||||
|
<DatasetActionBar datasetName={datasetName} />
|
||||||
</div>
|
</div>
|
||||||
</TopBar>
|
</TopBar>
|
||||||
<MainContent
|
<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