initial commit ; currently it wont build :C
This commit is contained in:
77
src/main/java/org/jdetect/ModelLoader.java
Normal file
77
src/main/java/org/jdetect/ModelLoader.java
Normal file
@@ -0,0 +1,77 @@
|
||||
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;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user