基于TensorFlow.js与Web Worker的智能证件照生成方案

功能简介

本文基于TensorFlow.js与Web Worker实现了常用的"证件照"功能,可以对照片实现抠图并替换背景。值得一提的是,正常抠图的操作应该由后端进行,这里只是主要演示该功能实现步骤,并不建议该功能由前端全权处理。

限于个人技术能力有限,当前功能实现的并不怎么良好,抠图不够精细且缺少对图片处理后的画质修复等操作,这些还请诸位大佬见谅了。

效果演示

原图

## 效果图

功能亮点

  • 智能人像分割:基于BodyPix模型实现精准人像抠图
  • 实时背景替换:支持动态颜色选择和边缘优化算法
  • 多尺寸适配:预设常用证件照尺寸+自定义毫米级精度
  • 高性能处理:Web Worker独立线程保障主线程流畅性
  • 模型加载优化:按需加载MobileNetV1量化模型
  • 边缘平滑算法:卷积核模糊处理提升证件照专业度

主要逻辑详解

加载模型

html 复制代码
// 初始化模型
async function loadModel() {
  try {
    console.log('开始加载模型...');
    if (!bodyPixModel) {
      bodyPixModel = await bodyPix.load({
        architecture: 'MobileNetV1',
        outputStride: 16,
        multiplier: 0.75,
        quantBytes: 2
      });
      console.log('模型加载成功');
    }
    return bodyPixModel;
  } catch (error) {
    console.error('模型加载失败:', error);
    throw error;
  }
}

该方法的主要作用便是加载bodyPixModel模型,这里用到的模型为MobileNetV1。关于bodyPix模型下load方法的各种参数可以自行搜索一下官方文档查看。

人像分割核心流程

html 复制代码
    // 使用 BodyPix 进行人像分割
    console.log('开始人像分割...');
    const segmentation = await model.segmentPerson(imgDataForSegmentation, {
      flipHorizontal: false,
      internalResolution: 'medium',
      segmentationThreshold: config.segmentationThreshold || 0.7
    });
    console.log('人像分割完成');
    
    // 创建输出 Canvas
    const outputCanvas = new OffscreenCanvas(img.width, img.height);
    const outputCtx = outputCanvas.getContext('2d');
    
    // 绘制原始图片
    outputCtx.drawImage(img, 0, 0);
    
    // 应用背景色
    const backgroundColor = hexToRgb(config.bgColor);
    const outputImageData = outputCtx.getImageData(0, 0, img.width, img.height);
    const pixelData = outputImageData.data;
    
    // 应用分割结果
    for (let i = 0; i < segmentation.data.length; i++) {
      const n = i * 4;
      if (segmentation.data[i] === 0) { // 背景部分
        pixelData[n] = backgroundColor.r;
        pixelData[n + 1] = backgroundColor.g;
        pixelData[n + 2] = backgroundColor.b;
        pixelData[n + 3] = 255;
      }
    }
    
    outputCtx.putImageData(outputImageData, 0, 0);

这里的主要逻辑是分割图像及应用选择的背景色。

边缘平滑算法

html 复制代码
// 边缘平滑处理
async function smoothEdges(ctx, width, height) {
  const imageData = ctx.getImageData(0, 0, width, height);
  const data = imageData.data;
  
  const kernel = [
    [0.075, 0.124, 0.075],
    [0.124, 0.204, 0.124],
    [0.075, 0.124, 0.075]
  ];
  
  const tempData = new Uint8ClampedArray(data);
  
  for (let y = 1; y < height - 1; y++) {
    for (let x = 1; x < width - 1; x++) {
      const idx = (y * width + x) * 4;
      
      if (isEdgePixel(data, idx, width)) {
        let r = 0, g = 0, b = 0, a = 0;
        
        for (let ky = -1; ky <= 1; ky++) {
          for (let kx = -1; kx <= 1; kx++) {
            const offset = ((y + ky) * width + (x + kx)) * 4;
            const weight = kernel[ky + 1][kx + 1];
            
            r += tempData[offset] * weight;
            g += tempData[offset + 1] * weight;
            b += tempData[offset + 2] * weight;
            a += tempData[offset + 3] * weight;
          }
        }
        
        data[idx] = r;
        data[idx + 1] = g;
        data[idx + 2] = b;
        data[idx + 3] = a;
      }
    }
  }
  
  ctx.putImageData(imageData, 0, 0);
}

