Customize for screenshots, fix install, usage

README.md

@@ -10,3 +10,17 @@ https://screen-shot-classifier.herokuapp.com/
 # Model :
 Download the model from : https://drive.google.com/open?id=1k99ndVPuxI3kDGs6or2rSUWPC5LSjF-9
 
+# Try standalone prediction code.
+python screenshot_predict.py [files]
+
+# To train a new model, first populate files in two subdirectories each of directories named test_set and training_set,
+# You can pick random files for testing and move them to the test directory.  Specify the number to move as 10-20% of samples
+cd training_set/chats
+ls | shuf -n 20 | xargs -i mv {} ../../test_set/chats
+cd ../others
+ls | shuf -n 20 | xargs -i mv {} ../../test_set/others
+cd ../..
+
+# Then run the training code.
+# It may be necessart to adjust `steps_per_epoch` in the code to match #images / batch_size
+python screenshot_train.py

requirements.txt

@@ -8,6 +8,7 @@ numpy>=1.9.2
 scipy>=0.15.1
 scikit-learn>=0.18
 keras==2.2.4
-tensorflow>=1.15.2
+tensorflow>=1.15.2,<2.0
 h5py==2.7.1
 Pillow>=2.2.2
+matplotlib

screenshot_predict.py

@@ -0,0 +1,37 @@
+#!/usr/bin/env python
+"""
+Generate predictions and show scores for screenshot filenames listed on command line.
+
+TODO:
+ Use labels from training directores, don't assume a classification name
+ Provide option to show images (show=True).
+"""
+
+import os
+os.environ['KERAS_BACKEND'] = 'tensorflow'
+
+import sys
+import matplotlib.pyplot as plt
+from keras.models import load_model
+from keras.preprocessing.image import image
+import numpy as np
+import glob
+
+classifier=load_model("model.h5")
+
+def load_image(img_path, show=True):
+    img_original = image.load_img(img_path)
+    img = image.load_img(img_path, target_size=(48, 54))
+    img_tensor = image.img_to_array(img)                    # (height, width, channels)
+    img_tensor = np.expand_dims(img_tensor, axis=0)         # (1, height, width, channels), add a dimension because the model expects this shape: (batch_size, height, width, channels)
+    img_tensor /= 255.                                      # imshow expects values in the range [0, 1]
+    if show:
+        plt.imshow(img_original)
+        plt.axis('off')
+        plt.show()
+    return img_tensor
+
+for img_file in sys.argv[1:]:
+  new_image = load_image(img_file, show=False)
+  pred = classifier.predict(new_image)[0][0]
+  print(f'score={pred:.6f} {" map " if pred < 0.5 else "other"} {img_file}')

screenshot_train.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python
+"""
+Generate tensorflow binary classification model for images, as model.h5.
+Assume images are screenshots, and don't bother with augmentation.
+Train the images found in in two subdirectories of training_set and test_set.
+"""
+
+import os
+os.environ['KERAS_BACKEND'] = 'tensorflow'
+
+import math
+from keras.models import Sequential
+from keras.layers import Convolution2D
+from keras.layers import MaxPooling2D
+from keras.layers import Flatten
+from keras.layers import Dense
+
+classifier = Sequential()
+
+batch_size = 16
+train_len = 72
+test_len = 14
+
+classifier.add(Convolution2D(batch_size, 3, 3, input_shape = (48, 54, 3), activation = 'relu'))
+classifier.add(MaxPooling2D(pool_size = (2, 2)))
+
+classifier.add(Convolution2D(batch_size, 3, 3, activation = 'relu'))
+classifier.add(MaxPooling2D(pool_size = (2, 2)))
+
+classifier.add(Flatten())
+
+classifier.add(Dense(output_dim = 128, activation = 'relu'))
+classifier.add(Dense(output_dim = 1, activation = 'sigmoid'))
+
+classifier.compile(optimizer = 'adam', loss = 'binary_crossentropy', metrics = ['accuracy'])
+
+from keras.preprocessing.image import ImageDataGenerator
+
+train_datagen = ImageDataGenerator(rescale = 1./255)
+
+test_datagen = ImageDataGenerator(rescale=1./255)
+
+# Assume all originals are 1080p: 1920x1080 pixels, and scale by 40x horizontally, 20x vertically to 48x54
+training_set = train_datagen.flow_from_directory(
+        'training_set',
+        target_size=(48, 54),
+        save_to_dir="tmp_resized_images",
+        batch_size=batch_size,
+        class_mode='binary')
+
+test_set = test_datagen.flow_from_directory(
+        'test_set',
+        target_size=(48, 54),
+        batch_size=batch_size,
+        class_mode='binary')
+
+classifier.fit_generator(
+        training_set,
+        steps_per_epoch=math.ceil(train_len / batch_size),
+        epochs=10,
+        validation_data=test_set,
+        validation_steps=math.ceil(test_len / batch_size))
+
+classifier.save("model.h5")