Add image-classification

public/workers/image-classification.js

@@ -0,0 +1,142 @@
+/* eslint-disable no-restricted-globals */
+import { pipeline } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.6.3'
+
+class MyImageClassificationPipeline {
+  static task = 'image-classification'
+  static instance = null
+  static modelId = null
+
+  static async getInstance(model, dtype = 'fp32', progress_callback = null) {
+    if (this.modelId !== model) {
+      // Dispose of previous pipeline if model changed
+      if (this.instance && this.instance.dispose) {
+        this.instance.dispose()
+      }
+      this.instance = null
+      this.modelId = null
+    }
+
+    if (!this.instance) {
+      try {
+        // Try WebGPU first
+        this.instance = await pipeline(this.task, model, {
+          dtype,
+          device: 'webgpu',
+          progress_callback
+        })
+      } catch (webgpuError) {
+        // Fallback to WASM if WebGPU fails
+        if (progress_callback) {
+          progress_callback({
+            status: 'fallback',
+            message: 'WebGPU failed, falling back to WASM'
+          })
+        }
+        try {
+          this.instance = await pipeline(this.task, model, {
+            dtype,
+            device: 'wasm',
+            progress_callback
+          })
+        } catch (wasmError) {
+          throw new Error(
+            `Both WebGPU and WASM failed. WebGPU error: ${webgpuError.message}. WASM error: ${wasmError.message}`
+          )
+        }
+      }
+      this.modelId = model
+    }
+
+    return this.instance
+  }
+
+  static dispose() {
+    if (this.instance && this.instance.dispose) {
+      this.instance.dispose()
+    }
+    this.instance = null
+    this.modelId = null
+  }
+}
+
+// Listen for messages from the main thread
+self.addEventListener('message', async (event) => {
+  try {
+    const { type, image, model, dtype, topK = 5 } = event.data
+
+    if (!model) {
+      self.postMessage({
+        status: 'error',
+        output: 'No model provided'
+      })
+      return
+    }
+
+    // Get the pipeline instance
+    const classifier = await MyImageClassificationPipeline.getInstance(
+      model,
+      dtype,
+      (x) => {
+        self.postMessage({ status: 'loading', output: x })
+      }
+    )
+
+    if (type === 'load') {
+      self.postMessage({
+        status: 'ready',
+        output: `Image classification model ${model}, dtype ${dtype} loaded`
+      })
+      return
+    }
+
+    if (type === 'classify') {
+      if (!image) {
+        self.postMessage({
+          status: 'error',
+          output: 'No image provided for classification'
+        })
+        return
+      }
+
+      try {
+        self.postMessage({ status: 'loading' })
+
+        // Run classification
+        const output = await classifier(image, {
+          topk: topK
+        })
+
+        // Format predictions
+        const predictions = output.map((item) => ({
+          label: item.label,
+          score: item.score
+        }))
+
+        self.postMessage({
+          status: 'output',
+          output: {
+            predictions
+          }
+        })
+      } catch (error) {
+        self.postMessage({
+          status: 'error',
+          output:
+            error.message || 'An error occurred during image classification'
+        })
+      }
+    } else if (type === 'dispose') {
+      MyImageClassificationPipeline.dispose()
+      self.postMessage({ status: 'disposed' })
+    }
+  } catch (error) {
+    self.postMessage({
+      status: 'error',
+      output:
+        error.message || 'An error occurred during pipeline initialization'
+    })
+  }
+})
+
+// Handle initialization
+self.postMessage({ status: 'ready' })

src/App.tsx

@@ -7,6 +7,7 @@ import { useModel } from './contexts/ModelContext'
 import { getModelsByPipeline } from './lib/huggingface'
 import TextGeneration from './components/pipelines/TextGeneration'
 import FeatureExtraction from './components/pipelines/FeatureExtraction'
+import ImageClassification from './components/pipelines/ImageClassification'
 import Sidebar from './components/Sidebar'
 import ModelReadme from './components/ModelReadme'
 import { PipelineLayout } from './components/PipelineLayout'
