More ui more ui

ui/src/app/api/caption/[...imagePath]/route.ts

@@ -0,0 +1,42 @@
+// src/app/api/img/[imagePath]/route.ts
+import { NextRequest, NextResponse } from 'next/server';
+import fs from 'fs';
+import path from 'path';
+import { getDatasetsRoot } from '@/app/api/datasets/utils';
+
+export async function GET(request: NextRequest, { params }: { params: { imagePath: string } }) {
+  const { imagePath } = await params;
+  try {
+    // Decode the path
+    const filepath = decodeURIComponent(imagePath);
+
+    // caption name is the filepath without extension but with .txt
+    const captionPath = filepath.replace(/\.[^/.]+$/, '') + '.txt';
+
+    // Get allowed directories
+    const allowedDir = await getDatasetsRoot();
+
+    // Security check: Ensure path is in allowed directory
+    const isAllowed = filepath.startsWith(allowedDir) && !filepath.includes('..');
+
+    if (!isAllowed) {
+      console.warn(`Access denied: ${filepath} not in ${allowedDir}`);
+      return new NextResponse('Access denied', { status: 403 });
+    }
+
+    // Check if file exists
+    if (!fs.existsSync(captionPath)) {
+      // send back blank string if caption file does not exist
+      return new NextResponse('');
+    }
+
+    // Read caption file
+    const caption = fs.readFileSync(captionPath, 'utf-8');
+
+    // Return caption
+    return new NextResponse(caption);
+  } catch (error) {
+    console.error('Error getting caption:', error);
+    return new NextResponse('Error getting caption', { status: 500 });
+  }
+}

ui/src/app/api/datasets/create/route.tsx

@@ -0,0 +1,22 @@
+import { NextResponse } from 'next/server';
+import fs from 'fs';
+import path from 'path';
+import { getDatasetsRoot } from '@/app/api/datasets/utils';
+
+export async function POST(request: Request) {
+  try {
+    const body = await request.json();
+    const { name } = body;
+    let datasetsPath = await getDatasetsRoot();
+    let datasetPath = path.join(datasetsPath, name);
+
+    // if folder doesnt exist, create it
+    if (!fs.existsSync(datasetPath)) {
+      fs.mkdirSync(datasetPath);
+    }
+
+    return NextResponse.json({ success: true });
+  } catch (error) {
+    return NextResponse.json({ error: 'Failed to create dataset' }, { status: 500 });
+  }
+}

ui/src/app/api/datasets/list/route.ts

@@ -0,0 +1,25 @@
+import { NextResponse } from 'next/server';
+import fs from 'fs';
+import { getDatasetsRoot } from '@/app/api/datasets/utils';
+
+export async function GET() {
+  try {
+    let datasetsPath = await getDatasetsRoot();
+
+    // if folder doesnt exist, create it
+    if (!fs.existsSync(datasetsPath)) {
+      fs.mkdirSync(datasetsPath);
+    }
+
+    // find all the folders in the datasets folder
+    let folders = fs
+      .readdirSync(datasetsPath, { withFileTypes: true })
+      .filter(dirent => dirent.isDirectory())
+      .filter(dirent => !dirent.name.startsWith('.'))
+      .map(dirent => dirent.name);
+
+    return NextResponse.json(folders);
+  } catch (error) {
+    return NextResponse.json({ error: 'Failed to fetch datasets' }, { status: 500 });
+  }
+}

ui/src/app/api/datasets/listImages/route.ts

