Java調(diào)用Pytorch模型實(shí)現(xiàn)圖像識(shí)別

更新時(shí)間：2023年06月12日 11:13:48 作者：歐內(nèi)的手好汗

這篇文章主要為大家詳細(xì)介紹了Java如何調(diào)用Pytorch實(shí)現(xiàn)圖像識(shí)別功能，文中的示例代碼講解詳細(xì)，具有一定的學(xué)習(xí)價(jià)值，感興趣的小伙伴可以了解一下

1. 依賴

<dependency>
    <groupId>ai.djl.pytorch</groupId>
    <artifactId>pytorch-engine</artifactId>
    <version>0.16.0</version>
</dependency>
<dependency>
    <groupId>ai.djl.pytorch</groupId>
    <artifactId>pytorch-native-auto</artifactId>
    <version>1.9.1</version>
    <scope>runtime</scope>
</dependency>
<dependency>
    <groupId>ai.djl.pytorch</groupId>
    <artifactId>pytorch-jni</artifactId>
    <version>1.9.1-0.16.0</version>
    <scope>runtime</scope>
</dependency>

2. 準(zhǔn)備模型

1.首先將模型按下面方法保存，放到項(xiàng)目resources中

import torch
# An instance of your model.
model = MyModel(num_classes = 80)
# Switch the model to eval model
model.eval()
# An example input you would normally provide to your model's forward() method.
example = torch.rand(1, 3, 224, 224)
# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model, example)
# Save the TorchScript model
traced_script_module.save("model.pt")

2.編寫(xiě)工具類，用于完成識(shí)別功能

public class HerbUtil {
    //規(guī)定輸入尺寸
    private static final int INPUT_SIZE = 224;
    //標(biāo)簽文件 一種類別名字占一行
    private List<String> herbNames;
    //用于識(shí)別
    Predictor<Image, Classifications> predictor;
    //模型
    private Model model;
    public HerbUtil() {
        //加載標(biāo)簽到herbNames中
        this.loadHerbNames();
        //初始化模型工作
        this.init();
    }
}