@@ -63,6 +64,7 @@ function App() {
                 {pipeline === 'text-classification' && <TextClassification />}
                 {pipeline === 'text-generation' && <TextGeneration />}
                 {pipeline === 'feature-extraction' && <FeatureExtraction />}
+                {pipeline === 'image-classification' && <ImageClassification />}
               </div>
             </div>
           </main>

src/components/PipelineLayout.tsx

@@ -2,6 +2,7 @@ import { useModel } from '../contexts/ModelContext'
 import { TextGenerationProvider } from '../contexts/TextGenerationContext'
 import { FeatureExtractionProvider } from '../contexts/FeatureExtractionContext'
 import { ZeroShotClassificationProvider } from '../contexts/ZeroShotClassificationContext'
+import { ImageClassificationProvider } from '../contexts/ImageClassificationContext'
 
 export const PipelineLayout = ({ children }: { children: React.ReactNode }) => {
   const { pipeline } = useModel()
@@ -20,6 +21,11 @@ export const PipelineLayout = ({ children }: { children: React.ReactNode }) => {
         </ZeroShotClassificationProvider>
       )
 
+    case 'image-classification':
+      return (
+        <ImageClassificationProvider>{children}</ImageClassificationProvider>
+      )
+
     default:
       return <>{children}</>
   }

src/components/Sidebar.tsx

@@ -6,6 +6,7 @@ import { useModel } from '../contexts/ModelContext'
 import TextGenerationConfig from './pipelines/TextGenerationConfig'
 import FeatureExtractionConfig from './pipelines/FeatureExtractionConfig'
 import ZeroShotClassificationConfig from './pipelines/ZeroShotClassificationConfig'
+import ImageClassificationConfig from './pipelines/ImageClassificationConfig'
 
 interface SidebarProps {
   isOpen: boolean
@@ -92,6 +93,9 @@ const Sidebar = ({ isOpen, onClose, setIsModalOpen }: SidebarProps) => {
             {pipeline === 'zero-shot-classification' && (
               <ZeroShotClassificationConfig />
             )}
+            {pipeline === 'image-classification' && (
+              <ImageClassificationConfig />
+            )}
           </div>
         </div>
       </div>

src/components/pipelines/ImageClassification.tsx

@@ -0,0 +1,489 @@
+import { useState, useRef, useCallback, useEffect } from 'react'
+import {
+  Upload,
+  Trash2,
+  Loader2,
+  X,
+  Eye,
+  EyeOff,
+  Image as ImageIcon
+} from 'lucide-react'
+import {
+  ImageClassificationWorkerInput,
+  WorkerMessage,
+  ImageExample
+} from '../../types'
+import { useModel } from '../../contexts/ModelContext'
+import { useImageClassification } from '../../contexts/ImageClassificationContext'
+
+// Sample images for quick testing (placeholder URLs)
+const SAMPLE_IMAGES = [
+  {
+    name: 'Cat',
+    url: 'https://images.unsplash.com/photo-1514888286974-6c03e2ca1dba?w=300&h=300&fit=crop'
+  },
+  {
+    name: 'Dog',
+    url: 'https://images.unsplash.com/photo-1552053831-71594a27632d?w=300&h=300&fit=crop'
+  },
+  {
+    name: 'Car',
+    url: 'https://images.unsplash.com/photo-1605559424843-9e4c228bf1c2?w=300&h=300&fit=crop'
+  },
+  {
+    name: 'Flower',
+    url: 'https://images.unsplash.com/photo-1490750967868-88aa4486c946?w=300&h=300&fit=crop'
+  }
+]
+
+function ImageClassification() {
+  const {
+    examples,
+    selectedExample,
+    setSelectedExample,
+    addExample,
+    removeExample,
+    updateExample,
+    clearExamples,
+    topK
+  } = useImageClassification()
+
+  const [isClassifying, setIsClassifying] = useState<boolean>(false)
+  const [showPreviews, setShowPreviews] = useState<boolean>(true)
+  const [dragOver, setDragOver] = useState<boolean>(false)
+  const [progress, setProgress] = useState<number | null>(null)
+
+  const {
+    activeWorker,
+    status,
+    modelInfo,
+    hasBeenLoaded,
+    selectedQuantization
+  } = useModel()
+
+  const fileInputRef = useRef<HTMLInputElement>(null)
+  const dropZoneRef = useRef<HTMLDivElement>(null)
+
+  const classifyImage = useCallback(
+    async (example: ImageExample) => {
+      if (!modelInfo || !activeWorker || !example.url) return
+
+      updateExample(example.id, { isLoading: true })
+      setIsClassifying(true)
+      setProgress(0)
+
+      const message: ImageClassificationWorkerInput = {
+        type: 'classify',
+        image: example.url,
+        model: modelInfo.id,
+        dtype: selectedQuantization ?? 'fp32',
+        topK
+      }
+
+      activeWorker.postMessage(message)
+    },
+    [modelInfo, activeWorker, selectedQuantization, topK, updateExample]
+  )
+
+  const handleFileSelect = useCallback(
+    (files: FileList | null) => {
+      if (!files) return
+
+      Array.from(files).forEach((file) => {
+        if (file.type.startsWith('image/')) {
+          addExample(file)
+        }
+      })
+    },
+    [addExample]
+  )
+
+  const handleDragOver = useCallback((e: React.DragEvent) => {
+    e.preventDefault()
+    setDragOver(true)
+  }, [])
+
+  const handleDragLeave = useCallback((e: React.DragEvent) => {
+    e.preventDefault()
+    setDragOver(false)
+  }, [])
+
+  const handleDrop = useCallback(
+    (e: React.DragEvent) => {
+      e.preventDefault()
+      setDragOver(false)
+      handleFileSelect(e.dataTransfer.files)
+    },
+    [handleFileSelect]
+  )
+
+  const handleClassifyAll = useCallback(() => {
+    const imagesToClassify = examples.filter(
+      (ex) => !ex.predictions && !ex.isLoading
+    )
+
+    imagesToClassify.forEach((example) => {
+      classifyImage(example)
+    })
+  }, [examples, classifyImage])
+
+  const handleSelectExample = useCallback(
+    (example: ImageExample) => {
+      setSelectedExample(example)
+    },
+    [setSelectedExample]
+  )
+
+  const handleLoadSampleImages = useCallback(async () => {
+    for (const sample of SAMPLE_IMAGES) {
+      try {
+        const response = await fetch(sample.url)
+        const blob = await response.blob()
+        const file = new File([blob], sample.name, { type: blob.type })
+        addExample(file)
+      } catch (error) {
+        console.error(`Failed to load sample image ${sample.name}:`, error)
+      }
+    }
+  }, [addExample])
+
+  useEffect(() => {
+    if (!activeWorker) return
+
+    const onMessageReceived = (e: MessageEvent<WorkerMessage>) => {
+      const { status, output, progress: workerProgress } = e.data
+
+      if (status === 'progress' && workerProgress !== undefined) {
+        setProgress(workerProgress)
+      } else if (status === 'output' && output?.predictions) {
+        // Find the example that was being processed
+        const processingExample = examples.find((ex) => ex.isLoading)
+        if (processingExample) {
+          updateExample(processingExample.id, {
+            predictions: output.predictions,
+            isLoading: false
+          })
+        }
+        setIsClassifying(false)
+        setProgress(null)
+      } else if (status === 'error') {
+        // Clear loading state for all examples
+        examples.forEach((ex) => {
+          if (ex.isLoading) {
+            updateExample(ex.id, { isLoading: false })
+          }
+        })
+        setIsClassifying(false)
+        setProgress(null)
+      }
+    }
+
+    activeWorker.addEventListener('message', onMessageReceived)
+    return () => activeWorker.removeEventListener('message', onMessageReceived)
+  }, [activeWorker, examples, updateExample])
+
+  const busy = status !== 'ready' || isClassifying
+
+  return (
+    <div className="flex flex-col h-full max-h-[92vh] w-full p-4">
+      <div className="flex items-center justify-between mb-4">
+        <h1 className="text-2xl font-bold">Image Classification</h1>
+        <div className="flex gap-2">
+          <button
+            onClick={handleLoadSampleImages}
+            disabled={!hasBeenLoaded || isClassifying}
+            className="px-3 py-2 bg-purple-100 hover:bg-purple-200 disabled:bg-gray-100 disabled:cursor-not-allowed rounded-lg transition-colors text-sm"
+            title="Load Sample Images"
+          >
+            Load Samples
+          </button>
+          <button
+            onClick={() => setShowPreviews(!showPreviews)}
+            className="p-2 bg-blue-100 hover:bg-blue-200 rounded-lg transition-colors"
+            title={showPreviews ? 'Hide Previews' : 'Show Previews'}
+          >
+            {showPreviews ? (
+              <EyeOff className="w-4 h-4" />
+            ) : (
+              <Eye className="w-4 h-4" />
+            )}
+          </button>
+          <button
+            onClick={clearExamples}
+            className="p-2 bg-red-100 hover:bg-red-200 rounded-lg transition-colors"
+            title="Clear All Images"
+          >
+            <Trash2 className="w-4 h-4" />
+          </button>
+        </div>
+      </div>
+
+      <div className="flex flex-col lg:flex-row gap-4 flex-1">
+        {/* Left Panel - Image Upload and List */}
+        <div className="lg:w-1/2 flex flex-col">
+          {/* Upload Area */}
+          <div className="mb-4">
+            <label className="block text-sm font-medium text-gray-700 mb-2">
+              Upload Images:
+            </label>
+            <div
+              ref={dropZoneRef}
+              onDragOver={handleDragOver}
+              onDragLeave={handleDragLeave}
+              onDrop={handleDrop}
+              className={`border-2 border-dashed rounded-lg p-6 text-center transition-colors cursor-pointer ${
+                dragOver
+                  ? 'border-blue-500 bg-blue-50'
+                  : 'border-gray-300 hover:border-gray-400'
+              } ${!hasBeenLoaded ? 'opacity-50 cursor-not-allowed' : ''}`}
+              onClick={() => hasBeenLoaded && fileInputRef.current?.click()}
+            >
+              <Upload className="w-8 h-8 mx-auto mb-2 text-gray-400" />
+              <p className="text-sm text-gray-600">
+                {dragOver
+                  ? 'Drop images here'
+                  : 'Click to upload or drag and drop images'}
+              </p>
+              <p className="text-xs text-gray-500 mt-1">
+                Supports JPG, PNG, GIF, WebP
+              </p>
+              <input
+                ref={fileInputRef}
+                type="file"
+                multiple
+                accept="image/*"
+                onChange={(e) => handleFileSelect(e.target.files)}
+                className="hidden"
+                disabled={!hasBeenLoaded}
+              />
+            </div>
+          </div>
+
+          {/* Classify Button */}
+          {examples.some((ex) => !ex.predictions) && (
+            <div className="mb-4">
+              <button
+                onClick={handleClassifyAll}
+                disabled={busy || !hasBeenLoaded}
+                className="px-6 py-2 bg-green-500 hover:bg-green-600 disabled:bg-gray-300 disabled:cursor-not-allowed text-white rounded-lg transition-colors flex items-center gap-2"
+              >
+                {isClassifying ? (
+                  <>
+                    <Loader2 className="w-4 h-4 animate-spin" />
+                    Classifying...
+                    {progress !== null && ` (${Math.round(progress * 100)}%)`}
+                  </>
+                ) : (
+                  'Classify Images'
+                )}
+              </button>
+            </div>
+          )}
+
+          {/* Images List */}
+          <div className="flex-1 overflow-y-auto border border-gray-300 rounded-lg bg-white">
+            <div className="p-4">
+              <h3 className="text-sm font-medium text-gray-700 mb-3">
+                Images ({examples.length})
+              </h3>
+              {examples.length === 0 ? (
+                <div className="text-gray-500 italic text-center py-8">
+                  No images uploaded yet. Upload some images above to get
+                  started.
+                </div>
+              ) : (
+                <div className="space-y-3">
+                  {examples.map((example) => (
+                    <div
+                      key={example.id}
+                      className={`p-3 border rounded-lg cursor-pointer transition-colors ${
+                        selectedExample?.id === example.id
+                          ? 'border-blue-500 bg-blue-50'
+                          : 'border-gray-200 hover:border-gray-300'
+                      }`}
+                      onClick={() => handleSelectExample(example)}
+                    >
+                      <div className="flex gap-3">
+                        {showPreviews && (
+                          <div className="flex-shrink-0">
+                            <img
+                              src={example.url}
+                              alt={example.name}
+                              className="w-16 h-16 object-cover rounded-lg"
+                            />
+                          </div>
+                        )}
+                        <div className="flex-1 min-w-0">
+                          <div className="flex justify-between items-start">
+                            <div className="flex-1 min-w-0">
+                              <div className="text-sm font-medium text-gray-800 truncate">
+                                {example.name}
+                              </div>
+                              <div className="flex items-center gap-2 mt-1">
+                                {example.isLoading ? (
+                                  <div className="flex items-center gap-1 text-xs text-blue-600">
+                                    <Loader2 className="w-3 h-3 animate-spin" />
+                                    Classifying...
+                                  </div>
+                                ) : example.predictions ? (
+                                  <div className="text-xs text-green-600">
+                                    ✓ Classified
+                                  </div>
+                                ) : (
+                                  <div className="text-xs text-gray-500">
+                                    Not classified
+                                  </div>
+                                )}
+                                {selectedExample?.id === example.id && (
+                                  <div className="text-xs text-blue-600">
+                                    Selected
+                                  </div>
+                                )}
+                              </div>
+                            </div>
+                            <button
+                              onClick={(e) => {
+                                e.stopPropagation()
+                                removeExample(example.id)
+                              }}
+                              className="ml-2 p-1 text-red-500 hover:text-red-700 transition-colors"
+                            >
+                              <X className="w-3 h-3" />
+                            </button>
+                          </div>
+                        </div>
+                      </div>
+                    </div>
+                  ))}
+                </div>
+              )}
+            </div>
+          </div>
+        </div>
+
+        {/* Right Panel - Preview and Results */}
+        <div className="lg:w-1/2 flex flex-col">
+          {/* Image Preview */}
+          {selectedExample && (
+            <div className="mb-4">
+              <h3 className="text-sm font-medium text-gray-700 mb-2">
+                Selected Image
+              </h3>
+              <div className="border border-gray-300 rounded-lg bg-white p-4">
+                <div className="flex flex-col items-center">
+                  <img
+                    src={selectedExample.url}
+                    alt={selectedExample.name}
+                    className="max-w-full max-h-64 object-contain rounded-lg mb-2"
+                  />
+                  <div className="text-sm text-gray-600 text-center">
+                    {selectedExample.name}
+                  </div>
+                </div>
+              </div>
+            </div>
+          )}
+
+          {/* Classification Results */}
+          <div className="flex-1 overflow-y-auto border border-gray-300 rounded-lg bg-white">
+            <div className="p-4">
+              <h3 className="text-sm font-medium text-gray-700 mb-3">
+                Classification Results
+                {selectedExample && ` - ${selectedExample.name}`}
+              </h3>
+              {!selectedExample ? (
+                <div className="text-gray-500 italic text-center py-8">
+                  <ImageIcon className="w-12 h-12 mx-auto mb-2 text-gray-300" />
+                  Select an image to see classification results
+                </div>
+              ) : selectedExample.isLoading ? (
+                <div className="text-center py-8">
+                  <Loader2 className="w-8 h-8 animate-spin mx-auto mb-2 text-blue-500" />
+                  <div className="text-sm text-gray-600">
+                    Classifying image...
+                  </div>
+                  {progress !== null && (
+                    <div className="text-xs text-gray-500 mt-1">
+                      {Math.round(progress * 100)}% complete
+                    </div>
+                  )}
+                </div>
+              ) : !selectedExample.predictions ? (
+                <div className="text-gray-500 italic text-center py-8">
+                  <button
+                    onClick={() => classifyImage(selectedExample)}
+                    disabled={busy || !hasBeenLoaded}
+                    className="px-4 py-2 bg-blue-500 hover:bg-blue-600 disabled:bg-gray-300 disabled:cursor-not-allowed text-white rounded-lg transition-colors"
+                  >
+                    Classify This Image
+                  </button>
+                </div>
+              ) : (
+                <div className="space-y-3">
+                  {selectedExample.predictions.map((prediction, index) => {
+                    const confidencePercent = (prediction.score * 100).toFixed(
+                      1
+                    )
+                    const isTopPrediction = index === 0
+
+                    return (
+                      <div
+                        key={index}
+                        className={`p-3 border rounded-lg ${
+                          isTopPrediction
+                            ? 'border-green-300 bg-green-50'
+                            : 'border-gray-200'
+                        }`}
+                      >
+                        <div className="flex justify-between items-center mb-2">
+                          <div className="flex items-center gap-2">
+                            <span className="text-sm font-medium text-gray-800">
+                              {prediction.label}
+                            </span>
+                          </div>
+                          <span className="text-sm font-medium text-gray-600">
+                            {confidencePercent}%
+                          </span>
+                        </div>
+                        <div className="w-full bg-gray-200 rounded-full h-2">
+                          <div
+                            className={`h-2 rounded-full transition-all duration-300 ${
+                              isTopPrediction
+                                ? 'bg-green-500'
+                                : prediction.score > 0.5
+                                  ? 'bg-blue-500'
+                                  : 'bg-gray-400'
+                            }`}
+                            style={{
+                              width: `${Math.max(prediction.score * 100, 2)}%`
+                            }}
+                          />
+                        </div>
+                      </div>
+                    )
+                  })}
+                </div>
+              )}
+            </div>
+          </div>
+        </div>
+      </div>
+
+      {!hasBeenLoaded && (
+        <div className="text-center text-gray-500 text-sm mt-2">
+          Please load an image classification model first to start classifying
+          images
+        </div>
+      )}
+
+      {hasBeenLoaded && examples.length === 0 && (
+        <div className="text-center text-blue-600 text-sm mt-2">
+          💡 Tip: Click "Load Samples" to try with example images, or upload
+          your own images above
+        </div>
+      )}
+    </div>
+  )
+}
+
+export default ImageClassification

src/components/pipelines/ImageClassificationConfig.tsx

@@ -0,0 +1,50 @@
+import React from 'react'
+import { useImageClassification } from '../../contexts/ImageClassificationContext'
+
+const ImageClassificationConfig = () => {
+  const { topK, setTopK } = useImageClassification()
+
+  return (
+    <div className="space-y-4">
+      <h3 className="text-lg font-semibold text-gray-900">
+        Image Classification Settings
+      </h3>
+
+      <div className="space-y-3">
+        <div>
+          <label className="block text-sm font-medium text-gray-700 mb-1">
+            Top K Predictions: {topK}
+          </label>
+          <input
+            type="range"
+            min="1"
+            max="10"
+            step="1"
+            value={topK}
+            onChange={(e) => setTopK(parseInt(e.target.value))}
+            className="w-full h-2 bg-gray-200 rounded-lg appearance-none cursor-pointer"
+          />
+          <div className="flex justify-between text-xs text-gray-400 mt-1">
+            <span>1</span>
+            <span>5</span>
+            <span>10</span>
+          </div>
+          <p className="text-xs text-gray-500 mt-1">
+            Number of top predictions to return for each image
+          </p>
+        </div>
+
+        <div className="p-3 bg-yellow-50 border border-yellow-200 rounded-lg">
+          <h4 className="text-sm font-medium text-yellow-800 mb-2">💡 Tips</h4>
+          <div className="text-xs text-yellow-700 space-y-1">
+            <p>• Use Top K = 3-5 for most cases</p>
+            <p>• Smaller images process faster</p>
+            <p>• Try quantized models for speed</p>
+          </div>
+        </div>
+      </div>
+    </div>
+  )
+}
+
+export default ImageClassificationConfig

src/contexts/ImageClassificationContext.tsx

@@ -0,0 +1,117 @@
+import React, { createContext, useContext, useState, useCallback } from 'react'
+import { ImageExample } from '../types'
+
+interface ImageClassificationContextType {
+  examples: ImageExample[]
+  selectedExample: ImageExample | null
+  setSelectedExample: (example: ImageExample | null) => void
+  addExample: (file: File) => void
+  removeExample: (id: string) => void
+  updateExample: (id: string, updates: Partial<ImageExample>) => void
+  clearExamples: () => void
+  topK: number
+  setTopK: (k: number) => void
+}
+
+const ImageClassificationContext = createContext<
+  ImageClassificationContextType | undefined
+>(undefined)
+
+export function useImageClassification() {
+  const context = useContext(ImageClassificationContext)
+  if (context === undefined) {
+    throw new Error(
+      'useImageClassification must be used within an ImageClassificationProvider'
+    )
+  }
+  return context
+}
+
+interface ImageClassificationProviderProps {
+  children: React.ReactNode
+}
+
+export function ImageClassificationProvider({
+  children
+}: ImageClassificationProviderProps) {
+  const [examples, setExamples] = useState<ImageExample[]>([])
+  const [selectedExample, setSelectedExample] = useState<ImageExample | null>(
+    null
+  )
+  const [topK, setTopK] = useState<number>(5)
+
+  const addExample = useCallback((file: File) => {
+    const id = Math.random().toString(36).substr(2, 9)
+    const url = URL.createObjectURL(file)
+
+    const newExample: ImageExample = {
+      id,
+      name: file.name,
+      url,
+      file,
+      predictions: undefined,
+      isLoading: false
+    }
+
+    setExamples((prev) => [...prev, newExample])
+  }, [])
+
+  const removeExample = useCallback((id: string) => {
+    setExamples((prev) => {
+      const updated = prev.filter((ex) => ex.id !== id)
+      // Clean up object URL to prevent memory leaks
+      const example = prev.find((ex) => ex.id === id)
+      if (example?.url) {
+        URL.revokeObjectURL(example.url)
+      }
+      return updated
+    })
+
+    // Clear selection if the selected example was removed
+    setSelectedExample((prev) => (prev?.id === id ? null : prev))
+  }, [])
+
+  const updateExample = useCallback(
+    (id: string, updates: Partial<ImageExample>) => {
+      setExamples((prev) =>
+        prev.map((ex) => (ex.id === id ? { ...ex, ...updates } : ex))
+      )
+
+      // Update selected example if it's the one being updated
+      setSelectedExample((prev) =>
+        prev?.id === id ? { ...prev, ...updates } : prev
+      )
+    },
+    []
+  )
+
+  const clearExamples = useCallback(() => {
+    // Clean up all object URLs to prevent memory leaks
+    examples.forEach((example) => {
+      if (example.url) {
+        URL.revokeObjectURL(example.url)
+      }
+    })
+
+    setExamples([])
+    setSelectedExample(null)
+  }, [examples])
+
+  const value: ImageClassificationContextType = {
+    examples,
+    selectedExample,
+    setSelectedExample,
+    addExample,
+    removeExample,
+    updateExample,
+    clearExamples,
+    topK,
+    setTopK
+  }
+
+  return (
+    <ImageClassificationContext.Provider value={value}>
+      {children}
+    </ImageClassificationContext.Provider>
+  )
+}

src/lib/huggingface.ts

@@ -34,9 +34,9 @@ const getModelInfo = async (
   const modelData: ModelInfoResponse = await response.json()
 
   const requiredFiles = [
-    'config.json',
-    'tokenizer.json',
-    'tokenizer_config.json'
+    'config.json'
+    // 'tokenizer.json',
+    // 'tokenizer_config.json'
   ]
 
   const siblingFiles = modelData.siblings?.map((s) => s.rfilename) || []

src/lib/workerManager.ts

@@ -17,6 +17,9 @@ export const getWorker = (pipeline: string) => {
       case 'feature-extraction':
         workerUrl = `${process.env.PUBLIC_URL}/workers/feature-extraction.js`
         break
+      case 'image-classification':
+        workerUrl = `${process.env.PUBLIC_URL}/workers/image-classification.js`
+        break
       default:
         return null
     }

src/types.ts

@@ -75,6 +75,28 @@ export interface FeatureExtractionWorkerInput {
   }
 }
 
+export interface ImageClassificationWorkerInput {
+  type: 'classify'
+  image: string | ImageData | HTMLImageElement | HTMLCanvasElement
+  model: string
+  dtype: QuantizationType
+  topK?: number
+}
+
+export interface ImageClassificationResult {
+  label: string
+  score: number
+}
+
+export interface ImageExample {
+  id: string
+  name: string
+  url: string
+  file?: File
+  predictions?: ImageClassificationResult[]
+  isLoading?: boolean
+}
+
 export interface EmbeddingExample {
   id: string
   text: string