mirror of
https://github.com/PhotonVision/photonvision
synced 2026-06-19 00:41:41 +00:00
Exclude license from being loaded as a model (#2063)
This commit is contained in:
@@ -226,8 +226,18 @@ public class NeuralNetworkModelManager {
|
||||
private static final Logger logger = new Logger(NeuralNetworkModelManager.class, LogGroup.Config);
|
||||
|
||||
public enum Family {
|
||||
RKNN,
|
||||
RUBIK
|
||||
RKNN(".rknn"),
|
||||
RUBIK(".tflite");
|
||||
|
||||
private final String fileExtension;
|
||||
|
||||
private Family(String fileExtension) {
|
||||
this.fileExtension = fileExtension;
|
||||
}
|
||||
|
||||
public String extension() {
|
||||
return fileExtension;
|
||||
}
|
||||
}
|
||||
|
||||
public enum Version {
|
||||
@@ -358,7 +368,11 @@ public class NeuralNetworkModelManager {
|
||||
try {
|
||||
Files.walk(modelsDirectory.toPath())
|
||||
.filter(Files::isRegularFile)
|
||||
.forEach(path -> loadModel(path));
|
||||
.filter(
|
||||
path ->
|
||||
supportedBackends.stream()
|
||||
.anyMatch(family -> path.toString().endsWith(family.extension())))
|
||||
.forEach(this::loadModel);
|
||||
} catch (IOException e) {
|
||||
logger.error("Failed to discover models at " + modelsDirectory.getAbsolutePath(), e);
|
||||
}
|
||||
|
||||
@@ -605,16 +605,13 @@ public class RequestHandler {
|
||||
return;
|
||||
}
|
||||
|
||||
String modelFileExtension = "";
|
||||
NeuralNetworkModelManager.Family family;
|
||||
|
||||
switch (Platform.getCurrentPlatform()) {
|
||||
case LINUX_QCS6490:
|
||||
modelFileExtension = "tflite";
|
||||
family = NeuralNetworkModelManager.Family.RUBIK;
|
||||
break;
|
||||
case LINUX_RK3588_64:
|
||||
modelFileExtension = "rknn";
|
||||
family = NeuralNetworkModelManager.Family.RKNN;
|
||||
break;
|
||||
default:
|
||||
@@ -625,19 +622,19 @@ public class RequestHandler {
|
||||
}
|
||||
|
||||
// If adding additional platforms, check platform matches
|
||||
if (!modelFile.extension().contains(modelFileExtension)) {
|
||||
if (!modelFile.extension().contains(family.extension())) {
|
||||
ctx.status(400);
|
||||
ctx.result(
|
||||
"The uploaded file was not of type '"
|
||||
+ modelFileExtension
|
||||
+ family.extension()
|
||||
+ "'. The uploaded file should be a ."
|
||||
+ modelFileExtension
|
||||
+ family.extension()
|
||||
+ " file.");
|
||||
logger.error(
|
||||
"The uploaded file was not of type '"
|
||||
+ modelFileExtension
|
||||
+ family.extension()
|
||||
+ "'. The uploaded file should be a ."
|
||||
+ modelFileExtension
|
||||
+ family.extension()
|
||||
+ " file.");
|
||||
return;
|
||||
}
|
||||
@@ -665,7 +662,7 @@ public class RequestHandler {
|
||||
.addModelProperties(
|
||||
new ModelProperties(
|
||||
modelPath,
|
||||
modelFile.filename().replaceAll("." + modelFileExtension, ""),
|
||||
modelFile.filename().replaceAll("." + family.extension(), ""),
|
||||
labels,
|
||||
width,
|
||||
height,
|
||||
|
||||
Reference in New Issue
Block a user