update nsfw detection model to 1.2

ladaapp · Jan 14, 2025 · 51d603d · 51d603d
1 parent 09b623f
commit 51d603d
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/README.md b/README.md
@@ -179,8 +179,8 @@ To train the models and create your own datasets you'll also need these files
 wget -P model_weights/3rd_party/ 'https://download.openmmlab.com/mmediting/restorers/basicvsr/spynet_20210409-c6c1bd09.pth'
 wget -P model_weights/3rd_party/ 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth'
 wget -P model_weights/3rd_party/ 'https://github.com/QualityAssessment/DOVER/releases/download/v0.1.0/DOVER.pth'
-wget -P model_weights/ 'https://github.com/ladaapp/lada/releases/download/v0.5.0-beta/lada_nsfw_detection_model_v1.1.pt'
-wget -P model_weights/ 'https://github.com/ladaapp/lada/releases/download/v0.5.0-beta/lada_watermark_detection_model_v1.1.pt'
+wget -P model_weights/ 'https://github.com/ladaapp/lada/releases/download/v0.5.0-beta3/lada_nsfw_detection_model_v1.2.pt'
+wget -P model_weights/ 'https://github.com/ladaapp/lada/releases/download/v0.5.0-beta2/lada_watermark_detection_model_v1.1.pt'
 wget -P model_weights/3rd_party/ 'https://github.com/notAI-tech/NudeNet/releases/download/v3.4-weights/640m.pt'
 ```
 

diff --git a/create_mosaic_removal_video_dataset.py b/create_mosaic_removal_video_dataset.py
@@ -501,7 +501,7 @@ def parse_args():
                         help="Save as images instead of videos")
 
     nsfw_detection = parser.add_argument_group('NSFW detection')
-    nsfw_detection.add_argument('--model', type=str, default="model_weights/lada_nsfw_detection_model_v1.1.pt",
+    nsfw_detection.add_argument('--model', type=str, default="model_weights/lada_nsfw_detection_model_v1.2.pt",
                         help="path to NSFW detection model")
     nsfw_detection.add_argument('--model-device', type=str, default="cuda", help="device to run the YOLO model on. E.g. 'cuda' or 'cuda:0'")
 

diff --git a/lada/yolo/train-yolo-nsfw-detection.py b/lada/yolo/train-yolo-nsfw-detection.py
@@ -11,4 +11,4 @@
 args = parser.parse_args()
 
 model = YOLO('yolo11m-seg.yaml')
-model.train(data=args.config, epochs=200, imgsz=640, name="train_nsfw_detection_yolo11m")
+model.train(data=args.config, epochs=250, imgsz=640, name="train_nsfw_detection_yolo11m")