Zenn
🖼️

サクっとスクショ取ってMobileNet認識するChrome拡張機能サンプル( React,TS, Chakra, TensorFlow)

2023/01/15に公開
Chrome
React
TensorFlow
Chrome 拡張機能
Chakra UI
tech

概要

Chrome拡張機能で作ってみたいアイディアがあったので、Chrome拡張機能の勉強がてら、サンプルまでを用意してみた. 手探り含めて大体2時間ぐらいで作り終える規模感. 非常にサンプルなどが充実しており詰まらずできてしまう.

成果物

https://github.com/UrusuLambda/ScreenShotAndClassifyChromeExtensionExample

sscce-gif

材料

  • React
  • TypeScript
  • TensorFlow.js / MobileNet
  • ChakraUI

上記を使用してChrome拡張機能のサンプルを実装.

手順

React/TypeScriptベースのChrome拡張機能を用意

Chrome拡張機能自体を学びたい場合は下記を参照.
https://developer.chrome.com/docs/extensions/mv3/getstarted/extensions-101/

  1. まずはcreate-react-appでプロジェクトを作成
npx create-react-app sample --template typescript
  1. public/manifest.jsonをChrome拡張機能用に編集
{
  "name": "Chrome ScreenShot and MobileNet Classify  Extension",
  "description": "Chrome Extension to capture screen and classify the screen with React ,Chakra,TensorFlow.js",
  "version": "1.0",
  "manifest_version": 3,
  "action": {
    "default_popup": "index.html",
    "default_title": "ScreenShot and Classify"
  },
  "short_name": "SSC",
  "icons": {
    "16": "logo192.png",
    "48": "logo192.png",
    "128": "logo192.png"
  },
  "permissions": ["activeTab"]
}
  1. buildして生成
npm run build
  1. Chromeの拡張機能のページを開き、デベロッパーモードに変更
  2. "パッケージ化されていない拡張機能を読み込む"からbuildディレクトリを指定
    setting-chrome
  3. 右上から開けることを確認
    react-sample
    ここまででChrome拡張機能としては準備完了.

ChakraUIを組み込む

これも秒.
参考 : https://chakra-ui.com/getting-started

  1. install
npm install @chakra-ui/react @emotion/react @emotion/styled framer-motion
  1. App.tsxを下記で置き換え. ChakraProviderで括ってあげることのみがポイント.
import * as React from "react";
import { ChakraProvider, Button } from "@chakra-ui/react";

function App() {
  return (
    <ChakraProvider>
      <Button bgColor="green.500" color="white" boxShadow="base">
        Click
      </Button>
    </ChakraProvider>
  );
}

export default App;
  1. 上記で実装は終わったので
npm run build
  1. 拡張機能のページで、再読み込みボタンをクリックして確認
    reload
    clickpage

Chrome拡張機能からScreenShotを取って表示

Buttonをクリックしたらスクリーンショットを取ってPopupに表示する.
これはChromeが提供するcaptureVisibleTabを使用する必要がある.
スクリーンショットを取るには、事前にmanifest.jsonのpermissionsにactiveTabを記載する必要があるが、前の手順ですでにmanifest.jsonは更新している.

  1. captureVisibleTabを使うにあたって下記を入れておく.
 npm install @types/chrome --save-dev
  1. 今のままでは画面が小さいので、index.cssを編集して、bodyに下記を追記.
  width: 600px;
  1. App.tsxを編集.

変更点は、

  • スクショを確認するImageを追加.
  • Imageのsrcが見るimageObjectをuseStateで用意.
  • handleCapture関数をButtonのonClickに付与.
  • captureVisibleTabで取得したBase64画像をsetして、Imageを更新.
import * as React from "react";

import { ChakraProvider, Button, Image } from "@chakra-ui/react";

function App() {
  let [imageObject, setImageObjcet] = React.useState("");

  const handleCapture = function () {
    (chrome.tabs as any).captureVisibleTab(
      null,
      { format: "jpeg" },
      (base64Data: any) => {
        setImageObjcet(base64Data);
      }
    );
  };

  return (
    <ChakraProvider>
      <Button
        bgColor="green.500"
        color="white"
        boxShadow="base"
        onClick={() => {
          handleCapture();
        }}
      >
        Click
      </Button>
      <Image
        src={imageObject}
        width="90%"
        height="300px"
        objectFit="contain"
      ></Image>
    </ChakraProvider>
  );
}

export default App;
  1. 確認. Clickを押すと、スクリーンショットが撮られます.

