Merge branch 'master' of https://github.com/deeppomf/DeepMindBreak

README.md

@@ -70,6 +70,10 @@ Decensored images will be saved to the "decensor_output_images" directory. Paste
 
 ## II. Train the pretrained model on custom dataset
 
+You must have a GPU for training since training on a CPU will take weeks.
+
+Your custom dataset should be 128 x 128 images of uncensored vaginas and penises cropped from hentai. The more images, the better: I used 70,000 images for training. Censoring these images yourself is unnecessary.
+
 Put your custom dataset for training in the "data/images" directory and convert images to npy format.
 
 ```
@@ -77,12 +81,14 @@ $ cd training_data
 $ python to_npy.py
 ```
 
-Train pretrained model on your custom dataset.
+To train, run
 
 ```
 $ python train.py
 ```
 
+If desired, you can train the pretrained model on your custom dataset.
+
 Training can be done separately for mosaics with train_mosaic.py, but decensor.py is not yet compatible with mosaic decensorship models.
 
 # To do

train.py

@@ -137,7 +137,7 @@ def get_points():
         	#rotate random amount between 0 and 90 degrees
         	m = scipy.ndimage.rotate(m, np.random.random()*90, reshape = False)
         	#set all elements greater than 0 to 1
-        	m[m > 0] = 1
+        	m[m > 0.5] = 1
 
         mask.append(m)