Skip to content

Commit a500341

Browse files
samfreundsamfreund
committed
fix tests
1 parent 9175577 commit a500341

1 file changed

Lines changed: 136 additions & 131 deletions

File tree

‎photon-core/src/main/java/org/photonvision/common/configuration/NeuralNetworkModelManager.java‎

Lines changed: 136 additions & 131 deletions
Original file line numberDiff line numberDiff line change
@@ -46,11 +46,15 @@
4646
/**
4747
* Manages the loading of neural network models.
4848
*
49-
* <p>Models are loaded from the filesystem at the <code>modelsFolder</code> location. PhotonVision
50-
* also supports shipping pre-trained models as resources in the JAR. If the model has already been
49+
* <p>
50+
* Models are loaded from the filesystem at the <code>modelsFolder</code>
51+
* location. PhotonVision
52+
* also supports shipping pre-trained models as resources in the JAR. If the
53+
* model has already been
5154
* extracted to the filesystem, it will not be extracted again.
5255
*
53-
* <p>Each model must have a corresponding {@link ModelProperties} entry in {@link
56+
* <p>
57+
* Each model must have a corresponding {@link ModelProperties} entry in {@link
5458
* NeuralNetworkPropertyManager}.
5559
*/
5660
public class NeuralNetworkModelManager {
@@ -60,95 +64,95 @@ public class NeuralNetworkModelManager {
6064
private final List<Family> supportedBackends = new ArrayList<>();
6165

6266
/**
63-
* This function stores the properties of the shipped object detection models. It is stored as a
67+
* This function stores the properties of the shipped object detection models.
68+
* It is stored as a
6469
* function so that it can be dynamic, to adjust for the models directory.
6570
*/
6671
private NeuralNetworkPropertyManager getShippedProperties(File modelsDirectory) {
6772
NeuralNetworkPropertyManager nnProps = new NeuralNetworkPropertyManager();
6873

69-
LinkedList<String> cocoLabels =
70-
new LinkedList<String>(
71-
List.of(
72-
"person",
73-
"bicycle",
74-
"car",
75-
"motorcycle",
76-
"airplane",
77-
"bus",
78-
"train",
79-
"truck",
80-
"boat",
81-
"traffic light",
82-
"fire hydrant",
83-
"stop sign",
84-
"parking meter",
85-
"bench",
86-
"bird",
87-
"cat",
88-
"dog",
89-
"horse",
90-
"sheep",
91-
"cow",
92-
"elephant",
93-
"bear",
94-
"zebra",
95-
"giraffe",
96-
"backpack",
97-
"umbrella",
98-
"handbag",
99-
"tie",
100-
"suitcase",
101-
"frisbee",
102-
"skis",
103-
"snowboard",
104-
"sports ball",
105-
"kite",
106-
"baseball bat",
107-
"baseball glove",
108-
"skateboard",
109-
"surfboard",
110-
"tennis racket",
111-
"bottle",
112-
"wine glass",
113-
"cup",
114-
"fork",
115-
"knife",
116-
"spoon",
117-
"bowl",
118-
"banana",
119-
"apple",
120-
"sandwich",
121-
"orange",
122-
"broccoli",
123-
"carrot",
124-
"hot dog",
125-
"pizza",
126-
"donut",
127-
"cake",
128-
"chair",
129-
"couch",
130-
"potted plant",
131-
"bed",
132-
"dining table",
133-
"toilet",
134-
"tv",
135-
"laptop",
136-
"mouse",
137-
"remote",
138-
"keyboard",
139-
"cell phone",
140-
"microwave",
141-
"oven",
142-
"toaster",
143-
"sink",
144-
"refrigerator",
145-
"book",
146-
"clock",
147-
"vase",
148-
"scissors",
149-
"teddy bear",
150-
"hair drier",
151-
"toothbrush"));
74+
LinkedList<String> cocoLabels = new LinkedList<String>(
75+
List.of(
76+
"person",
77+
"bicycle",
78+
"car",
79+
"motorcycle",
80+
"airplane",
81+
"bus",
82+
"train",
83+
"truck",
84+
"boat",
85+
"traffic light",
86+
"fire hydrant",
87+
"stop sign",
88+
"parking meter",
89+
"bench",
90+
"bird",
91+
"cat",
92+
"dog",
93+
"horse",
94+
"sheep",
95+
"cow",
96+
"elephant",
97+
"bear",
98+
"zebra",
99+
"giraffe",
100+
"backpack",
101+
"umbrella",
102+
"handbag",
103+
"tie",
104+
"suitcase",
105+
"frisbee",
106+
"skis",
107+
"snowboard",
108+
"sports ball",
109+
"kite",
110+
"baseball bat",
111+
"baseball glove",
112+
"skateboard",
113+
"surfboard",
114+
"tennis racket",
115+
"bottle",
116+
"wine glass",
117+
"cup",
118+
"fork",
119+
"knife",
120+
"spoon",
121+
"bowl",
122+
"banana",
123+
"apple",
124+
"sandwich",
125+
"orange",
126+
"broccoli",
127+
"carrot",
128+
"hot dog",
129+
"pizza",
130+
"donut",
131+
"cake",
132+
"chair",
133+
"couch",
134+
"potted plant",
135+
"bed",
136+
"dining table",
137+
"toilet",
138+
"tv",
139+
"laptop",
140+
"mouse",
141+
"remote",
142+
"keyboard",
143+
"cell phone",
144+
"microwave",
145+
"oven",
146+
"toaster",
147+
"sink",
148+
"refrigerator",
149+
"book",
150+
"clock",
151+
"vase",
152+
"scissors",
153+
"teddy bear",
154+
"hair drier",
155+
"toothbrush"));
152156

153157
nnProps.addModelProperties(
154158
new ModelProperties(
@@ -200,16 +204,10 @@ private NeuralNetworkPropertyManager getShippedProperties(File modelsDirectory)
200204
*/
201205
private NeuralNetworkModelManager() {
202206
switch (Platform.getCurrentPlatform()) {
203-
case LINUX_QCS6490 -> {
204-
if (LoadJNI.JNITypes.RUBIK_DETECTOR.hasLoaded()) {
205-
supportedBackends.add(Family.RUBIK);
206-
}
207-
}
208-
case LINUX_RK3588_64 -> {
209-
if (LoadJNI.JNITypes.RKNN_DETECTOR.hasLoaded()) {
210-
supportedBackends.add(Family.RKNN);
211-
}
212-
}
207+
case LINUX_QCS6490 -> supportedBackends.add(Family.RUBIK);
208+
209+
case LINUX_RK3588_64 -> supportedBackends.add(Family.RKNN);
210+
213211
default -> {
214212
logger.warn(
215213
"No supported neural network backends found for this platform: "
@@ -221,7 +219,8 @@ private NeuralNetworkModelManager() {
221219
}
222220

223221
/**
224-
* Returns the singleton instance of the NeuralNetworkModelManager. Call getInstance() to use the
222+
* Returns the singleton instance of the NeuralNetworkModelManager. Call
223+
* getInstance() to use the
225224
* default (no reset), or getInstance(true) to reset.
226225
*
227226
* @return The singleton instance
@@ -231,7 +230,8 @@ public static NeuralNetworkModelManager getInstance() {
231230
}
232231

233232
/**
234-
* Returns the singleton instance of the NeuralNetworkModelManager, optionally resetting it.
233+
* Returns the singleton instance of the NeuralNetworkModelManager, optionally
234+
* resetting it.
235235
*
236236
* @param reset If true, resets the instance
237237
* @return The singleton instance
@@ -279,18 +279,23 @@ public List<String> getSupportedBackends() {
279279
/**
280280
* Stores model information, such as the model file, labels, and version.
281281
*
282-
* <p>The first model in the list is the default model.
282+
* <p>
283+
* The first model in the list is the default model.
283284
*/
284285
private Map<Family, ArrayList<Model>> models;
285286

286287
/**
287-
* Retrieves the model with the specified name, assuming it is available under a supported
288+
* Retrieves the model with the specified name, assuming it is available under a
289+
* supported
288290
* backend.
289291
*
290-
* <p>If this method returns `Optional.of(..)` then the model should be safe to load.
292+
* <p>
293+
* If this method returns `Optional.of(..)` then the model should be safe to
294+
* load.
291295
*
292296
* @param modelUID the unique identifier of the model to retrieve
293-
* @return an Optional containing the model if found, or an empty Optional if not found
297+
* @return an Optional containing the model if found, or an empty Optional if
298+
* not found
294299
*/
295300
public Optional<Model> getModel(String modelUID) {
296301
if (models == null) {
@@ -300,8 +305,8 @@ public Optional<Model> getModel(String modelUID) {
300305
// Check if the model exists in any supported backend
301306
for (Family backend : supportedBackends) {
302307
if (models.containsKey(backend)) {
303-
Optional<Model> model =
304-
models.get(backend).stream().filter(m -> m.getUID().equals(modelUID)).findFirst();
308+
Optional<Model> model = models.get(backend).stream().filter(m -> m.getUID().equals(modelUID))
309+
.findFirst();
305310
if (model.isPresent()) {
306311
return model;
307312
}
@@ -326,8 +331,8 @@ private void loadModel(Path path) {
326331
models = new HashMap<>();
327332
}
328333

329-
ModelProperties properties =
330-
ConfigManager.getInstance().getConfig().neuralNetworkPropertyManager().getModel(path);
334+
ModelProperties properties = ConfigManager.getInstance().getConfig().neuralNetworkPropertyManager()
335+
.getModel(path);
331336

332337
if (properties == null) {
333338
logger.error(
@@ -371,7 +376,9 @@ private void loadModel(Path path) {
371376
/**
372377
* Discovers DNN models from the specified folder.
373378
*
374-
* <p>This makes the assumption that all of the models have their properties stored in the
379+
* <p>
380+
* This makes the assumption that all of the models have their properties stored
381+
* in the
375382
* database
376383
*/
377384
public void discoverModels() {
@@ -390,9 +397,8 @@ public void discoverModels() {
390397
files
391398
.filter(Files::isRegularFile)
392399
.filter(
393-
path ->
394-
supportedBackends.stream()
395-
.anyMatch(family -> path.toString().endsWith(family.extension())))
400+
path -> supportedBackends.stream()
401+
.anyMatch(family -> path.toString().endsWith(family.extension())))
396402
.forEach(this::loadModel);
397403
} catch (IOException e) {
398404
logger.error("Failed to discover models at " + modelsDirectory.getAbsolutePath(), e);
@@ -415,7 +421,8 @@ public void discoverModels() {
415421
}
416422

417423
/**
418-
* Extracts models from the JAR and copies them to disk. Also copies properties into the database.
424+
* Extracts models from the JAR and copies them to disk. Also copies properties
425+
* into the database.
419426
*/
420427
public void extractModels() {
421428
File modelsDirectory = ConfigManager.getInstance().getModelsDirectory();
@@ -431,7 +438,8 @@ public void extractModels() {
431438
}
432439
}
433440

434-
// Used for checking if the model to be extracted is supported for this architecture
441+
// Used for checking if the model to be extracted is supported for this
442+
// architecture
435443
ArrayList<String> supportedModelFileNames = new ArrayList<String>();
436444
for (ModelProperties model : supportedProperties.getModels()) {
437445
supportedModelFileNames.add(model.modelPath().getFileName().toString());
@@ -444,17 +452,16 @@ public void extractModels() {
444452
String resource = "models";
445453

446454
try {
447-
String jarPath =
448-
getClass().getProtectionDomain().getCodeSource().getLocation().toURI().getPath();
455+
String jarPath = getClass().getProtectionDomain().getCodeSource().getLocation().toURI().getPath();
449456
try (JarFile jarFile = new JarFile(jarPath)) {
450457
Enumeration<JarEntry> entries = jarFile.entries();
451458
while (entries.hasMoreElements()) {
452459
JarEntry entry = entries.nextElement();
453460
if (!entry.getName().startsWith(resource + "/") || entry.isDirectory()) {
454461
continue;
455462
}
456-
Path outputPath =
457-
modelsDirectory.toPath().resolve(entry.getName().substring(resource.length() + 1));
463+
Path outputPath = modelsDirectory.toPath()
464+
.resolve(entry.getName().substring(resource.length() + 1));
458465

459466
// Check if the file already exists or if it is a supported model file
460467
if ((Files.exists(outputPath))
@@ -519,24 +526,22 @@ public File exportSingleModel(String modelPath) {
519526
return null;
520527
}
521528

522-
ModelProperties properties =
523-
ConfigManager.getInstance()
524-
.getConfig()
525-
.neuralNetworkPropertyManager()
526-
.getModel(Path.of(modelPath));
529+
ModelProperties properties = ConfigManager.getInstance()
530+
.getConfig()
531+
.neuralNetworkPropertyManager()
532+
.getModel(Path.of(modelPath));
527533

528534
String fileName = "";
529535
String suffix = modelFile.getName().substring(modelFile.getName().lastIndexOf('.'));
530536
if (properties != null) {
531-
fileName =
532-
String.format(
533-
"%s-%s-%s-%dx%d-%s",
534-
properties.nickname().replace(" ", ""),
535-
properties.family(),
536-
properties.version(),
537-
properties.resolutionWidth(),
538-
properties.resolutionHeight(),
539-
String.join("_", properties.labels()));
537+
fileName = String.format(
538+
"%s-%s-%s-%dx%d-%s",
539+
properties.nickname().replace(" ", ""),
540+
properties.family(),
541+
properties.version(),
542+
properties.resolutionWidth(),
543+
properties.resolutionHeight(),
544+
String.join("_", properties.labels()));
540545
} else {
541546
fileName = new File(modelPath).getName();
542547
}

0 commit comments

Comments
 (0)