ss

TensorFlow.jsのMobileNetで取得した画像を分類

ここからTensorFlow.jsの出番.

  1. install
npm i @tensorflow-models/mobilenet  
npm install @tensorflow/tfjs
  1. App.tsxにて下記を追加

最後に全体を載せますが、差分は下記3点.

import * as tf from "@tensorflow/tfjs";
import { MobileNet } from "@tensorflow-models/mobilenet";
import * as mobilenet from "@tensorflow-models/mobilenet";

また、Imageの要素が変わったときにuseEffectで更新をかける処理をhandleCaptureの前に挿入.

  let [classifyResult, setClassifyResult] = React.useState<
    Array<{ className: string; probability: number }>
  >([]);

  React.useEffect(() => {
    (async () => {
      if (imageObject) {
        await tf.setBackend("webgl");
        const model: MobileNet = await mobilenet.load();

        let img_element = document.createElement("img");
        img_element.onload = async function () {
          const classifyPredictions: Array<{
            className: string;
            probability: number;
          }> = await model.classify(img_element);
          setClassifyResult(classifyPredictions);
        };
        img_element.src = imageObject;
      }
    })();
  }, [imageObject]);

結果画像を下記にて表示. Imageの後に追加.

{imageObject ? (
        <TableContainer
          width="90%"
          margin="10px"
          borderRadius="3px"
          boxShadow="base"
        >
          <Table variant="striped" colorScheme="teal">
            <TableCaption>MobileNet Classify Result</TableCaption>
            <Thead>
              <Tr>
                <Th>Classs</Th>
                <Th isNumeric>Probability</Th>
              </Tr>
            </Thead>
            <Tbody>
              {classifyResult.map((result, _index) => (
                <Tr>
                  <Td>{result.className}</Td>
                  <Td isNumeric>{result.probability}</Td>
                </Tr>
              ))}
            </Tbody>
          </Table>
        </TableContainer>
      ) : null}

これらを入れると下記のようになるはず.

import * as React from "react";

import {
  ChakraProvider,
  Button,
  Image,
  Table,
  TableCaption,
  TableContainer,
  Tbody,
  Td,
  Th,
  Thead,
  Tr,
} from "@chakra-ui/react";

import * as tf from "@tensorflow/tfjs";
import { MobileNet } from "@tensorflow-models/mobilenet";
import * as mobilenet from "@tensorflow-models/mobilenet";

function App() {
  let [imageObject, setImageObjcet] = React.useState("");

  let [classifyResult, setClassifyResult] = React.useState<
    Array<{ className: string; probability: number }>
  >([]);

  React.useEffect(() => {
    (async () => {
      if (imageObject) {
        await tf.setBackend("webgl");
        const model: MobileNet = await mobilenet.load();

        let img_element = document.createElement("img");
        img_element.onload = async function () {
          const classifyPredictions: Array<{
            className: string;
            probability: number;
          }> = await model.classify(img_element);
          setClassifyResult(classifyPredictions);
        };
        img_element.src = imageObject;
      }
    })();
  }, [imageObject]);

  const handleCapture = function () {
    (chrome.tabs as any).captureVisibleTab(
      null,
      { format: "jpeg" },
      (base64Data: any) => {
        setImageObjcet(base64Data);
      }
    );
  };

  return (
    <ChakraProvider>
      <Button
        bgColor="green.500"
        color="white"
        boxShadow="base"
        onClick={() => {
          handleCapture();
        }}
      >
        Click
      </Button>
      <Image
        src={imageObject}
        width="90%"
        height="300px"
        objectFit="contain"
      ></Image>
      {imageObject ? (
        <TableContainer
          width="90%"
          margin="10px"
          borderRadius="3px"
          boxShadow="base"
        >
          <Table variant="striped" colorScheme="teal">
            <TableCaption>MobileNet Classify Result</TableCaption>
            <Thead>
              <Tr>
                <Th>Classs</Th>
                <Th isNumeric>Probability</Th>
              </Tr>
            </Thead>
            <Tbody>
              {classifyResult.map((result, _index) => (
                <Tr>
                  <Td>{result.className}</Td>
                  <Td isNumeric>{result.probability}</Td>
                </Tr>
              ))}
            </Tbody>
          </Table>
        </TableContainer>
      ) : null}
    </ChakraProvider>
  );
}

export default App;
  1. 確認.

classify

以上!

あとは色々とデザインを整えれば、良い感じの拡張機能になります.

Discussion