77 lines
2.3 KiB
Java
77 lines
2.3 KiB
Java
package org.jdetect;
|
|
|
|
import org.opencv.core.*;
|
|
import org.opencv.dnn.Dnn;
|
|
import org.opencv.dnn.Net;
|
|
|
|
import java.io.BufferedReader;
|
|
import java.io.IOException;
|
|
import java.io.InputStream;
|
|
import java.io.InputStreamReader;
|
|
import java.util.ArrayList;
|
|
import java.util.List;
|
|
import java.nio.file.Files;
|
|
import java.nio.file.Paths;
|
|
|
|
|
|
public class ModelLoader {
|
|
private static Net net;
|
|
private static List<String> classes;
|
|
private static List<String> outputLayers;
|
|
|
|
public static synchronized void loadModel() throws IOException {
|
|
if (net != null) return;
|
|
|
|
// Load network
|
|
InputStream weightsStream = ModelLoader.class.getResourceAsStream("/yolov4-tiny.weights");
|
|
InputStream configStream = ModelLoader.class.getResourceAsStream("/yolov4-tiny.cfg");
|
|
|
|
assert weightsStream != null;
|
|
byte[] weights = weightsStream.readAllBytes();
|
|
assert configStream != null;
|
|
byte[] config = configStream.readAllBytes();
|
|
|
|
// Create temp files
|
|
String weightsPath = "temp_weights.weights";
|
|
String configPath = "temp_config.cfg";
|
|
|
|
Files.write(Paths.get(weightsPath), weights);
|
|
Files.write(Paths.get(configPath), config);
|
|
|
|
net = Dnn.readNetFromDarknet(configPath, weightsPath);
|
|
|
|
// Load classes
|
|
classes = new ArrayList<>();
|
|
try (InputStream classesStream = ModelLoader.class.getResourceAsStream("/coco.names")) {
|
|
assert classesStream != null;
|
|
try (BufferedReader br = new BufferedReader(new InputStreamReader(classesStream))) {
|
|
String line;
|
|
while ((line = br.readLine()) != null) {
|
|
classes.add(line.trim());
|
|
}
|
|
}
|
|
}
|
|
|
|
// Get output layers
|
|
List<String> layerNames = net.getLayerNames();
|
|
MatOfInt unconnectedOutLayers = net.getUnconnectedOutLayers();
|
|
int[] indices = unconnectedOutLayers.toArray();
|
|
|
|
outputLayers = new ArrayList<>();
|
|
for (int idx : indices) {
|
|
outputLayers.add(layerNames.get(idx - 1));
|
|
}
|
|
}
|
|
|
|
public static Net getNet() {
|
|
return net;
|
|
}
|
|
|
|
public static List<String> getClasses() {
|
|
return classes;
|
|
}
|
|
|
|
public static List<String> getOutputLayers() {
|
|
return outputLayers;
|
|
}
|
|
} |