@@ -0,0 +1,67 @@
+import { NextResponse } from 'next/server';
+import fs from 'fs';
+import path from 'path';
+import { getDatasetsRoot } from '@/app/api/datasets/utils';
+
+export async function POST(request: Request) {
+  const datasetsPath = await getDatasetsRoot();
+  const body = await request.json();
+  const { datasetName } = body;
+  const datasetFolder = path.join(datasetsPath, datasetName);
+  
+  try {
+    // Check if folder exists
+    if (!fs.existsSync(datasetFolder)) {
+      return NextResponse.json(
+        { error: `Folder '${datasetName}' not found` },
+        { status: 404 }
+      );
+    }
+
+    // Find all images recursively
+    const imageFiles = findImagesRecursively(datasetFolder);
+    
+    // Format response
+    const result = imageFiles.map(imgPath => ({
+      img_path: imgPath
+    }));
+    
+    return NextResponse.json({ images: result });
+  } catch (error) {
+    console.error('Error finding images:', error);
+    return NextResponse.json(
+      { error: 'Failed to process request' },
+      { status: 500 }
+    );
+  }
+}
+
+/**
+ * Recursively finds all image files in a directory and its subdirectories
+ * @param dir Directory to search
+ * @returns Array of absolute paths to image files
+ */
+function findImagesRecursively(dir: string): string[] {
+  const imageExtensions = ['.png', '.jpg', '.jpeg', '.webp'];
+  let results: string[] = [];
+  
+  const items = fs.readdirSync(dir);
+  
+  for (const item of items) {
+    const itemPath = path.join(dir, item);
+    const stat = fs.statSync(itemPath);
+    
+    if (stat.isDirectory()) {
+      // If it's a directory, recursively search it
+      results = results.concat(findImagesRecursively(itemPath));
+    } else {
+      // If it's a file, check if it's an image
+      const ext = path.extname(itemPath).toLowerCase();
+      if (imageExtensions.includes(ext)) {
+        results.push(itemPath);
+      }
+    }
+  }
+  
+  return results;
+}

ui/src/app/api/datasets/utils.ts

@@ -0,0 +1,17 @@
+import { PrismaClient } from '@prisma/client';
+import { defaultDatasetsFolder } from '@/paths';
+
+const prisma = new PrismaClient();
+
+export const getDatasetsRoot = async () => {
+  let row = await prisma.settings.findFirst({
+    where: {
+      key: 'DATASETS_FOLDER',
+    },
+  });
+  let datasetsPath = defaultDatasetsFolder;
+  if (row?.value && row.value !== '') {
+    datasetsPath = row.value;
+  }
+  return datasetsPath;
+};

ui/src/app/api/img/[...imagePath]/route.ts

@@ -0,0 +1,66 @@
+// src/app/api/img/[imagePath]/route.ts
+import { NextRequest, NextResponse } from 'next/server';
+import fs from 'fs';
+import path from 'path';
+import { getDatasetsRoot } from '@/app/api/datasets/utils';
+
+export async function GET(request: NextRequest, { params }: { params: { imagePath: string } }) {
+  const { imagePath } = await params;
+  try {
+    // Decode the path
+    const filepath = decodeURIComponent(imagePath);
+    console.log('Serving image:', filepath);
+
+    // Get allowed directories
+    const allowedDir = await getDatasetsRoot();
+
+    // Security check: Ensure path is in allowed directory
+    const isAllowed = filepath.startsWith(allowedDir) && !filepath.includes('..');
+
+    if (!isAllowed) {
+      console.warn(`Access denied: ${filepath} not in ${allowedDir}`);
+      return new NextResponse('Access denied', { status: 403 });
+    }
+
+    // Check if file exists
+    if (!fs.existsSync(filepath)) {
+      console.warn(`File not found: ${filepath}`);
+      return new NextResponse('File not found', { status: 404 });
+    }
+
+    // Get file info
+    const stat = fs.statSync(filepath);
+    if (!stat.isFile()) {
+      return new NextResponse('Not a file', { status: 400 });
+    }
+
+    // Determine content type
+    const ext = path.extname(filepath).toLowerCase();
+    const contentTypeMap: { [key: string]: string } = {
+      '.jpg': 'image/jpeg',
+      '.jpeg': 'image/jpeg',
+      '.png': 'image/png',
+      '.gif': 'image/gif',
+      '.webp': 'image/webp',
+      '.svg': 'image/svg+xml',
+      '.bmp': 'image/bmp',
+    };
+
+    const contentType = contentTypeMap[ext] || 'application/octet-stream';
+
+    // Read file as buffer
+    const fileBuffer = fs.readFileSync(filepath);
+
+    // Return file with appropriate headers
+    return new NextResponse(fileBuffer, {
+      headers: {
+        'Content-Type': contentType,
+        'Content-Length': String(stat.size),
+        'Cache-Control': 'public, max-age=86400',
+      },
+    });
+  } catch (error) {
+    console.error('Error serving image:', error);
+    return new NextResponse('Internal Server Error', { status: 500 });
+  }
+}

