From 742a4c8cef6db2d3096e865b44f1e9d403882e89 Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Wed, 12 Aug 2026 21:37:35 -0600
Subject: [PATCH] Added a Download Full Dataset and Download Captions option to
 the datasets page so they captions or dataset can be easily downloaded.

---
 ui/src/app/api/zip/route.ts                | 85 +++++++++++++++++++++-
 ui/src/app/datasets/[datasetName]/page.tsx |  2 +
 ui/src/components/DatasetActionBar.tsx     | 76 +++++++++++++++++++
 3 files changed, 160 insertions(+), 3 deletions(-)
 create mode 100644 ui/src/components/DatasetActionBar.tsx

diff --git a/ui/src/app/api/zip/route.ts b/ui/src/app/api/zip/route.ts
index 04941ea..ad17f52 100644
--- a/ui/src/app/api/zip/route.ts
+++ b/ui/src/app/api/zip/route.ts
@@ -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 });
     }
diff --git a/ui/src/app/datasets/[datasetName]/page.tsx b/ui/src/app/datasets/[datasetName]/page.tsx
index bc6cb51..ead3011 100644
--- a/ui/src/app/datasets/[datasetName]/page.tsx
+++ b/ui/src/app/datasets/[datasetName]/page.tsx
@@ -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
diff --git a/ui/src/components/DatasetActionBar.tsx b/ui/src/components/DatasetActionBar.tsx
new file mode 100644
index 0000000..cfba9d4
--- /dev/null
+++ b/ui/src/components/DatasetActionBar.tsx
@@ -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>
+  );
+}