基于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界面(我这里样式写的不是特别好)。

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

相关推荐
大大。2 分钟前
微信小程序 左右滑动块,自定义的switch组件,带文字状态的开关
java·前端·javascript
库库林_沙琪马13 分钟前
Spring Boot Validation 接口校验:从零到掌握
java·前端·spring boot
best_virtuoso40 分钟前
会话对象 HttpSession 一、HttpSession原理
java·前端
hamburgerDaddy11 小时前
从零开始用react + tailwindcs + express + mongodb实现一个聊天程序(二)
前端·javascript·mongodb·react.js·前端框架·express
小童不学前端1 小时前
前端面试真题 2025最新版
前端·面试
bin91532 小时前
DeepSeek 助力 Vue 开发:打造丝滑的文本输入框(Text Input)
前端·javascript·vue.js·前端框架·ecmascript·deepseek
pixle02 小时前
Three.js 快速入门教程【六】相机控件 OrbitControls
前端·javascript·3d
rkmhr_sef3 小时前
Go-Gin Web 框架完整教程
前端·golang·gin
Ama_tor3 小时前
网页制作08-html,css,javascript初认识のhtml使用框架结构,请先建立站点!
前端·css