这里使用3*3卷积核对边缘像素进行了加权平均处理,从而消除锯齿效果。这里用到的isEdgePixel()方法目的是判断一个像素是否为图像的边缘像素,方法是通过比较该像素与其相邻像素的alpha通道值(也就是透明度)。

html 复制代码
function isEdgePixel(data, idx, width) {
  const alpha = data[idx + 3];
  const leftAlpha = data[idx - 4 + 3];
  const rightAlpha = data[idx + 4 + 3];
  const topAlpha = data[idx - width * 4 + 3];
  const bottomAlpha = data[idx + width * 4 + 3];
  
  return (alpha !== leftAlpha || 
          alpha !== rightAlpha || 
          alpha !== topAlpha || 
          alpha !== bottomAlpha);
}

完整代码

IDPhoto.vue

html 复制代码
<template>
  <div class="id-photo-container">
    <!-- 左侧工具栏 -->
    <div class="tools-panel">
      <el-form :model="photoConfig" label-position="top">
        <!-- 预设尺寸选择 -->
        <el-form-item label="证件照尺寸">
          <el-select v-model="photoConfig.selectedSize" placeholder="选择尺寸">
            <el-option
              v-for="size in presetSizes"
              :key="size.value"
              :label="size.label"
              :value="size.value"
            />
          </el-select>
        </el-form-item>

        <!-- 自定义尺寸输入 --> 
        <el-form-item label="自定义尺寸(mm)">
          <div class="custom-size">
            <el-input-number 
              v-model="photoConfig.customWidth" 
              :min="20" 
              :max="1000"
              placeholder="宽度"
            />
            <span class="separator">×</span>
            <el-input-number 
              v-model="photoConfig.customHeight" 
              :min="20" 
              :max="1000"
              placeholder="高度"
            />
          </div>
        </el-form-item>

        <!-- 背景颜色选择 -->
        <el-form-item label="背景颜色">
          <el-color-picker v-model="photoConfig.bgColor" />
        </el-form-item>

        <!-- 操作按钮 -->
        <div class="action-buttons">
          <el-button type="primary" @click="uploadImage">
            上传图片
          </el-button>
          <el-button 
            type="success" 
            :disabled="!hasImage"
            @click="downloadPhoto"
          >
            下载证件照
          </el-button>
        </div>
      </el-form>
    </div>

    <!-- 右侧预览区域 -->
    <div class="preview-panel">
      <div 
        class="preview-area"
        :style="{ backgroundColor: photoConfig.bgColor }"
      >
        <img 
          v-if="previewUrl"
          :src="previewUrl"
          ref="previewImage"
          @load="handleImageLoad"
        />
        <div v-else class="placeholder">
          请上传图片
        </div>
      </div>
    </div>

    <!-- 隐藏的文件输入框 -->
    <input
      type="file"
      ref="fileInput"
      accept="image/*"
      style="display: none"
      @change="handleFileChange"
    />

    <el-loading 
      v-model:visible="loading"
      text="处理中..."
      background="rgba(255, 255, 255, 0.8)"
    />
  </div>
</template>

<script setup>
import { ref, reactive, onMounted, onUnmounted, watch } from 'vue'
import { ElMessage } from 'element-plus'

// 预设尺寸选项
const presetSizes = [
  { label: '一寸照片 (25×35mm)', value: '25x35' },
  { label: '二寸照片 (35×49mm)', value: '35x49' },
  { label: '小二寸 (35×45mm)', value: '35x45' },
  { label: '大二寸 (35×53mm)', value: '35x53' }
]

