当你执行像图像预处理、模型推理等计算密集型任务时。Web Worker 可以帮助你将这些繁重的计算任务转移到后台线程,以避免阻塞主线程,进而提升 Web 应用的响应速度和用户体验。

为什么使用 Web Worker?

  1. 异步执行计算密集型任务:模型推理和图像处理通常会占用大量的 CPU 资源,使用 Web Worker 可以将这些任务移到后台线程中执行,不会阻塞主线程的渲染和交互。
  2. 提升响应速度:通过将计算密集型任务移到 Web Worker,主线程可以保持流畅的用户交互,提高应用的响应速度。
  3. 更好的用户体验:使用 Web Worker 处理任务可以避免主线程被阻塞,带来更流畅的用户体验。

1.在src/components中新建removeBgc-useWorker目录

2.在removeBgc-useWorker下新建index.css:

.remove-bgc-container {
  min-height: 500px;
  display: flex;
  justify-content: center;
  align-items: center;
}

3.在removeBgc-useWorker下新建imageWorker.js:

/* eslint-disable no-restricted-globals */
import { AutoModel, AutoProcessor, RawImage } from '@xenova/transformers';

const initialize = async () => {
  const model = await AutoModel.from_pretrained('briaai/RMBG-1.4', {
    config: { model_type: 'custom' },
  });

  const processor = await AutoProcessor.from_pretrained('briaai/RMBG-1.4', {
    config: {
      do_normalize: true,
      do_pad: false,
      do_rescale: true,
      do_resize: true,
      image_mean: [0.5, 0.5, 0.5],
      feature_extractor_type: "ImageFeatureExtractor",
      image_std: [1, 1, 1],
      resample: 2,
      rescale_factor: 0.00392156862745098,
      size: { width: 1024, height: 1024 },
    }
  });

  self.model = model;
  self.processor = processor;
  self.postMessage({ type: 'initialized' });
};

// Listen to messages from the main thread
self.onmessage = async (event) => {
  if (event.data.type === 'init') {
    await initialize();
  } else if (event.data.type === 'analyze') {
    const { url } = event.data;
    const image = await RawImage.fromURL(url);
    const { pixel_values } = await self.processor(image);
    const { output } = await self.model({ input: pixel_values });
    const mask = await RawImage.fromTensor(output[0].mul(255).to('uint8')).resize(image.width, image.height);
    const canvas = new OffscreenCanvas(image.width, image.height);
    const ctx = canvas.getContext('2d');
    ctx.drawImage(await image.toCanvas(), 0, 0);
    const pixelData = ctx.getImageData(0, 0, image.width, image.height);
    for (let i = 0; i < mask.data.length; ++i) {
      pixelData.data[4 * i + 3] = mask.data[i];
    }
    ctx.putImageData(pixelData, 0, 0);
    const analyzedImageBlob = await canvas.convertToBlob(); // Convert canvas to Blob
    self.postMessage({ type: 'analyzed', analyzedImageBlob });
  }
};

4.在removeBgc-useWorker下新建index.jsx:

import React, { useState, useRef, useEffect } from "react";
import { Button, Image, Spin, Upload } from "antd";
import './index.css';

const RemoveBgc = () => {
  const [isUploaded, setIsUploaded] = useState(false);
  const [isAnalyzed, setIsAnalyzed] = useState(false);
  const [loading, setLoading] = useState(false);
  const [isModelReady, setIsModelReady] = useState(false);
  const [isAnalysing, setIsAnalysing] = useState(false);
  const [imageSrc, setImageSrc] = useState('');
  const [analyzedImageSrc, setAnalyzedImageSrc] = useState('');
  const worker = useRef(null);
  const [count, setCount] = useState(0);

  useEffect(() => {
    worker.current = new Worker(new URL('./imageWorker.js', import.meta.url));

    worker.current.onmessage = (event) => {
      if (event.data.type === 'initialized') {
        setIsModelReady(true);
      } else if (event.data.type === 'analyzed') {
        const { analyzedImageBlob } = event.data;
        const url = URL.createObjectURL(analyzedImageBlob);
        setAnalyzedImageSrc(url);
        setIsAnalyzed(true);
        setIsAnalysing(false);
        setLoading(false);
      }
    };

    worker.current.postMessage({ type: 'init' });

    return () => {
      if (worker.current) {
        worker.current.terminate();
      }
    };
  }, []);

  const onUpload = ({ file }) => {
    if (file.status === 'removed') {
      setImageSrc('');
      setAnalyzedImageSrc('');
      setIsUploaded(false);
      setIsAnalyzed(false);
      return;
    }
    const reader = new FileReader();
    reader.onload = (e) => {
      setIsUploaded(true);
      setImageSrc(e.target.result);
      analysisImage(e.target.result);
    };
    reader.readAsDataURL(file);
  };

  const analysisImage = (url) => {
    setIsAnalysing(true);
    setLoading(true);
    worker.current.postMessage({ type: 'analyze', url });
  };

  return (
    <div className="remove-bgc-container">
      <Spin spinning={loading}>
        <Upload onChange={onUpload} beforeUpload={() => false} maxCount={1}>
          <Button loading={(!isModelReady || isAnalysing)} disabled={(!isModelReady || isAnalysing)} onClick={() => { }}>
            {isModelReady ? isAnalysing ? '正在处理...' : '点击上传' : '模型加载中'}
          </Button>
        </Upload>
        <div className="image-container">
          {isUploaded && <div>
            您上传的图片: <Image src={imageSrc} width={200}></Image>
          </div>}
          {isAnalyzed && <div>
            处理后的图片: <Image src={analyzedImageSrc} width={200}></Image>
          </div>}
        </div>
      </Spin>
      <Button onClick={() => setCount(count => count + 1)}>点击测试一下dom</Button>
      <div>{count}</div>
    </div>
  );
};

export default RemoveBgc;

5.修改App.jsx:

import React from "react";
// import Transform from "./components/transform";
// import RemoveBgc from "./components/removeBgc";
import RemoveBgc from "./components/removerBgc-useWorker";

function App() {
  return (
    <div>
      {/* <Transform></Transform> */}
      <RemoveBgc></RemoveBgc>
    </div>
  );
}

export default App;

6.效果图:可以点击dom做阻塞测试

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