安全帽检测SDK
安全帽检测。
- 支持类别:
- safe
- unsafe
SDK功能
- 安全帽检测,给出检测框和置信度
- 三个模型:
- 小模型(mobilenet0.25)
- 中模型(mobilenet1.0)
- 大模型(darknet53)
运行小模型例子 - SmallSafetyHelmetDetectExample
-
测试图片
运行中模型例子 - MediumSafetyHelmetDetectExample
-
测试图片
运行大模型例子 - LargeSafetyHelmetDetectExample
-
测试图片
/**
* 安全帽检测例子
*
* 目录:http://aias.top/
*
* @author Calvin
*/
public final class Yolov5sExample {
private static final Logger logger = LoggerFactory.getLogger(Yolov5sExample.class);
private Yolov5sExample() {}
public static void main(String[] args) throws IOException, ModelException, TranslateException {
Path imageFile = Paths.get("src/test/resources/demo.jpg");
Image image = ImageFactory.getInstance().fromFile(imageFile);
Criteria<Image, DetectedObjects> criteria = new Yolov5sDetect().criteria();
try (ZooModel model = ModelZoo.loadModel(criteria);
Predictor<Image, DetectedObjects> predictor = model.newPredictor()) {
DetectedObjects detections = predictor.predict(image);
List<DetectedObjects.DetectedObject> items = detections.items();
List<String> names = new ArrayList<>();
List<Double> prob = new ArrayList<>();
List<BoundingBox> boxes = new ArrayList<>();
for (int i = 0; i < items.size(); i++) {
DetectedObjects.DetectedObject item = items.get(i);
if (item.getProbability() < 0.5f) {
continue;
}
if (item.getClassName().equals("person")) {
continue;
}
names.add(item.getClassName());
prob.add(item.getProbability());
boxes.add(item.getBoundingBox());
}
detections = new DetectedObjects(names, prob, boxes);
ImageUtils.saveBoundingBoxImage(image, detections, "helmet_head_person_s.png", "build/output");
logger.info("{}", detections);
}
}
}
运行成功后,命令行应该看到下面的信息:
[INFO ] - [
class: "safe 0.9983590245246887", probability: 0.99835, bounds: [x=0.244, y=0.000, width=0.086, height=0.150]
class: "unsafe 0.998088538646698", probability: 0.99808, bounds: [x=0.226, y=0.204, width=0.115, height=0.263]
class: "safe 0.997364342212677", probability: 0.99736, bounds: [x=0.584, y=0.247, width=0.162, height=0.302]
class: "safe 0.9963852167129517", probability: 0.99638, bounds: [x=0.319, y=0.000, width=0.076, height=0.133]
class: "safe 0.9952006936073303", probability: 0.99520, bounds: [x=0.757, y=0.262, width=0.111, height=0.264]
]
目录:
Git地址:
https://github.com/mymagicpower/AIAS
https://gitee.com/mymagicpower/AIAS