// 照片配置
const photoConfig = reactive({
  selectedSize: '35x45',
  customWidth: 35,
  customHeight: 45,
  bgColor: '#FFFFFF',
  modelQuality: 'medium', // 可选: 'low', 'medium', 'high'
  segmentationThreshold: 0.7, // 分割阈值,可调整精度
  edgeBlur: 3 // 边缘模糊半径
})

// 组件引用
const fileInput = ref(null)
const previewImage = ref(null)

// 状态变量
const previewUrl = ref('')
const hasImage = ref(false)

// 图片处理 Worker
let imageWorker = null

// 添加加载状态
const loading = ref(false)

// 初始化 Worker
onMounted(() => {
  try {
    // 使用 ?worker 查询参数来告诉 Vite 这是一个 worker 文件
    imageWorker = new Worker(
      new URL('../../workers/idphoto.worker.js?worker', import.meta.url),
      { type: 'module' }
    )
    
    imageWorker.onmessage = (e) => {
      console.log('收到Worker响应:', e.data)
      if (e.data.status === 'success') {
        const blobUrl = URL.createObjectURL(e.data.result)
        previewUrl.value = blobUrl
        loading.value = false // 确保加载状态被重置
      } else {
        console.error('Worker处理失败:', e.data.error)
        ElMessage.error(`处理图片时出错: ${e.data.error}`)
        loading.value = false
      }
    }

    imageWorker.onerror = (error) => {
      console.error('Worker错误:', error)
      ElMessage.error('图片处理服务初始化失败')
      loading.value = false
    }
  } catch (error) {
    console.error('创建Worker失败:', error)
    ElMessage.error('初始化图片处理服务失败')
    loading.value = false
  }
})

// 清理 Worker
onUnmounted(() => {
  if (imageWorker) {
    imageWorker.terminate()
  }
})

// 上传图片
const uploadImage = () => {
  fileInput.value.click()
}

// 处理文件选择
const handleFileChange = async (event) => {
  const file = event.target.files[0]
  if (!file) return

  if (!file.type.startsWith('image/')) {
    ElMessage.error('请上传图片文件')
    return
  }

  if (file.size > 10 * 1024 * 1024) { // 10MB 限制
    ElMessage.error('图片大小不能超过10MB')
    return
  }

  loading.value = true
  ElMessage.info('正在处理图片,首次使用可能需要加载模型...')

  const reader = new FileReader()
  reader.onload = (e) => {
    const img = new Image()
    img.onload = () => {
      console.log('图片加载成功,尺寸:', img.width, 'x', img.height)
      try {
        imageWorker.postMessage({
          imageData: e.target.result,
          config: {
            width: img.width,
            height: img.height,
            bgColor: photoConfig.bgColor,
            segmentationThreshold: photoConfig.segmentationThreshold,
            modelQuality: photoConfig.modelQuality
          }
        })
      } catch (error) {
        loading.value = false
        console.error('发送数据到Worker时出错:', error)
        ElMessage.error('处理图片时出错')
      }
    }
    img.onerror = (error) => {
      loading.value = false
      console.error('图片加载失败:', error)
      ElMessage.error('图片加载失败')
    }
    img.src = e.target.result
    hasImage.value = true
  }
  reader.onerror = (error) => {
    loading.value = false
    console.error('读取文件失败:', error)
    ElMessage.error('读取文件失败')
  }
  reader.readAsDataURL(file)
}

// 处理图片加载
const handleImageLoad = () => {
  // 这里可以添加图片加载后的处理逻辑
}

// 下载证件照
const downloadPhoto = async () => {
  if (!hasImage.value) return

  try {
    const response = await fetch(previewUrl.value)
    const blob = await response.blob()
    
    const link = document.createElement('a')
    link.download = '证件照.png'
    link.href = URL.createObjectURL(blob)
    link.click()
    
    ElMessage.success('下载成功')
  } catch (error) {
    ElMessage.error('下载图片时出错')
    console.error(error)
  }
}

// 添加背景色变化监听
watch(() => photoConfig.bgColor, (newColor) => {
  if (hasImage.value && previewUrl.value) {
    const img = new Image()
    img.onload = () => {
      imageWorker.postMessage({
        imageData: previewUrl.value,
        config: {
          width: img.width,
          height: img.height,
          bgColor: newColor
        }
      })
    }
    img.src = previewUrl.value
  }
})
</script>