3.將標(biāo)簽文件放到resources中，載入標(biāo)簽

   private void loadHerbNames() {
        BufferedReader reader = null;
        herbNames = new ArrayList<>();
        try {
            InputStream in = HerbUtil.class.getClassLoader().getResourceAsStream("names.txt");
            reader = new BufferedReader(new InputStreamReader(in));
            String name = null;
            while ((name = reader.readLine()) != null) {
                herbNames.add(name);
            }
            System.out.println(herbNames);
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (reader != null) {
                try {
                    reader.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }

4.初始化模型

   private void init() {
        Translator<Image, Classifications> translator = ImageClassificationTranslator.builder()
                //下面的transform根據(jù)自己的改
                .addTransform(new RandomResizedCrop(INPUT_SIZE, INPUT_SIZE, 0.6, 1,
                        3. / 4, 4. / 3))
                .addTransform(new ToTensor())
                .addTransform(new Normalize(
                        new float[] {0.5f, 0.5f, 0.5f},
                        new float[] {0.5f, 0.5f, 0.5f}))
                //如果你的模型最后一層沒(méi)有經(jīng)過(guò)softmax就啟用它
                .optApplySoftmax(true)
                //載入所有標(biāo)簽進(jìn)去
                .optSynset(herbNames)
                //最終顯示概率最高的5個(gè)
                .optTopK(5)
                .build();
        //隨便起名
        Model model = Model.newInstance("model", Device.cpu());
        try {
            InputStream inputStream = HerbUtil.class.getClassLoader().getResourceAsStream("model.pt");
            if (inputStream == null) {
                throw new RuntimeException("找不到模型文件");
            }
            model.load(inputStream);
            predictor = model.newPredictor(translator);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

5.我開(kāi)頭提到的圖片預(yù)處理的代碼

   private Image resizeImage(InputStream inputStream) {
        BufferedImage input = null;
        try {
            input = ImageIO.read(inputStream);
        } catch (IOException e) {
            e.printStackTrace();
        }
        int iw = input.getWidth(), ih = input.getHeight();
        int w = 224, h = 224;
        double scale = Math.min(1. *  w / iw, 1. * h / ih);
        int nw = (int) (iw * scale), nh = (int) (ih * scale);
        java.awt.Image img;
        //只有太長(zhǎng)或太寬才會(huì)保留橫縱比，填充顏色
        boolean needResize = 1. * iw / ih > 1.4 || 1. * ih / iw > 1.4;
        if (needResize) {
            img = input.getScaledInstance(nw, nh, BufferedImage.SCALE_SMOOTH);
        } else {
            img = input.getScaledInstance(INPUT_SIZE, INPUT_SIZE, BufferedImage.SCALE_SMOOTH);
        }
        BufferedImage out = new BufferedImage(INPUT_SIZE, INPUT_SIZE, BufferedImage.TYPE_INT_RGB);
        Graphics g = out.getGraphics();
        //先將整個(gè)224*224區(qū)域填充128 128 128顏色
        g.setColor(new Color(128, 128, 128));
        g.fillRect(0, 0, INPUT_SIZE, INPUT_SIZE);
        out.getGraphics().drawImage(img, 0, needResize ? (INPUT_SIZE - nh) / 2 : 0, null);
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        try {
            ImageOutputStream imageOutputStream = ImageIO.createImageOutputStream(outputStream);
            ImageIO.write(out, "jpg", imageOutputStream);
            //去D盤(pán)看效果
            //ImageIO.write(out, "jpg", new File("D:\\out.jpg"));
            InputStream is = new ByteArrayInputStream(outputStream.toByteArray());
            return ImageFactory.getInstance().fromInputStream(is);
        } catch (IOException e) {
            e.printStackTrace();
            throw new RuntimeException("圖片轉(zhuǎn)換失敗");
        }
    }

6.識(shí)別功能

    public List<Classifications.Classification> predict(InputStream inputStream) {
        List<Classifications.Classification> result = new ArrayList<>();
        Image input = this.resizeImage(inputStream);
        try {
            Classifications output = predictor.predict(input);
            System.out.println("推測(cè)為：" + output.best().getClassName()
                    + ", 概率：" + output.best().getProbability());
            System.out.println(output);
            result = output.topK();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return result;
    }

3. 測(cè)試

    @Test
    public void test7() {
        HerbUtil herbUtil = new HerbUtil();
        String path = "E:\\深度學(xué)習(xí)專用\\data\\train\\當(dāng)歸\\24.jpeg";
        try {
            File file = new File(path);
            InputStream inputStream = new FileInputStream(file);
            herbUtil.predict(inputStream);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

輸出：

加入到項(xiàng)目中后，工具類直接Autowire注入或者方法都寫(xiě)static的，Controller接收前端MultipartFile，將其inputstream用于推測(cè)

如果你想加載網(wǎng)絡(luò)圖片，那就去網(wǎng)上搜索怎么把它轉(zhuǎn)成inputstream吧

測(cè)試多線程一起predict時(shí)報(bào)錯(cuò)了

4.更新

當(dāng)我打包成jar到centos7的linux中運(yùn)行時(shí)，報(bào)錯(cuò)UnsatisfiedLinkError，經(jīng)過(guò)大神的指導(dǎo)，問(wèn)題出來(lái)我引的依賴。

修改后的依賴：

    <properties>
        <java.version>8</java.version>
        <jna.version>5.3.0</jna.version>
    </properties>
    <dependencies>
        <dependency>
            <groupId>ai.djl.pytorch</groupId>
            <artifactId>pytorch-engine</artifactId>
            <version>0.16.0</version>
        </dependency>
        <dependency>
            <groupId>ai.djl.pytorch</groupId>
            <artifactId>pytorch-native-cpu-precxx11</artifactId>
            <classifier>linux-x86_64</classifier>
            <version>1.9.1</version>
            <scope>runtime</scope>
        </dependency>
        <dependency>
            <groupId>ai.djl.pytorch</groupId>
            <artifactId>pytorch-jni</artifactId>
            <version>1.9.1-0.16.0</version>
            <scope>runtime</scope>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-web</artifactId>
        </dependency>
    </dependencies>

到此這篇關(guān)于Java調(diào)用Pytorch模型實(shí)現(xiàn)圖像識(shí)別的文章就介紹到這了,更多相關(guān)Java Pytorch圖像識(shí)別內(nèi)容請(qǐng)搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！

您可能感興趣的文章: