Skip to content

Commit

Permalink
Add tests for training
Browse files Browse the repository at this point in the history
  • Loading branch information
@abagshaw
abagshaw committed May 22, 2017
1 parent c49f518 commit 36bb03d
Show file tree
Hide file tree
Showing 6 changed files with 165 additions and 32 deletions.
3 changes: 3 additions & 0 deletions .gitignore
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,5 +39,8 @@ summary/
#Built graphs
built_graph/

#Training checkpoints
ckpt/*

#pytest cache
.cache/
106 changes: 74 additions & 32 deletions test/test_darkflow.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,26 +12,36 @@
# locally if you don't want this happening!)

#Settings
imgWidth = 640
imgHeight = 424
buildPath = os.environ.get("TRAVIS_BUILD_DIR")

if buildPath is None:
print()
print("TRAVIS_BUILD_DIR environment variable was not found - is this running on TravisCI?")
print("If you want to test this locally, set TRAVIS_BUILD_DIR to the base directory of the cloned darkflow repository.")
exit()
testImgPath = os.path.join(buildPath, "sample_img", "sample_person.jpg")
expectedDetectedObjectsV1 = [{"label": "dog","confidence": 0.46,"topleft": {"x": 84, "y": 249},"bottomright": {"x": 208,"y": 367}},
{"label": "person","confidence": 0.60,"topleft": {"x": 159, "y": 102},"bottomright": {"x": 304,"y": 365}}]

expectedDetectedObjectsV2 = [{"label":"person","confidence":0.82,"topleft":{"x":189,"y":96},"bottomright":{"x":271,"y":380}},
{"label":"dog","confidence":0.79,"topleft":{"x":69,"y":258},"bottomright":{"x":209,"y":354}},
{"label":"horse","confidence":0.89,"topleft":{"x":397,"y":127},"bottomright":{"x":605,"y":352}}]
testImg = {"path": os.path.join(buildPath, "sample_img", "sample_person.jpg"), "width": 640, "height": 424,
"expected-objects": {"yolo-small": [{"label": "dog", "confidence": 0.46, "topleft": {"x": 84, "y": 249}, "bottomright": {"x": 208, "y": 367}},
{"label": "person", "confidence": 0.60, "topleft": {"x": 159, "y": 102}, "bottomright": {"x": 304, "y": 365}}],
"yolo": [{"label": "person", "confidence": 0.82, "topleft": {"x": 189, "y": 96}, "bottomright": {"x": 271, "y": 380}},
{"label": "dog", "confidence": 0.79, "topleft": {"x": 69, "y": 258}, "bottomright": {"x": 209, "y": 354}},
{"label": "horse", "confidence": 0.89, "topleft": {"x": 397, "y": 127}, "bottomright": {"x": 605, "y": 352}}]}}

trainImgBikePerson = {"path": os.path.join(buildPath, "test", "training", "images", "1.jpg"), "width": 500, "height": 375,
"expected-objects": {"tiny-yolo-voc-TRAINED": [{"label": "bicycle", "confidence": 0.46, "topleft": {"x": 121, "y": 126}, "bottomright": {"x": 234, "y": 244}},
{"label": "cow", "confidence": 0.54, "topleft": {"x": 262, "y": 218}, "bottomright": {"x": 385, "y": 311}},
{"label": "person", "confidence": 0.70, "topleft": {"x": 132, "y": 34}, "bottomright": {"x": 232, "y": 167}}]}}

trainImgHorsePerson = {"path": os.path.join(buildPath, "test", "training", "images", "2.jpg"), "width": 500, "height": 332,
"expected-objects": {"tiny-yolo-voc-TRAINED": [{"label": "horse", "confidence": 0.97, "topleft": {"x": 157, "y": 95}, "bottomright": {"x": 420, "y": 304}},
{"label": "person", "confidence": 0.89, "topleft": {"x": 258, "y": 53}, "bottomright": {"x": 300, "y": 218}}]}}


posCompareThreshold = 0.05 #Comparisons must match be within 5% of width/height when compared to expected value
threshCompareThreshold = 0.1 #Comparisons must match within 0.1 of expected threshold for each prediction
yoloDownloadV1 = "https://pjreddie.com/media/files/yolo-small.weights"
yoloDownloadV2 = "https://pjreddie.com/media/files/yolo.weights"
yolo_small_Download = "https://pjreddie.com/media/files/yolo-small.weights" #YOLOv1
yolo_Download = "https://pjreddie.com/media/files/yolo.weights" #YOLOv2
tiny_yolo_voc_Download = "https://pjreddie.com/media/files/tiny-yolo-voc.weights" #YOLOv2

def download_file(url, savePath):
fileName = savePath.split("/")[-1]
Expand All @@ -47,19 +57,23 @@ def download_file(url, savePath):
else:
print("Found existing " + fileName + " file.")

yoloWeightPathV1 = os.path.join(buildPath, "bin", yoloDownloadV1.split("/")[-1])
yoloCfgPathV1 = os.path.join(buildPath, "cfg", "v1", "{0}.cfg".format(os.path.splitext(os.path.basename(yoloWeightPathV1))[0]))
yolo_small_WeightPath = os.path.join(buildPath, "bin", yolo_small_Download.split("/")[-1])
yolo_small_CfgPath = os.path.join(buildPath, "cfg", "v1", "{0}.cfg".format(os.path.splitext(os.path.basename(yolo_small_WeightPath))[0]))

yoloWeightPathV2 = os.path.join(buildPath, "bin", yoloDownloadV2.split("/")[-1])
yoloCfgPathV2 = os.path.join(buildPath, "cfg", "{0}.cfg".format(os.path.splitext(os.path.basename(yoloWeightPathV2))[0]))
yolo_WeightPath = os.path.join(buildPath, "bin", yolo_Download.split("/")[-1])
yolo_CfgPath = os.path.join(buildPath, "cfg", "{0}.cfg".format(os.path.splitext(os.path.basename(yolo_WeightPath))[0]))

pbPath = os.path.join(buildPath, "built_graph", os.path.splitext(os.path.basename(yoloWeightPathV2))[0] + ".pb")
metaPath = os.path.join(buildPath, "built_graph", os.path.splitext(os.path.basename(yoloWeightPathV2))[0] + ".meta")
tiny_yolo_voc_WeightPath = os.path.join(buildPath, "bin", tiny_yolo_voc_Download.split("/")[-1])
tiny_yolo_voc_CfgPath = os.path.join(buildPath, "cfg", "{0}.cfg".format(os.path.splitext(os.path.basename(tiny_yolo_voc_WeightPath))[0]))

pbPath = os.path.join(buildPath, "built_graph", os.path.splitext(os.path.basename(yolo_WeightPath))[0] + ".pb")
metaPath = os.path.join(buildPath, "built_graph", os.path.splitext(os.path.basename(yolo_WeightPath))[0] + ".meta")

generalConfigPath = os.path.join(buildPath, "cfg")

download_file(yoloDownloadV1, yoloWeightPathV1) #Check if we need to download (and if so download) the YOLOv1 weights
download_file(yoloDownloadV2, yoloWeightPathV2) #Check if we need to download (and if so download) the YOLOv2 weights
download_file(yolo_small_Download, yolo_small_WeightPath) #Check if we need to download (and if so download) the yolo-small weights (YOLOv1)
download_file(yolo_Download, yolo_WeightPath) #Check if we need to download (and if so download) the yolo weights (YOLOv2)
download_file(tiny_yolo_voc_Download, tiny_yolo_voc_WeightPath) #Check if we need to download (and if so download) the tiny-yolo-voc weights (YOLOv2)

def executeCLI(commandString):
print()
Expand Down Expand Up @@ -95,43 +109,45 @@ def compareObjectData(defaultObjects, newObjects, width, height):
return True

#Delete all images that won't be tested on so forwarding the whole folder doesn't take forever
filelist = [f for f in os.listdir(os.path.dirname(testImgPath)) if os.path.isfile(os.path.join(os.path.dirname(testImgPath), f)) and f != os.path.basename(testImgPath)]
filelist = [f for f in os.listdir(os.path.dirname(testImg["path"])) if os.path.isfile(os.path.join(os.path.dirname(testImg["path"]), f)) and f != os.path.basename(testImg["path"])]
for f in filelist:
os.remove(os.path.join(os.path.dirname(testImgPath), f))
os.remove(os.path.join(os.path.dirname(testImg["path"]), f))


#TESTS FOR INFERENCE
def test_CLI_IMG_YOLOv2():
#Test predictions outputted to an image using the YOLOv2 model through CLI
#NOTE: This test currently does not verify anything about the image created (i.e. proper labeling, proper positioning of prediction boxes, etc.)
# it simply verifies that the code executes properly and that the expected output image is indeed created in ./test/img/out

testString = "flow --imgdir {0} --model {1} --load {2} --config {3} --threshold 0.4".format(os.path.dirname(testImgPath), yoloCfgPathV2, yoloWeightPathV2, generalConfigPath)
testString = "flow --imgdir {0} --model {1} --load {2} --config {3} --threshold 0.4".format(os.path.dirname(testImg["path"]), yolo_CfgPath, yolo_WeightPath, generalConfigPath)
executeCLI(testString)

outputImgPath = os.path.join(os.path.dirname(testImgPath), "out", os.path.basename(testImgPath))
outputImgPath = os.path.join(os.path.dirname(testImg["path"]), "out", os.path.basename(testImg["path"]))
assert os.path.exists(outputImgPath), "Expected output image: {0} was not found.".format(outputImgPath)

def test_CLI_JSON_YOLOv2():
#Test predictions outputted to a JSON file using the YOLOv2 model through CLI
#NOTE: This test verifies that the code executes properly, the JSON file is created properly and the predictions generated are within a certain
# margin of error when compared to the expected predictions.

testString = "flow --imgdir {0} --model {1} --load {2} --config {3} --threshold 0.4 --json".format(os.path.dirname(testImgPath), yoloCfgPathV2, yoloWeightPathV2, generalConfigPath)
testString = "flow --imgdir {0} --model {1} --load {2} --config {3} --threshold 0.4 --json".format(os.path.dirname(testImg["path"]), yolo_CfgPath, yolo_WeightPath, generalConfigPath)
executeCLI(testString)

outputJSONPath = os.path.join(os.path.dirname(testImgPath), "out", os.path.splitext(os.path.basename(testImgPath))[0] + ".json")
outputJSONPath = os.path.join(os.path.dirname(testImg["path"]), "out", os.path.splitext(os.path.basename(testImg["path"]))[0] + ".json")
assert os.path.exists(outputJSONPath), "Expected output JSON file: {0} was not found.".format(outputJSONPath)

with open(outputJSONPath) as json_file:
loadedPredictions = json.load(json_file)

assert compareObjectData(expectedDetectedObjectsV2, loadedPredictions, imgWidth, imgHeight), "Generated object predictions from JSON were not within margin of error compared to expected values."
assert compareObjectData(testImg["expected-objects"]["yolo"], loadedPredictions, testImg["width"], testImg["height"]), "Generated object predictions from JSON were not within margin of error compared to expected values."

def test_CLI_SAVEPB_YOLOv2():
#Save .pb and .meta as generated from the YOLOv2 model through CLI
#NOTE: This test verifies that the code executes properly, and the .pb and .meta files are successfully created. A subsequent test will verify the
#NOTE: This test verifies that the code executes properly, and the .pb and .meta files are successfully created. The subsequent test will verify the
# contents of those files.

testString = "flow --model {0} --load {1} --config {2} --threshold 0.4 --savepb".format(yoloCfgPathV2, yoloWeightPathV2, generalConfigPath)
testString = "flow --model {0} --load {1} --config {2} --threshold 0.4 --savepb".format(yolo_CfgPath, yolo_WeightPath, generalConfigPath)

with pytest.raises(SystemExit):
executeCLI(testString)
Expand All @@ -146,18 +162,44 @@ def test_RETURNPREDICT_PBLOAD_YOLOv2():

options = {"pbLoad": pbPath, "metaLoad": metaPath, "threshold": 0.4}
tfnet = TFNet(options)
imgcv = cv2.imread(testImgPath)
imgcv = cv2.imread(testImg["path"])
loadedPredictions = tfnet.return_predict(imgcv)

assert compareObjectData(expectedDetectedObjectsV2, loadedPredictions, imgWidth, imgHeight), "Generated object predictions from return_predict() were not within margin of error compared to expected values."
assert compareObjectData(testImg["expected-objects"]["yolo"], loadedPredictions, testImg["width"], testImg["height"]), "Generated object predictions from return_predict() were not within margin of error compared to expected values."

def test_RETURNPREDICT_YOLOv1():
#Test YOLOv1 using normal .weights and .cfg
#NOTE: This test verifies that the code executes properly, and that the predictions generated are within the accepted margin of error to the expected predictions.

options = {"model": yoloCfgPathV1, "load": yoloWeightPathV1, "config": generalConfigPath, "threshold": 0.4}
options = {"model": yolo_small_CfgPath, "load": yolo_small_WeightPath, "config": generalConfigPath, "threshold": 0.4}
tfnet = TFNet(options)
imgcv = cv2.imread(testImg["path"])
loadedPredictions = tfnet.return_predict(imgcv)

assert compareObjectData(testImg["expected-objects"]["yolo-small"], loadedPredictions, testImg["width"], testImg["height"]), "Generated object predictions from return_predict() were not within margin of error compared to expected values."

#TESTS FOR TRAINING
def test_TRAIN_FROM_WEIGHTS_CLI__LOAD_CHECKPOINT_RETURNPREDICT_YOLOv2():
#Test training using pre-generated weights for tiny-yolo-voc
#NOTE: This test verifies that the code executes properly, and that the expected checkpoint file (tiny-yolo-voc-20.meta in this case) is generated.
# In addition, predictions are generated using the checkpoint file to verify that training completed successfully.

testString = "flow --model {0} --load {1} --train --dataset {2} --annotation {3} --epoch 20".format(tiny_yolo_voc_CfgPath, tiny_yolo_voc_WeightPath, os.path.join(buildPath, "test", "training", "images"), os.path.join(buildPath, "test", "training", "annotations"))
with pytest.raises(SystemExit):
executeCLI(testString)

checkpointPath = os.path.join(buildPath, "ckpt", "tiny-yolo-voc-20.meta")
assert os.path.exists(checkpointPath), "Expected output checkpoint file: {0} was not found.".format(checkpointPath)

options = {"model": tiny_yolo_voc_CfgPath, "load": 20, "config": generalConfigPath, "threshold": 0.4}
tfnet = TFNet(options)
imgcv = cv2.imread(testImgPath)

#Make sure predictions match the expected values for image with bike and person
imgcv = cv2.imread(trainImgBikePerson["path"])
loadedPredictions = tfnet.return_predict(imgcv)
assert compareObjectData(trainImgBikePerson["expected-objects"]["tiny-yolo-voc-TRAINED"], loadedPredictions, trainImgBikePerson["width"], trainImgBikePerson["height"]), "Generated object predictions from training were not within margin of error compared to expected values for the image with the bike and person.\nTraining may not have completed successfully."

assert compareObjectData(expectedDetectedObjectsV1, loadedPredictions, imgWidth, imgHeight), "Generated object predictions from return_predict() were not within margin of error compared to expected values."
#Make sure predictions match the expected values for image with horse and person
imgcv = cv2.imread(trainImgHorsePerson["path"])
loadedPredictions = tfnet.return_predict(imgcv)
assert compareObjectData(trainImgHorsePerson["expected-objects"]["tiny-yolo-voc-TRAINED"], loadedPredictions, trainImgHorsePerson["width"], trainImgHorsePerson["height"]), "Generated object predictions from training were not within margin of error compared to expected values for the image with the bike and person.\nTraining may not have completed successfully."
44 changes: 44 additions & 0 deletions test/training/annotations/1.xml
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
<annotation>
<folder>VOC2007</folder>
<filename>1.jpg</filename>
<source>
<database>The VOC2007 Database</database>
<annotation>PASCAL VOC2007</annotation>
<image>flickr</image>
<flickrid>336426776</flickrid>
</source>
<owner>
<flickrid>Elder Timothy Chaves</flickrid>
<name>Tim Chaves</name>
</owner>
<size>
<width>500</width>
<height>375</height>
<depth>3</depth>
</size>
<segmented>0</segmented>
<object>
<name>person</name>
<pose>Left</pose>
<truncated>0</truncated>
<difficult>0</difficult>
<bndbox>
<xmin>135</xmin>
<ymin>25</ymin>
<xmax>236</xmax>
<ymax>188</ymax>
</bndbox>
</object>
<object>
<name>bicycle</name>
<pose>Left</pose>
<truncated>0</truncated>
<difficult>0</difficult>
<bndbox>
<xmin>95</xmin>
<ymin>85</ymin>
<xmax>232</xmax>
<ymax>253</ymax>
</bndbox>
</object>
</annotation>
44 changes: 44 additions & 0 deletions test/training/annotations/2.xml
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
<annotation>
<folder>VOC2007</folder>
<filename>2.jpg</filename>
<source>
<database>The VOC2007 Database</database>
<annotation>PASCAL VOC2007</annotation>
<image>flickr</image>
<flickrid>329950741</flickrid>
</source>
<owner>
<flickrid>Lothar Lenz</flickrid>
<name>Lothar Lenz</name>
</owner>
<size>
<width>500</width>
<height>332</height>
<depth>3</depth>
</size>
<segmented>0</segmented>
<object>
<name>person</name>
<pose>Left</pose>
<truncated>0</truncated>
<difficult>0</difficult>
<bndbox>
<xmin>235</xmin>
<ymin>51</ymin>
<xmax>309</xmax>
<ymax>222</ymax>
</bndbox>
</object>
<object>
<name>horse</name>
<pose>Left</pose>
<truncated>0</truncated>
<difficult>0</difficult>
<bndbox>
<xmin>157</xmin>
<ymin>106</ymin>
<xmax>426</xmax>
<ymax>294</ymax>
</bndbox>
</object>
</annotation>
Binary file added test/training/images/1.jpg
View file
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Binary file added test/training/images/2.jpg
View file
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.

0 comments on commit 36bb03d

Please sign in to comment.