<style scoped>
.id-photo-container {
  display: flex;
  gap: 20px;
  padding: 20px;
  height: 100%;
}

.tools-panel {
  width: 300px;
  padding: 20px;
  background: #f5f7fa;
  border-radius: 8px;
}

.preview-panel {
  flex: 1;
  display: flex;
  justify-content: center;
  align-items: center;
  background: #f5f7fa;
  border-radius: 8px;
  overflow: hidden;
}

.preview-area {
  width: 80%;
  height: 80%;
  display: flex;
  justify-content: center;
  align-items: center;
  background: #fff;
  border-radius: 4px;
  box-shadow: 0 2px 12px rgba(0, 0, 0, 0.1);
}

.preview-area img {
  max-width: 100%;
  max-height: 100%;
  object-fit: contain;
}

.placeholder {
  color: #909399;
  font-size: 16px;
}

.custom-size {
  display: flex;
  align-items: center;
  gap: 10px;
}

.separator {
  color: #909399;
}

.action-buttons {
  display: flex;
  gap: 10px;
  margin-top: 20px;
}
</style>

idphoto.workder.js

html 复制代码
import * as tf from '@tensorflow/tfjs'
import * as bodyPix from '@tensorflow-models/body-pix'

let bodyPixModel = null;

// 初始化模型
async function loadModel() {
  try {
    console.log('开始加载模型...');
    if (!bodyPixModel) {
      bodyPixModel = await bodyPix.load({
        architecture: 'MobileNetV1',
        outputStride: 16,
        multiplier: 0.75,
        quantBytes: 2
      });
      console.log('模型加载成功');
    }
    return bodyPixModel;
  } catch (error) {
    console.error('模型加载失败:', error);
    throw error;
  }
}

// 处理图片的 Worker
self.onmessage = async function(e) {
  console.log('Worker 收到消息:', e.data);
  const { imageData, config } = e.data;

  try {
    if (!imageData || !config) {
      throw new Error('缺少必要的参数')
    }

    // 加载模型
    const model = await loadModel();
    console.log('模型准备就绪');
    
    // 创建图片元素
    const img = await createImageBitmap(
      await fetch(imageData).then(r => r.blob())
    );
    
    // 创建离屏 Canvas
    const canvas = new OffscreenCanvas(img.width, img.height);
    const ctx = canvas.getContext('2d');
    
    if (!ctx) {
      throw new Error('无法创建Canvas上下文')
    }
    
    // 绘制原始图片
    ctx.drawImage(img, 0, 0);
    console.log('图片绘制完成');
    
    // 获取图片数据
    const imgDataForSegmentation = ctx.getImageData(0, 0, img.width, img.height);
    
    // 使用 BodyPix 进行人像分割
    console.log('开始人像分割...');
    const segmentation = await model.segmentPerson(imgDataForSegmentation, {
      flipHorizontal: false,
      internalResolution: 'medium',
      segmentationThreshold: config.segmentationThreshold || 0.7
    });
    console.log('人像分割完成');
    
    // 创建输出 Canvas
    const outputCanvas = new OffscreenCanvas(img.width, img.height);
    const outputCtx = outputCanvas.getContext('2d');
    
    // 绘制原始图片
    outputCtx.drawImage(img, 0, 0);
    
    // 应用背景色
    const backgroundColor = hexToRgb(config.bgColor);
    const outputImageData = outputCtx.getImageData(0, 0, img.width, img.height);
    const pixelData = outputImageData.data;
    
    // 应用分割结果
    for (let i = 0; i < segmentation.data.length; i++) {
      const n = i * 4;
      if (segmentation.data[i] === 0) { // 背景部分
        pixelData[n] = backgroundColor.r;
        pixelData[n + 1] = backgroundColor.g;
        pixelData[n + 2] = backgroundColor.b;
        pixelData[n + 3] = 255;
      }
    }
    
    outputCtx.putImageData(outputImageData, 0, 0);
    
    // 优化边缘
    await smoothEdges(outputCtx, img.width, img.height);
    
    // 转换为 Blob
    const resultBlob = await outputCanvas.convertToBlob({
      type: 'image/png'
    });
    
    console.log('处理完成,发送结果');
    self.postMessage({
      status: 'success',
      result: resultBlob
    });
    
  } catch (error) {
    console.error('Worker处理错误:', error);
    self.postMessage({
      status: 'error',
      error: error.message || '处理图片时发生未知错误'
    });
  }
};