ui/src/app/api/settings/route.ts

@@ -1,6 +1,6 @@
 import { NextResponse } from 'next/server';
 import { PrismaClient } from '@prisma/client';
-import { defaultTrainFolder } from '@/paths';
+import { defaultTrainFolder, defaultDatasetsFolder } from '@/paths';
 
 const prisma = new PrismaClient();
 
@@ -15,6 +15,10 @@ export async function GET() {
     if (!settingsObject.TRAINING_FOLDER || settingsObject.TRAINING_FOLDER === '') {
       settingsObject.TRAINING_FOLDER = defaultTrainFolder;
     }
+    // if DATASETS_FOLDER is not set, use default
+    if (!settingsObject.DATASETS_FOLDER || settingsObject.DATASETS_FOLDER === '') {
+      settingsObject.DATASETS_FOLDER = defaultDatasetsFolder;
+    }
     return NextResponse.json(settingsObject);
   } catch (error) {
     return NextResponse.json({ error: 'Failed to fetch settings' }, { status: 500 });
@@ -24,7 +28,7 @@ export async function GET() {
 export async function POST(request: Request) {
   try {
     const body = await request.json();
-    const { HF_TOKEN, TRAINING_FOLDER } = body;
+    const { HF_TOKEN, TRAINING_FOLDER, DATASETS_FOLDER } = body;
 
     // Upsert both settings
     await Promise.all([
@@ -38,6 +42,11 @@ export async function POST(request: Request) {
         update: { value: TRAINING_FOLDER },
         create: { key: 'TRAINING_FOLDER', value: TRAINING_FOLDER },
       }),
+      prisma.settings.upsert({
+        where: { key: 'DATASETS_FOLDER' },
+        update: { value: DATASETS_FOLDER },
+        create: { key: 'DATASETS_FOLDER', value: DATASETS_FOLDER },
+      }),
     ]);
 
     return NextResponse.json({ success: true });

ui/src/app/dashboard/page.tsx

@@ -5,7 +5,7 @@ import GpuMonitor from '@/components/GPUMonitor';
 export default function Dashboard() {
   return (
     <div className="space-y-6">
-      <h1 className="text-3xl font-bold">Dashboard</h1>
+      <h1 className="text-xl font-bold mb-8">Dashboard</h1>
       <GpuMonitor />
     </div>
   );

ui/src/app/datasets/[datasetName]/page.tsx

@@ -0,0 +1,73 @@
+'use client';
+
+import { useEffect, useState, use } from 'react';
+import Card from '@/components/Card';
+import { Modal } from '@/components/Modal';
+import Link from 'next/link';
+import { TextInput } from '@/components/formInputs';
+import { useRouter } from 'next/router';
+import DatasetImageCard from '@/components/DatasetImageCard';
+
+export default function DatasetPage({ params }: { params: { datasetName: string } }) {
+  const [imgList, setImgList] = useState<{ img_path: string }[]>([]);
+  const usableParams = use(params as any) as { datasetName: string };
+  const datasetName = usableParams.datasetName;
+  const [newDatasetName, setNewDatasetName] = useState('');
+  const [isNewDatasetModalOpen, setIsNewDatasetModalOpen] = useState(false);
+  const [status, setStatus] = useState<'idle' | 'loading' | 'success' | 'error'>('idle');
+
+  const refreshImageList = (dbName: string) => {
+    setStatus('loading');
+    fetch('/api/datasets/listImages', {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+      },
+      body: JSON.stringify({ datasetName: dbName }),
+    })
+      .then(res => res.json())
+      .then(data => {
+        console.log('Images:', data.images);
+        // sort
+        data.images.sort((a: { img_path: string }, b: { img_path: string }) => a.img_path.localeCompare(b.img_path));
+        setImgList(data.images);
+        setStatus('success');
+      })
+      .catch(error => {
+        console.error('Error fetching images:', error);
+        setStatus('error');
+      });
+  };
+  useEffect(() => {
+    if (datasetName) {
+      refreshImageList(datasetName);
+    }
+  }, [datasetName]);
+  return (
+    <>
+      <div className="space-y-6">
+        <div className="flex justify-between items-center">
+          <div>
+            <h1 className="text-xl font-bold mb-8">Dataset: {datasetName}</h1>
+          </div>
+        </div>
+        <Card title={`Images (${imgList.length})`}>
+          {status === 'loading' && <p>Loading...</p>}
+          {status === 'error' && <p>Error fetching images</p>}
+          {status === 'success' && (
+            <div className="grid grid-cols-1 sm:grid-cols-2 md:grid-cols-3 lg:grid-cols-4 gap-4">
+              {imgList.length === 0 && <p>No images found</p>}
+              {imgList.map(img => (
+                <DatasetImageCard
+                  key={img.img_path}
+                  alt="image"
+                  imageUrl={img.img_path}
+                />
+              ))}
+            </div>
+          )}
+        </Card>
+      </div>
+    </>
+  );
+}

ui/src/app/datasets/page.tsx

@@ -0,0 +1,124 @@
+'use client';
+
+import { useEffect, useState } from 'react';
+import Card from '@/components/Card';
+import { Modal } from '@/components/Modal';
+import Link from 'next/link';
+import { TextInput } from '@/components/formInputs';
+
+export default function Datasets() {
+  const [datasets, setDatasets] = useState([]);
+  const [newDatasetName, setNewDatasetName] = useState('');
+  const [isNewDatasetModalOpen, setIsNewDatasetModalOpen] = useState(false);
+  const [status, setStatus] = useState<'idle' | 'loading' | 'success' | 'error'>('idle');
+
+  const refreshDatasets = () => {
+    setStatus('loading');
+    fetch('/api/datasets/list')
+      .then(res => res.json())
+      .then(data => {
+        console.log('Datasets:', data);
+        // sort
+        data.sort((a: string, b: string) => a.localeCompare(b));
+        setDatasets(data);
+        setStatus('success');
+      })
+      .catch(error => {
+        console.error('Error fetching datasets:', error);
+        setStatus('error');
+      });
+  };
+  useEffect(() => {
+    refreshDatasets();
+  }, []);
+  return (
+    <>
+      <div className="space-y-6">
+        <div className="flex justify-between items-center">
+          <div>
+            <h1 className="text-xl font-bold mb-8">Datasets</h1>
+          </div>
+          <div>
+            <button
+              onClick={() => {
+                setIsNewDatasetModalOpen(true);
+              }}
+              className="bg-blue-500 hover:bg-blue-700 text-white font-bold py-2 px-4 rounded"
+            >
+              New Dataset
+            </button>
+          </div>
+        </div>
+        <Card title={`Datasets (${datasets.length})`}>
+          {status === 'loading' && <p>Loading...</p>}
+          {status === 'error' && <p>Error fetching datasets</p>}
+          {status === 'success' && (
+            <div className="space-y-1">
+              {datasets.length === 0 && <p>No datasets found</p>}
+              {datasets.map((dataset: string) => (
+                <Link href={`/datasets/${dataset}`} className="bg-gray-800 hover:bg-gray-700 py-2 px-4 rounded-lg cursor-pointer block" key={dataset}>
+                  {dataset}
+                </Link>
+              ))}
+            </div>
+          )}
+        </Card>
+      </div>
+      <Modal
+        isOpen={isNewDatasetModalOpen}
+        onClose={() => setIsNewDatasetModalOpen(false)}
+        title="New Dataset"
+        size="md"
+      >
+        <div className="space-y-4 text-gray-200">
+          <form
+            onSubmit={e => {
+              e.preventDefault();
+              console.log('Creating new dataset');
+              // make post with name to create new dataset
+              fetch('/api/datasets/create', {
+                method: 'POST',
+                headers: {
+                  'Content-Type': 'application/json',
+                },
+                body: JSON.stringify({ name: newDatasetName }),
+              })
+                .then(res => res.json())
+                .then(data => {
+                  console.log('New dataset created:', data);
+                  refreshDatasets();
+                  setNewDatasetName('');
+                  setIsNewDatasetModalOpen(false);
+                })
+                .catch(error => {
+                  console.error('Error creating new dataset:', error);
+                });
+            }}
+          >
+            <div className="text-sm text-gray-600">
+              This will create a new folder with the name below in your dataset folder.
+            </div>
+            <div>
+              <TextInput label="Dataset Name" value={newDatasetName} onChange={value => setNewDatasetName(value)} />
+            </div>
+
+            <div className="mt-6 flex justify-end space-x-3">
+              <button
+                className="rounded-md bg-gray-700 px-4 py-2 text-gray-200 hover:bg-gray-600 focus:outline-none focus:ring-2 focus:ring-gray-500"
+                onClick={() => setIsNewDatasetModalOpen(false)}
+              >
+                Cancel
+              </button>
+              <button
+                className="rounded-md bg-blue-600 px-4 py-2 text-white hover:bg-blue-700 focus:outline-none focus:ring-2 focus:ring-blue-500"
+                type="submit"
+              >
+                Confirm
+              </button>
+            </div>
+          </form>
+        </div>
+      </Modal>
+    </>
+  );
+}

ui/src/app/settings/page.tsx

@@ -6,6 +6,7 @@ export default function Settings() {
   const [settings, setSettings] = useState({
     HF_TOKEN: '',
     TRAINING_FOLDER: '',
+    DATASETS_FOLDER: '',
   });
   const [status, setStatus] = useState<'idle' | 'saving' | 'success' | 'error'>('idle');
 
@@ -17,6 +18,7 @@ export default function Settings() {
         setSettings({
           HF_TOKEN: data.HF_TOKEN || '',
           TRAINING_FOLDER: data.TRAINING_FOLDER || '',
+          DATASETS_FOLDER: data.DATASETS_FOLDER || '',
         });
       })
       .catch(error => console.error('Error fetching settings:', error));
@@ -65,7 +67,8 @@ export default function Settings() {
                 <a href="https://huggingface.co/settings/tokens" target="_blank" rel="noreferrer">
                   {' '}
                   Huggingface
-                </a> if you need to access gated/private models.
+                </a>{' '}
+                if you need to access gated/private models.
               </div>
             </label>
             <input
@@ -83,7 +86,8 @@ export default function Settings() {
             <label htmlFor="TRAINING_FOLDER" className="block text-sm font-medium mb-2">
               Training Folder Path
               <div className="text-gray-500 text-sm ml-1">
-                We will store your training information here. Must be an absolute path. If blank, it will default to the output folder in the project root.
+                We will store your training information here. Must be an absolute path. If blank, it will default to the
+                output folder in the project root.
               </div>
             </label>
             <input
@@ -96,6 +100,28 @@ export default function Settings() {
               placeholder="Enter training folder path"
             />
           </div>
+
+          <div>
+            <label htmlFor="DATASETS_FOLDER" className="block text-sm font-medium mb-2">
+              Dataset Folder Path
+              <div className="text-gray-500 text-sm ml-1">
+                Where we store and find your datasets.{' '}
+                <span className="text-orange-800">
+                  Warning: This software may modify datasets so it is recommended you keep a backup somewhere else or
+                  have a dedicated folder for this software.
+                </span>
+              </div>
+            </label>
+            <input
+              type="text"
+              id="DATASETS_FOLDER"
+              name="DATASETS_FOLDER"
+              value={settings.DATASETS_FOLDER}
+              onChange={handleChange}
+              className="w-full px-4 py-2 bg-gray-800 border border-gray-700 rounded-lg focus:ring-2 focus:ring-gray-600 focus:border-transparent"
+              placeholder="Enter datasets folder path"
+            />
+          </div>
         </div>
 
         <button