100% Private Skin Screening: Building an Edge AI Vision App with WebGPU and Transformers.js

What if you could screen for skin health issues without ever uploading a single photo to a corporate server? In the era of massive data breaches and privacy concerns, "sending data to the cloud" is becoming a liability, especially for sensitive medical imagery.

Today, we are diving deep into the world of Edge AI and Privacy-First Machine Learning. We will build a skin lesion screening application that runs entirely in the browser using WebGPU acceleration, Transformers.js, and WebLLM. By leveraging on-device computation, we ensure that user data stays strictly within the browser sandbox.

Keywords: Edge AI, WebGPU Acceleration, Privacy-Preserving AI, Transformers.js Tutorial, On-device Machine Learning.

---

The Architecture: Zero-Server Data Flow

Traditional AI apps send images to a Python backend. Our approach flips the script. We download the model weights once and execute the inference locally using the user's GPU.

graph TD
    A[User Uploads Image] --> B{Browser Environment}
    B --> C[WebGPU Tensors]
    C --> D[Transformers.js Vision Model]
    D --> E[Skin Lesion Classification]
    E --> F[WebLLM Assistant]
    F --> G[Local Privacy-First Report]
    B -.->|No Data Transmitted| H[External Internet]
    style H fill:#f96,stroke:#333,stroke-dasharray: 5 5

---

Why this Stack? 🥑

1. Transformers.js: Allows us to run state-of-the-art vision models (like ViT or Swin Transformer) directly in JavaScript.
2. WebGPU: The successor to WebGL, providing near-native performance for neural network operations.
3. WebLLM: Brings large language models to the browser, allowing us to generate natural language summaries of the classification results.
4. React: Our UI layer for a responsive, modern experience.

---

Prerequisites

Before we start, ensure your browser (Chrome 113+ or Edge) supports WebGPU.

• Node.js installed
• A basic understanding of React hooks
• Curiosity for high-performance web tech!

---

Step 1: Setting up the Vision Pipeline

First, we need to initialize our image classification model. We'll use a pre-trained Vision Transformer (ViT) fine-tuned on medical datasets.

import { pipeline, env } from '@xenova/transformers';

// Enable WebGPU if available
env.allowLocalModels = false;
env.useBrowserCache = true;

const useSkinClassifier = () => {
  const [classifier, setClassifier] = useState(null);

  useEffect(() => {
    const initModel = async () => {
      // Initialize the pipeline with WebGPU execution provider
      const pipe = await pipeline('image-classification', 'Xenova/vit-base-patch16-224', {
        device: 'webgpu', 
      });
      setClassifier(() => pipe);
    };
    initModel();
  }, []);

  return classifier;
};

Step 2: Processing Images Locally

When a user selects a file, we convert it into a format Transformers.js understands without any multipart/form-data uploads.

const handleUpload = async (event: React.ChangeEvent<HTMLInputElement>) => {
  const file = event.target.files?.[0];
  if (!file || !classifier) return;

  const url = URL.createObjectURL(file);

  // Running inference 100% locally!
  const output = await classifier(url);

  console.log("Classification Results:", output);
  // Example output: [{ label: 'Melanocytic nevi', score: 0.98 }]
};

Step 3: Adding the AI Assistant with WebLLM

To make the screening "human-readable," we use WebLLM to explain the results. This allows the app to provide context while keeping the "AI logic" on the edge.

import { CreateMLCEngine } from "@mlc-ai/web-llm";

async function explainResults(label: string) {
  const engine = await CreateMLCEngine("Llama-3-8B-Instruct-q4f16_1-MLC");
  const response = await engine.chat.completions.create({
    messages: [
      { role: "system", content: "You are a helpful medical assistant. Explain what this skin condition label means in simple terms." },
      { role: "user", content: `Explain the label: ${label}` }
    ]
  });
  return response.choices[0].message.content;
}

---

The "Official" Way: Advanced Production Patterns 🚀

While building local-first apps is exciting, deploying medical-grade AI requires rigorous version control, model quantization, and robust fallback mechanisms.

For more production-ready examples and advanced patterns on optimizing WebGPU shaders for mobile browsers, check out the detailed guides at WellAlly Tech Blog. It’s my go-to resource for scaling Edge AI applications beyond simple prototypes.

---

Challenges & Optimization ⚡️

1. Model Size

Standard Vision Transformers can be ~300MB. For a production app, use Quantization (Int8 or O4) to reduce the model size to ~80MB without significant accuracy loss.

2. Warm-up Time

The first time the model runs, WebGPU compiles the shaders.
Tip: Run a "dummy inference" with a blank 1x1 pixel image as soon as the app loads to prevent UI lag during actual usage.

---

Conclusion

We’ve just built a foundation for a 100% private, browser-based medical screening tool. By combining Transformers.js and WebGPU, we respect user privacy while providing high-performance AI capabilities.

The future of AI isn't just in the cloud—it's right there in your browser's console. 🛠️

What’s next?

• [ ] Try swapping the ViT model for a specialized MobileNetV3 for even faster speeds.
• [ ] Implement PWA features so the app works offline.

If you enjoyed this tutorial, drop a comment below and let me know what Edge AI project you're working on! Happy coding! 🚀

---

Disclaimer: This tool is for educational purposes and is not a substitute for professional medical advice. Always consult a dermatologist.