Private Detector图片检测模型

联合创作 · 2023-09-25 22:50

Private Detector 是 Bumble 开源的图像检测模型,可以用于检测色情图片。内部版本已经进行了大量的重构,并作为一个完全开源的项目发布,以允许更广泛的社区使用和微调他们自己的 Private Detector 模型。你可以下载预训练的 SavedModel 模型。



模型:


SavedModel 模型可以在 private_detector.zip 中的 save_model/ 目录中找到。该模型基于 Efficientnet-v2,并在 Bumble 内部的数据集上进行了训练。


推理:


推理是非常简单的,在 inference.py 中已经给出了一个例子



python3 inference.py \
--model saved_model/ \
--image_paths \
Yes_samples/1.jpg \
Yes_samples/2.jpg \
Yes_samples/3.jpg \
Yes_samples/4.jpg \
Yes_samples/5.jpg \
No_samples/1.jpg \
No_samples/2.jpg \
No_samples/3.jpg \
No_samples/4.jpg \
No_samples/5.jpg \

附加训练


开发者可以根据自己的数据自行微调模型,这样做非常简单,可以在 private_detector.zipsaved_checkpoint/目录中找到所需要的 checkpoint 文件。


设置一个 JSON 文件,其中包含指向每个类的图像路径列表的链接:




{
"Yes": {
"path": "/home/sofarrell/private_detector/Yes.txt",
"label": 0
},
"No": {
"path": "/home/sofarrell/private_detector/No.txt",
"label": 1
}
}


与每个列出图像的图像路径的.txt文件




/home/sofarrell/private_detector_images/Yes/1093840880_309463828.jpg
/home/sofarrell/private_detector_images/Yes/657954182_3459624.jpg
/home/sofarrell/private_detector_images/Yes/1503714421_3048734.jpg


可以使用 conda 创建训练环境:




conda env create -f environment.yaml
conda activate private_detector


然后像这样重新训练:




python3 ./train.py \
--train_json /home/sofarrell/private_detector/train_classes.json \
--eval_json /home/sofarrell/private_detector/eval_classes.json \
--checkpoint_dir saved_checkpoint/ \
--train_id retrained_private_detector

浏览 33
点赞
评论
收藏
分享

手机扫一扫分享

编辑 分享
举报
评论
图片
表情
推荐
点赞
评论
收藏
分享

手机扫一扫分享

编辑 分享
举报