// 边缘平滑处理
async function smoothEdges(ctx, width, height) {
  const imageData = ctx.getImageData(0, 0, width, height);
  const data = imageData.data;
  
  const kernel = [
    [0.075, 0.124, 0.075],
    [0.124, 0.204, 0.124],
    [0.075, 0.124, 0.075]
  ];
  
  const tempData = new Uint8ClampedArray(data);
  
  for (let y = 1; y < height - 1; y++) {
    for (let x = 1; x < width - 1; x++) {
      const idx = (y * width + x) * 4;
      
      if (isEdgePixel(data, idx, width)) {
        let r = 0, g = 0, b = 0, a = 0;
        
        for (let ky = -1; ky <= 1; ky++) {
          for (let kx = -1; kx <= 1; kx++) {
            const offset = ((y + ky) * width + (x + kx)) * 4;
            const weight = kernel[ky + 1][kx + 1];
            
            r += tempData[offset] * weight;
            g += tempData[offset + 1] * weight;
            b += tempData[offset + 2] * weight;
            a += tempData[offset + 3] * weight;
          }
        }
        
        data[idx] = r;
        data[idx + 1] = g;
        data[idx + 2] = b;
        data[idx + 3] = a;
      }
    }
  }
  
  ctx.putImageData(imageData, 0, 0);
}

function isEdgePixel(data, idx, width) {
  const alpha = data[idx + 3];
  const leftAlpha = data[idx - 4 + 3];
  const rightAlpha = data[idx + 4 + 3];
  const topAlpha = data[idx - width * 4 + 3];
  const bottomAlpha = data[idx + width * 4 + 3];
  
  return (alpha !== leftAlpha || 
          alpha !== rightAlpha || 
          alpha !== topAlpha || 
          alpha !== bottomAlpha);
}

function hexToRgb(hex) {
  const result = /^#?([a-f\d]{2})([a-f\d]{2})([a-f\d]{2})$/i.exec(hex);
  return result ? {
    r: parseInt(result[1], 16),
    g: parseInt(result[2], 16),
    b: parseInt(result[3], 16)
  } : null;
} 

以上便是证件照功能的全部逻辑,其实我这里写的相当简陋且具有很大的扩展空间。各位大佬如果有精力,可以在这个基础上 增加服装替换、美颜等功能,以及可以进一步优化UI界面(我这里样式写的不是特别好)。

总之,感谢阅读了,愿你我都能在技术之路上更进一步!

相关推荐
前端大卫1 小时前
Vue3 + Element-Plus 自定义虚拟表格滚动实现方案【附源码】
前端
却尘1 小时前
Next.js 请求最佳实践 - vercel 2026一月发布指南
前端·react.js·next.js
ccnocare1 小时前
浅浅看一下设计模式
前端
Lee川1 小时前
🎬 从标签到屏幕:揭秘现代网页构建与适配之道
前端·面试
Ticnix2 小时前
ECharts初始化、销毁、resize 适配组件封装(含完整封装代码)
前端·echarts
纯爱掌门人2 小时前
终焉轮回里,藏着 AI 与人类的答案
前端·人工智能·aigc
twl2 小时前
OpenClaw 深度技术解析
前端
崔庆才丨静觅2 小时前
比官方便宜一半以上!Grok API 申请及使用
前端
星光不问赶路人2 小时前
vue3使用jsx语法详解
前端·vue.js
天蓝色的鱼鱼2 小时前
shadcn/ui,给你一个真正可控的UI组件库
前端