java使用Deep Java Library(djl)搭配TorchScript搭建图片分类

一、前置要求

1.1、下载TorchScript类型的模型，注意这里是TorchScript类型，有些模型在说明中会说明是否为该格式的文件。可以从huggingface下载，在huggingface注意未区分PyTorch和TorchScript，在模型下方的标签都标记的为PyTorch，需要看具体的描述是否说该模型为TorchScript。
1.2、pom文件引入依赖，注意和引擎相关的包需要搭配引用，例如ai.djl.pytorch的native和jni包与engine版本要对上。pom.xml引入包如下：

<properties>
		<maven.compiler.source>11</maven.compiler.source>
		<maven.compiler.target>11</maven.compiler.target>
		<djl.version>0.27.0</djl.version>
	</properties>


	<dependencies>
		<!-- https://mvnrepository.com/artifact/ai.djl/api -->
		<dependency>
			<groupId>ai.djl</groupId>
			<artifactId>api</artifactId>
			<version>${djl.version}</version>
		</dependency>

		<!-- https://mvnrepository.com/artifact/ai.djl/model-zoo -->
		<dependency>
			<groupId>ai.djl</groupId>
			<artifactId>model-zoo</artifactId>
			<version>${djl.version}</version>
		</dependency>

		<!-- https://mvnrepository.com/artifact/ai.djl.pytorch/pytorch-engine -->
		<dependency>
			<groupId>ai.djl.pytorch</groupId>
			<artifactId>pytorch-engine</artifactId>
			<version>${djl.version}</version>
		</dependency>

		<dependency>
			<groupId>ai.djl</groupId>
			<artifactId>basicdataset</artifactId>
			<version>${djl.version}</version>
		</dependency>
	
		<dependency>
			<groupId>ai.djl.pytorch</groupId>
			<artifactId>pytorch-engine</artifactId>
			<version>${djl.version}</version>
		</dependency>
		
		<dependency>
			<groupId>ai.djl.pytorch</groupId>
			<artifactId>pytorch-jni</artifactId>
			<version>2.1.1-0.27.0</version>
		</dependency>
		<dependency>
			<groupId>ai.djl.pytorch</groupId>
			<artifactId>pytorch-native-cpu</artifactId>
			<classifier>win-x86_64</classifier>
			<version>2.1.1</version>
		</dependency>

		<dependency>
			<groupId>ai.djl</groupId>
			<artifactId>djl-zero</artifactId>
			<version>${djl.version}</version>
		</dependency>

		<dependency>
			<groupId>org.apache.logging.log4j</groupId>
			<artifactId>log4j-slf4j-impl</artifactId>
			<version>2.21.0</version>
		</dependency>
	</dependencies>

二、java代码

将下载好的模型放到对应的位置，其中模型文件包含两个部分，一个是pt结尾的文件，当然结尾不一定是这个，可能是其他的，可以使用压缩文件打开这个模型文件看看是否包含了constants.pkl等文件，这个可以用作确认是否为TorchScript的一个标志。然后还需要一个synset.txt文件。

//这里也可以使用在线的模型
private static final URL MODEL_URL = NSFWUtil.class.getClassLoader().getResource("model/xxx.pt");

	public static void main(String[] args) throws MalformedModelException, IOException, ModelNotFoundException, TranslateException {
		getNSFW4JSON("image path");
	}
	
	/**
	 * 
	 * @param imagePath 文件地址
	 * @throws ModelNotFoundException 
	 * @throws MalformedModelException
	 * @throws IOException
	 * @throws TranslateException
	 * @return nsfw的json
	 */
	public static Classifications  getNSFW4JSON(String imagePath) throws ModelNotFoundException, MalformedModelException, IOException, TranslateException {
		Image img = ImageFactory.getInstance().fromFile(Paths.get(imagePath));
		Translator<Image, Classifications> translator =
                ImageClassificationTranslator.builder()
                		.addTransform(new Resize(224, 224))
                        .addTransform(new ToTensor())
                        .optApplySoftmax(true)
                        .build();
		Criteria<Image, Classifications> criteria = Criteria.builder()
               .setTypes(Image.class, Classifications.class)
               .optModelUrls(MODEL_URL.toString())
               .optTranslator(translator)
               .optEngine("PyTorch") // Use PyTorch engine
               .optProgress(new ProgressBar())
               .build();
		try (ZooModel<Image, Classifications> model = criteria.loadModel())
		{
           Predictor<Image, Classifications> predictor = model.newPredictor();
           return predictor.predict(img);
       }
	}
	
	/**
	 * 
	 * @param in 输入流
	 * @throws ModelNotFoundException 
	 * @throws MalformedModelException
	 * @throws IOException
	 * @throws TranslateException
	 * @return nsfw的json
	 */
	public static Classifications  getNSFW4JSON(InputStream in) throws ModelNotFoundException, MalformedModelException, IOException, TranslateException {
		Image img = BufferedImageFactory.getInstance().fromInputStream(in);
		Translator<Image, Classifications> translator =
                ImageClassificationTranslator.builder()
                		.addTransform(new Resize(224, 224))
                        .addTransform(new ToTensor())
                        .optApplySoftmax(true)
                        .build();
		Criteria<Image, Classifications> criteria = Criteria.builder()
               .setTypes(Image.class, Classifications.class)
               .optModelUrls(MODEL_URL.toString())
               .optTranslator(translator)
               .optEngine("PyTorch") // Use PyTorch engine
               .optProgress(new ProgressBar())
               .build();
		try (ZooModel<Image, Classifications> model = criteria.loadModel())
		{
           Predictor<Image, Classifications> predictor = model.newPredictor();
           return predictor.predict(img);
       }
	}

三、总结

3.1、代码中的ImageClassificationTranslator在其他很多时候是自己在定义具体的方法实现，这里我们是图片分类，所以我们用的是官方提供的Translator。
3.2、就目前来说框架帮我们实现了很多的代码，能写的代码不是很多，难点在于如何找到能用的模型，目前很多模型还是PyTorch类型的，无法在JAVA或者C++环境调用。
3.3、可以试一下的模型nsfw,记住下synset.txt