Compare commits

..
Author SHA1 Message Date
Dan Petro d1149115e8 Saving 2019-05-02 16:46:50 -07:00
Dan Petro bce00aa714 Saving progress 2019-04-30 10:30:05 -07:00
Dan Petro bb75967025 Saving progress.
Works now!
2019-04-30 09:56:08 -07:00
Dan Petro d3c4e7fc49 Saving progress. Reset this 2019-04-26 09:59:25 -07:00
24 changed files with 93 additions and 567 deletions

No files matched your search

-5
View File
@@ -113,8 +113,3 @@ venv.bak/
*.model
*.csv
*.txt
results.html
!requirements.txt
!requirements-gpu.txt
!tests/models/*
!tests/data/*
-47
View File
@@ -1,47 +0,0 @@
# Eyeballer
Give those screenshots of yours a quick eyeballing.
Eyeballer is meant for large-scope network penetration tests where you need to find "interesting" targets from a huge set of web-based hosts. Go ahead and use your favorite screenshotting tool like normal (EyeWitness or GoWitness) and then run them through Eyeballer to tell you what's likely to contain vulnerabilities, and what isn't.
## Setup
Download required packages on pip:
```
sudo pip3 install -r requirements.txt
```
Or if you want GPU support:
```
sudo pip3 install -r requirements-gpu.txt
```
## Predicting Labels
To eyeball some screenshots, just run the "predict" mode:
```
eyeballer.py --weights YOUR_WEIGHTS.h5 predict YOUR_FILE.png
```
Or for a whole directory of files:
```
eyeballer.py --weights YOUR_WEIGHTS.h5 predict PATH_TO/YOUR_FILES/
```
## Training
To train a new model, run:
```
eyeballer.py train
```
You'll want a machine with a good GPU for this to run in a reasonable amount of time. Setting that up is outside the scope of this readme, however.
This will output a new model file (weights.h5 by default).
## Evaluation
You just trained a new model, cool! Let's see how well it performs against some images it's never seen before, across a variety of metrics:
```
eyeballer.py --weights YOUR_WEIGHTS.h5 evaluate
```
+2
View File
@@ -0,0 +1,2 @@
from eyeballer.eyeballer import EyeballModel # noqa: F401
from eyeballer.visualization import OcclusionGenerator # noqa: F401
+9 -54
View File
@@ -1,11 +1,9 @@
#!/usr/bin/env python3
import click
import csv
from eyeballer.model import EyeballModel, DATA_LABELS
from eyeballer.visualization import HeatMap
from jinja2 import Template
from model import EyeballModel, DATA_LABELS
from visualization import HeatMap
@click.group(invoke_without_command=True)
@@ -38,15 +36,14 @@ def train(ctx, graphs, batchsize, epochs):
@cli.command()
@click.argument('screenshot')
@click.option('--heatmap', default=False, is_flag=True, help="Create a heatmap graphfor the prediction")
@click.option('--threshold', default=.5, type=float, help="Threshold confidence for labeling")
@click.pass_context
def predict(ctx, screenshot, heatmap, threshold):
def predict(ctx, screenshot, heatmap):
model = EyeballModel(**ctx.obj['model_kwargs'])
results = model.predict(screenshot)
if heatmap:
# Generate a heatmap
HeatMap(screenshot, model, threshold).generate()
HeatMap(screenshot, model).generate()
if not results:
print("Error: Input file does not exist")
@@ -58,60 +55,18 @@ def predict(ctx, screenshot, heatmap, threshold):
labelwriter = csv.DictWriter(csvfile, fieldnames=fieldnames)
labelwriter.writeheader()
labelwriter.writerows(results)
print("Output written to results.csv")
buildHTML(processResults(results, threshold))
print("HTML written to results.html")
def processResults(results, threshold):
'''Filter the initial results dictionary and reformat it for use in JS.
Keyword arguments:
results -- dictionary output from predict function
'''
jsResults = {}
for result in results:
positiveTags = []
for label, label_info in result.items():
if (label == 'filename'):
pass
elif label_info > threshold:
positiveTags.append(label)
jsResults[result['filename']] = positiveTags
return(jsResults)
def buildHTML(jsResults):
'''Build HTML around the JS Dictionary that is passed from processResults.
Keyword arguments:
jsResults -- dictionary output from processResults function
'''
html_output = ""
with open("prediction_output_template.html") as template_file:
template = Template(template_file.read())
html_output = template.render(jsResults=jsResults)
with open('results.html', 'w') as file:
file.write(html_output)
def pretty_print_evaluation(results):
"""Print a human-readable summary of the evaluation"""
# We use 4.2% to handle all the way from " 0.00%" (7chars) to "100.00%" (7chars)
for label in DATA_LABELS:
print("{} Precision Score: {:4.2%}".format(label, results[label]['precision']))
print("{} Recall Score: {:4.2%}".format(label, results[label]['recall']))
print("'None of the above' Precision: {:4.2%}".format(results['none_of_the_above_precision']))
print("'None of the above' Recall: {:4.2%}".format(results['none_of_the_above_recall']))
print("All or nothing Accuracy: {:4.2%}".format(results['all_or_nothing_accuracy']))
print("Overall Binary Accuracy: {:4.2%}".format(results['total_binary_accuracy']))
print("Top 10 worst predictions: {}".format(results['top_10_worst'][1]))
print(f"{label} Precision Score: {results[label]['precision']:4.2%}")
print(f"{label} Recall Score: {results[label]['recall']:4.2%}")
print(f"Overall Binary Accuracy: {results['total_binary_accuracy']:4.2%}")
print(f"All or nothing accuracy: {results['all_or_nothing_accuracy']:4.2%}")
print(f"Top 10 worst predictions: {results['top_10_worst'][1]}")
@cli.command()
+53 -112
View File
@@ -1,21 +1,23 @@
import os
import random
import sys
import progressbar
# Prevent Tkinter Dependency
import matplotlib
matplotlib.use('agg') # noqa: E402
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import Augmentor
import tensorflow as tf
from tensorflow.keras.applications.mobilenet import preprocess_input
from keras import regularizers
from keras.applications import MobileNet
from keras.applications.mobilenet import preprocess_input
from keras.callbacks import ModelCheckpoint
from keras.layers import Dropout, Dense
from keras.layers import GlobalAveragePooling2D
from keras.models import Model
from keras.optimizers import Adam
from keras.preprocessing import image
from keras.preprocessing.image import ImageDataGenerator
from sklearn.metrics import classification_report, accuracy_score, hamming_loss
from eyeballer.augmentation import EyeballerAugmentation
from augmentation import EyeballerAugmentation
DATA_LABELS = ["custom404", "login", "homepage", "oldlooking"]
@@ -27,11 +29,11 @@ class EyeballModel:
"""
graphs_directory = "graphs/"
checkpoint_file = "weights.h5"
image_dir = "images/"
image_dir = "images_224x224/"
image_width, image_height = 224, 224
input_shape = (image_width, image_height, 3)
def __init__(self, weights_file, print_summary=False, seed=None):
def __init__(self, print_summary=False, weights_file="weights.h5", seed=None):
"""Constructor for model class.
Keyword arguments:
@@ -39,24 +41,23 @@ class EyeballModel:
weights_file -- A filename for where to load the model's weights.
seed -- PRNG seed, useful for repeating a previous run and using the same data. Training/Validation split is determined randomly.
"""
# # Build the model
self.model = tf.keras.Sequential()
pretrained_layer = tf.keras.applications.mobilenet.MobileNet(weights='imagenet', include_top=False, input_shape=self.input_shape)
self.model.add(pretrained_layer)
self.model.add(tf.keras.layers.GlobalAveragePooling2D())
self.model.add(tf.keras.layers.Dense(256, activation="relu"))
self.model.add(tf.keras.layers.Dropout(0.5))
self.model.add(tf.keras.layers.Dense(128, activation="relu"))
self.model.add(tf.keras.layers.Dropout(0.2))
self.model.add(tf.keras.layers.Dense(len(DATA_LABELS), activation="sigmoid"))
# Build the model
base_model = MobileNet(weights='imagenet', include_top=False, input_shape=self.input_shape)
x = base_model.output
x = GlobalAveragePooling2D()(x)
x = Dense(256, activation="relu", name="HiddenLayer1")(x)
x = Dropout(0.5)(x)
x = Dense(128, activation="relu", name="HiddenLayer2", kernel_regularizer=regularizers.l2(0.01))(x)
x = Dropout(0.2)(x)
output_layer = Dense(len(DATA_LABELS), activation="sigmoid", name="OutputLayer")(x)
self.model = Model(inputs=base_model.input, outputs=output_layer)
self.model.compile(optimizer=tf.keras.optimizers.Adam(0.0005),
loss="binary_crossentropy",
metrics=["accuracy"])
# for layer in pretrained_layer.layers:
# for layer in base_model.layers:
# layer.trainable = False
adam = Adam(lr=0.0005)
self.model.compile(optimizer=adam, loss="binary_crossentropy", metrics=["accuracy"])
if print_summary:
print(self.model.summary())
@@ -72,21 +73,15 @@ class EyeballModel:
self.random_seed = True
self.seed = random.randint(0, 999999)
print("No seed set, ", end='')
print("using seed: {}".format(self.seed))
print(f"using seed: {self.seed}")
random.seed(self.seed)
self.training_labels = self.training_labels.sample(frac=1)
if weights_file is not None and os.path.isfile(weights_file):
try:
self.model.load_weights(weights_file)
except OSError:
print("ERROR: Unable to open weights file '{}'".foramt(weights_file))
sys.exit(-1)
self.model.load_weights(weights_file)
print("Loaded model from file.")
else:
if weights_file is not None:
raise FileNotFoundError
print("WARN: No model loaded from file. Generating random model")
print("No model loaded from file")
# Data augmentation
augmentor = Augmentor.Pipeline()
@@ -112,7 +107,7 @@ class EyeballModel:
"""
print("Training with seed: " + str(self.seed))
data_generator = tf.keras.preprocessing.image.ImageDataGenerator(
data_generator = ImageDataGenerator(
preprocessing_function=self.preprocess_training_function,
validation_split=0.2,
samplewise_center=True)
@@ -141,12 +136,8 @@ class EyeballModel:
class_mode="other")
# Model checkpoint - Saves model weights when validation accuracy improves
callbacks = [tf.keras.callbacks.ModelCheckpoint(self.checkpoint_file,
monitor='val_loss',
verbose=1,
save_best_only=True,
save_weights_only=True,
mode='min')]
callbacks = [ModelCheckpoint(self.checkpoint_file, monitor='val_loss', verbose=1, save_best_only=True,
save_weights_only=True, mode='min')]
history = self.model.fit_generator(
training_generator,
@@ -199,11 +190,11 @@ class EyeballModel:
result["oldlooking"] = prediction[0][3]
return result
def predict(self, path, threshold=0.5):
"""Predict the labels for a single file or directory of files
def predict(self, path):
"""Predict the labels for a single file
Keyword arguments:
path -- The path to the file(s) that we'll be evaluating.
path -- The path to the file that we'll be evaluating
"""
# Is this a single file, or a directory?
screenshots = []
@@ -213,24 +204,15 @@ class EyeballModel:
screenshots = os.listdir(path)
screenshots = [os.path.join(path, s) for s in screenshots]
else:
raise FileNotFoundError
return None
results = []
pbar = progressbar.ProgressBar()
for screenshot in pbar(screenshots):
for screenshot in screenshots:
# Load the image into memory
img = None
try:
img = tf.keras.preprocessing.image.load_img(screenshot, target_size=(self.image_width, self.image_height))
img = tf.keras.preprocessing.image.img_to_array(img)
img = np.expand_dims(img, axis=0)
img = preprocess_input(img)
except IsADirectoryError:
print("\nWARN: Skipping directory: ", screenshot)
continue
except OSError:
print("\nWARN: Skipping empty or corrupt file: ", screenshot)
continue
img = image.load_img(screenshot, target_size=(self.image_width, self.image_height))
img = image.img_to_array(img)
img = np.expand_dims(img, axis=0)
img = preprocess_input(img)
prediction = self.model.predict(img, batch_size=1)
result = dict()
@@ -240,6 +222,7 @@ class EyeballModel:
result["homepage"] = prediction[0][2]
result["oldlooking"] = prediction[0][3]
results.append(result)
return results
def evaluate(self, threshold=0.5):
@@ -249,7 +232,7 @@ class EyeballModel:
threshold -- Value between 0->1. The cutoff where the numerical prediction becomes boolean. Default: 0.5
"""
# Prepare the data
data_generator = tf.keras.preprocessing.image.ImageDataGenerator(
data_generator = ImageDataGenerator(
preprocessing_function=preprocess_input)
evaluation_generator = data_generator.flow_from_dataframe(
self.evaluation_labels,
@@ -264,7 +247,7 @@ class EyeballModel:
if not self.random_seed:
print("Using validation set...")
# Data augmentation
data_generator = tf.keras.preprocessing.image.ImageDataGenerator(
data_generator = ImageDataGenerator(
preprocessing_function=self.preprocess_training_function,
samplewise_center=True,
validation_split=0.2)
@@ -282,71 +265,29 @@ class EyeballModel:
else:
print("Using evaluation set...")
predictions = self.model.predict_generator(
evaluation_generator,
verbose=1,
steps=len(evaluation_generator))
predictions = self.model.predict_generator(evaluation_generator, verbose=1, steps=len(evaluation_generator))
self._save_prediction_histograms(predictions)
print("predictions: ", type(predictions))
predictions = predictions > threshold
ground_truth = self.evaluation_labels[DATA_LABELS].to_numpy()
filenames = self.evaluation_labels[["filename"]].to_numpy()
ground_truth = evaluation_generator.data
stats = classification_report(ground_truth, predictions, target_names=DATA_LABELS, output_dict=True)
stats["total_binary_accuracy"] = 1 - hamming_loss(ground_truth, predictions)
stats["all_or_nothing_accuracy"] = accuracy_score(ground_truth, predictions)
stats["top_10_best"] = self._top_images(filenames, ground_truth, predictions, best=True)
stats["top_10_worst"] = self._top_images(filenames, ground_truth, predictions, best=False)
stats["none_of_the_above_recall"] = self._none_of_the_above_recall(ground_truth, predictions)
stats["none_of_the_above_precision"] = self._none_of_the_above_precision(ground_truth, predictions)
stats["top_10_best"] = self._top_images(evaluation_generator, predictions, best=True)
stats["top_10_worst"] = self._top_images(evaluation_generator, predictions, best=False)
return stats
def _none_of_the_above_recall(self, labels, predictions):
"""Returns the recall score for the 'none of the above' images.
That means, all the images that don't have a category.
"""
total_count = 0
correct_count = 0
for item in zip(labels.astype(int), predictions.astype(int)):
# Is this a none of the above?
if not item[0].any():
total_count += 1
if not item[1].any():
correct_count += 1
if total_count == 0:
print("WARNING: None of the Above Recall is NaN")
return 0
return correct_count / total_count
def _none_of_the_above_precision(self, labels, predictions):
"""Returns the precision score for the 'none of the above' images.
That means, all the images that don't have a category.
"""
total_count = 0
correct_count = 0
for item in zip(labels.astype(int), predictions.astype(int)):
# Is this a none of the above prediction?
if not item[1].any():
total_count += 1
if not item[0].any():
correct_count += 1
if total_count == 0:
print("WARNING: None of the Above Precision is NaN")
return 0
return correct_count / total_count
def _top_images(self, filenames, ground_truth, predictions, top_k=10, best=False):
def _top_images(self, labels, predictions, top_k=10, best=False):
"""Collect top-k best or top-k worst predicted images
Keyword arguments:
ground_truth -- The correct labels
labels -- The keras generator that contains the filenames and data labels
predictions -- The numpy array of predictions
top_k -- Top k elements
best -- True/False. Calculate either the best or worst images
:Return -- Tuple of top-k indicies and top-k filenames
"""
true_labels = np.array(ground_truth).astype(float)
true_labels = np.array(labels.data).astype(float)
predictions = predictions.astype(float)
differences = np.abs(predictions - true_labels).sum(axis=1)
@@ -359,7 +300,7 @@ class EyeballModel:
indicies = np.flipud(indicies)
for i in indicies[:top_k]:
top_file_list.append(filenames[i][0])
top_file_list.append(labels.filenames[i])
return indicies[:top_k], top_file_list
+7 -8
View File
@@ -3,11 +3,11 @@ import math
from keras.preprocessing import image
from copy import deepcopy
from matplotlib import pyplot as plt
from eyeballer.model import DATA_LABELS
from model import DATA_LABELS
class HeatMap():
def __init__(self, screenshot_path, model, threshold=0.5, boxsize=28, step=7):
def __init__(self, screenshot_path, model, boxsize=28, step=7):
self.model = model
img = image.load_img(screenshot_path, target_size=(self.model.image_width, self.model.image_height))
@@ -20,7 +20,6 @@ class HeatMap():
self.y = 0
self.screenshot_path = screenshot_path
self.gamma = 3
self.threshold = threshold
def generate(self, output_file="heatmap.png"):
""" Make a single heatmap image and return it """
@@ -31,13 +30,13 @@ class HeatMap():
for label in DATA_LABELS:
boxsize = self.boxsize
# Ignore this label if it didn't predict positively (true label)
if results[label] > self.threshold:
if results[label] > 0.5:
worst_score = 1
while worst_score > self.threshold:
while worst_score > 0.5:
heatmap, worst_score = self._get_heatmap(label, boxsize)
if worst_score > self.threshold:
if worst_score > 0.5:
boxsize += 28
print("Didn't get a good image for {}. Trying again with a bigger boxsize: {}".format(label, boxsize))
print(f"Didn't get a good image for {label}. Trying again with a bigger boxsize: {boxsize}")
else:
heatmaps.append(heatmap)
labels.append(label)
@@ -60,7 +59,7 @@ class HeatMap():
subplots.imshow(heatmap, alpha=0.5, interpolation='none')
subplots.set_title(labels[0])
plt.savefig(output_file)
print("Heatmap image written to: {}".format(output_file))
print(f"Heatmap image written to: {output_file}")
def _get_heatmap(self, label, boxsize):
worst_score = 1
-189
View File
@@ -1,189 +0,0 @@
<!DOCTYPE HTML>
<meta charset="utf-8"/>
<html>
<title>Eyeballer Results</title>
<body>
<div id="btnsCont">
<button class="btn all" onclick="filterSelection('all')"> Show all</button>
<button class="btn" onclick="filterSelection('custom404')"> 404</button>
<button class="btn" onclick="filterSelection('login')"> Login</button>
<button class="btn" onclick="filterSelection('homepage')"> Home Page</button>
<button class="btn" onclick="filterSelection('oldlooking')"> Old Looking</button>
<select onchange="numOnPage=this.value; filterSelection([]);" style="float:right; margin:5px; margin-right:30px; height:30px; width:auto;">
<option value="25">25</option>
<option value="50">50</option>
<option value="100">100</option>
<option value="250">250</option>
</select>
</div>
<div id="btnsPage">
</div>
<div id="container">
</div>
<button onclick="topFunction()" id="topBtn" title="Go to top">Top</button>
<script>
var results = {{jsResults}};
var selectedTags = new Array();
var btnContainer = document.getElementById("btnsCont");
var categoryBtns = btnContainer.getElementsByClassName("btn");
var numOnPage = 100;
var indexesToDisplay = [];
filterSelection([]);
for (var i = 0; i < categoryBtns.length; i++) {
categoryBtns[i].addEventListener("click", function () {
if (this.className.indexOf("active") > -1) {
this.className = this.className.replace(" active", "");
}
else {
this.className += " active";
}
});
}
function filterSelection(c) {
indexesToDisplay = [];
var currentIndex = -1;
if (c == "all") location.reload();
if (selectedTags.includes(c)!=true) selectedTags=selectedTags.concat(c);
else selectedTags.splice(selectedTags.indexOf(c),1);
for (var filename in results){
currentIndex++;
var tags = results[filename];
var isMissingTag = false;
for (var i in selectedTags){
if (!tags.includes(selectedTags[i])){
isMissingTag = true;
break;
}
}
if (isMissingTag === false){
indexesToDisplay = indexesToDisplay.concat(currentIndex);
}
}
buildPageButtons(indexesToDisplay);
buildHTML(0, numOnPage);
}
function buildPageButtons(indexes){
numOfPages = Math.ceil(indexes.length/numOnPage);
var html = '';
for (var i = 0; i < numOfPages; i++){
html += ('<button class="btnpage" onclick="buildHTML(' + (i*numOnPage) + ', ' + ((i+1)*numOnPage) + '); activatePage(this);">' + (i+1) +'</button>');
}
document.getElementById('btnsPage').innerHTML = html;
}
function activatePage(pressedBtn){
var activeNow = document.getElementsByClassName('activePage')
for (var i = 0; i < activeNow.length; i++){
activeNow[i].className = activeNow[i].className.replace(" activePage", "");
}
pressedBtn.className += ' activePage';
}
function buildHTML(rangeMin, rangeMax){
var html = '';
if (rangeMax > indexesToDisplay.length) rangeMax=indexesToDisplay.length;
for (var i = rangeMin; i < rangeMax; i++){
index=indexesToDisplay[i];
filename = Object.keys(results)[index];
html += ('<div class="filterDiv"><a href=' + filename + ' target="_blank"> ' + results[filename].toString().replace(/,/g, ' ') + '</a><img class="res_pic" src="' + filename + '"></div>');
}
document.getElementById('container').innerHTML = html;
}
window.onscroll = function() {scrollFunction()};
function scrollFunction() {
if (document.body.scrollTop > 20 || document.documentElement.scrollTop > 20) {
document.getElementById("topBtn").style.display = "block";
} else {
document.getElementById("topBtn").style.display = "none";
}
}
function topFunction() {
document.body.scrollTop = 0; // For Safari
document.documentElement.scrollTop = 0; // For Chrome, Firefox, IE and Opera
}
</script>
<style>
#topBtn {
display: none; /* Hidden by default */
position: fixed; /* Fixed/sticky position */
bottom: 20px; /* Place the button at the bottom of the page */
right: 30px; /* Place the button 30px from the right */
z-index: 99; /* Make sure it does not overlap */
border: none; /* Remove borders */
outline: none; /* Remove outline */
background-color: #193b72; /* Set a background color */
color: white; /* Text color */
cursor: pointer; /* Add a mouse pointer on hover */
padding: 15px; /* Some padding */
border-radius: 10px; /* Rounded corners */
font-size: 18px; /* Increase font size */
font-weight:bold;
border:solid;
}
#topBtn:hover {
background-color: #555;
}
.container {
overflow: hidden;
}
#btnsPage {
margin:10px;
}
.btnpage{
width:40px;
height:30px;
background-color: #f1f1f1;
outline: none;
}
.res_pic{
width:100%;
height:auto;
}
a{
font-weight:600;
padding:3px;
text-decoration: none;
color:#ffffff;
font-size:16px;
font-family:Sans-serif;
}
.name{
float:left;
}
.filterDiv {
float: left;
background-color: #777;
color: #ffffff;
width: 22%;
border-style:solid;
border-color:#666;
text-align: center;
margin: 10px;
line-height: 30px;
display: block;
}
.btn {
border: none;
outline: none;
padding: 12px 16px;
margin-left:10px;
background-color: #f1f1f1;
cursor: pointer;
}
.btn:hover {
background-color: #ddd;
}
.active {
background-color: #666;
color: white;
}
.activePage{
border: solid;
border-color:#666666;
color: black;
}
.btn.all {
background-color: #193b72;
color: white;
}
</style>
</body>
</html>
-10
View File
@@ -1,10 +0,0 @@
Augmentor
click
matplotlib
numpy
pandas
pillow
sklearn
tensorflow-gpu==2.0.0-beta0
jinja2
progressbar2
-10
View File
@@ -1,10 +0,0 @@
Augmentor
click
matplotlib
numpy
pandas
pillow
sklearn
tensorflow==2.0.0-beta0
jinja2
progressbar2
+22
View File
@@ -0,0 +1,22 @@
from setuptools import find_packages, setup
setup(
name='eyeballer',
version='0.0.1',
packages=find_packages(exclude=["tests"]),
url='https://github.com/BishopFox/eyeballer',
license='GNU General Public License v3.0',
author='Daniel Petro',
author_email='',
description='',
install_requires=[
'keras',
'sklearn',
'numpy',
'pandas',
'matplotlib',
'pillow', 'click'
],
test_requires=[],
python_requires='>=3.6', # We use some python 3.6+ only features, like ordered dicts by default and f-strings
)
-1
View File
@@ -1 +0,0 @@
#TODO
-128
View File
@@ -1,128 +0,0 @@
import errno
import os
import unittest
import sys
from eyeballer.model import EyeballModel
class PredictTest(unittest.TestCase):
def setUp(self):
class DummyFile(object):
def write(self, x): pass
def flush(self): pass
sys.stdout = DummyFile()
weights_file = "tests/models/test_weights.h5"
if not os.path.isfile(weights_file):
print("Error: Symlink the latest weights file to " + weights_file)
raise FileNotFoundError(
errno.ENOENT,
os.strerror(errno.ENOENT),
weights_file)
model_kwargs = {
"weights_file": weights_file,
"print_summary": False,
"seed": None
}
self.model = EyeballModel(**model_kwargs)
def test_different_seed_predict(self):
model_kwargs = {
"weights_file": None,
"print_summary": False,
"seed": 12345678
}
same_seed_model = EyeballModel(**model_kwargs)
screenshot = "tests/data/404.png"
results_one = self.model.predict(screenshot)
results_two = same_seed_model.predict(screenshot)
self.assertNotEqual(results_one, results_two)
def test_same_seed_predict(self):
model_kwargs = {
"weights_file": None,
"print_summary": False,
"seed": 12345678
}
same_seed_model = EyeballModel(**model_kwargs)
screenshot = "tests/data/404.png"
results_one = same_seed_model.predict(screenshot)
results_two = same_seed_model.predict(screenshot)
self.assertEqual(results_one, results_two)
def test_predict_custom404(self):
screenshot = "tests/data/404.png"
results = self.model.predict(screenshot)[0]
self.assertGreater(results["custom404"], 0.5)
def test_predict_not_custom404(self):
screenshot = "tests/data/nothing.png"
results = self.model.predict(screenshot)[0]
self.assertLess(results["custom404"], 0.5)
def test_predict_login(self):
screenshot = "tests/data/login.png"
results = self.model.predict(screenshot)[0]
self.assertGreater(results["login"], 0.5)
def test_predict_not_login(self):
screenshot = "tests/data/nothing.png"
results = self.model.predict(screenshot)[0]
print(screenshot, results)
self.assertLess(results["login"], 0.5)
def test_predict_homepage(self):
screenshot = "tests/data/homepage.png"
results = self.model.predict(screenshot)[0]
self.assertGreater(results["homepage"], 0.5)
def test_predict_not_homepage(self):
screenshot = "tests/data/nothing.png"
results = self.model.predict(screenshot)[0]
self.assertLess(results["homepage"], 0.5)
def test_predict_oldlooking(self):
screenshot = "tests/data/old-looking.png"
results = self.model.predict(screenshot)[0]
self.assertGreater(results["oldlooking"], 0.5)
def test_predict_not_oldlooking(self):
screenshot = "tests/data/nothing.png"
results = self.model.predict(screenshot)[0]
self.assertLess(results["oldlooking"], 0.5)
def test_file_doesnt_exist(self):
screenshot = "tests/data/doesnotexist.png"
try:
self.model.predict(screenshot)[0]
self.fail("FileNotFoundError was expected but not found")
except FileNotFoundError:
pass
def test_folder(self):
screenshots = "tests/data/"
results = self.model.predict(screenshots)
self.assertEqual(len(results), 7)
def test_file_is_empty(self):
"""
We're just testing that it doesn't crash, basically
"""
screenshot = "tests/data/empty.png"
self.model.predict(screenshot)
def test_file_is_invalid(self):
"""
We're just testing that it doesn't crash, basically
"""
screenshot = "tests/data/invalid.png"
self.model.predict(screenshot)
-1
View File
@@ -1 +0,0 @@
#TODO
-1
View File
@@ -1 +0,0 @@
from tests.PredictTest import *
Binary file not shown.

Before

Width:  |  Height:  |  Size: 101 KiB

View File
Whitespace-only changes.
Binary file not shown.

Before

Width:  |  Height:  |  Size: 479 KiB

-1
View File
@@ -1 +0,0 @@
THIS TEXT IS NOT A VALID PNG
Binary file not shown.

Before

Width:  |  Height:  |  Size: 27 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 10 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 202 B

Binary file not shown.

Before

Width:  |  Height:  |  Size: 347 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 277 B

Binary file not shown.