mergy merge

coxapp · Jan 21, 2019 · e84ce41 · e84ce41
1 parent c595905
commit e84ce41
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 1 deletion.
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1 @@
+weights.best.hdf5
diff --git a/README.md b/README.md
@@ -1 +1,12 @@
 # nsfw_model
+
+Trained on 100+ Gigs of NSFW data to identify:
+- `porn` - pornographic images
+- `hentai` - hentai images, but also includes pornographic drawings
+- `sexy` - sexually explicit images, but not pornography
+- `neutral` - safe for work neutral images
+- `drawings` - safe for work drawings (including anime)
+
+
+### Thanks
+Special thanks to the [nsfw_data_scraper](https://github.com/GantMan/nsfw_data_scrapper).
diff --git a/train_model.py b/train_model.py
@@ -42,6 +42,10 @@
     Dense(5, activation='softmax')
 ])
 
+if os.path.exists("weights.best.hdf5"):
+        print ("loading ", "weights.best.hdf5")
+        model.load_weights('weights.best.hdf5')
+
 # checkpoint
 filepath = "weights.best.hdf5"
 checkpoint = ModelCheckpoint(
@@ -86,7 +90,7 @@
 
 validation_generator = validation_datagen.flow_from_directory(
     test_dir,
-    target_size=(256, 512562),
+    target_size=(256, 256),
     class_mode='categorical',
     batch_size=GENERATOR_BATCH_SIZE
 )