Skip to content

Commit

Permalink
Added testing for directory with images
Browse files Browse the repository at this point in the history
  • Loading branch information
jaredmejia committed Apr 16, 2021
1 parent e354cc6 commit 67c0a27
Showing 1 changed file with 92 additions and 12 deletions.
104 changes: 92 additions & 12 deletions network/object_detection/test_model.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import cv2
import distutils.util
import numpy as np
import os
import PIL
Expand All @@ -18,18 +19,26 @@
import references.transforms as T

from box_hndlr import iso_package, get_box, draw_box
from bbox import PackageDataset, get_transform
from bbox import PackageDataset, get_transform, show_pred_box

def main():
model_name = sys.argv[1] # should be .pt file in ./
test_im_dir = sys.argv[2] # assume test_im_dir/Normal and test_im_dir/Segmentation
package_dir = bool(distutils.util.strtobool(sys.argv[3])) if len(sys.argv) > 3 else True
output_iou = bool(distutils.util.strtobool(sy.argv[4])) if len(sys.argv) > 4 else False

"""
Planning to split test data into includes box and doesn't include package before evaluation
As is, the metrics are thrown off by the loss of predictions where package isn't in the frame
since we currently require the target to include a bbox even if there is no package
-- evaluate normally for sub test set with packages in image
-- let accuracy be % of predictions in which no package was correctly predicted for sub test set without packages
If directory contains images of packages (package_dir == True), then new directory test_im_dir/NormalEval
is made, containing the original images with the bounding box drawn on.
If the directory contains images without packages (package_dir == False), then the new directory
test_im_dir/SegmentationEval is made. This new directory contains images with the bounding box drawn on,
only for the images in which a bounding box was predicted.
For both new directories, the name of the images in the new directory is the name of the original image
with 'eval' appended. Additionally, the new directories will contain a txt file containing information
about the box corner coordinates, the box center coordinates, the area of the box, and the confidence
level of the prediction for all the images with predictions made, as well as a summary of the evaluation.
"""

# loading model
Expand All @@ -38,13 +47,84 @@ def main():

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

# creating dataloader
test_dataset = PackageDataset(test_im_dir, get_transform(train=False))
data_loader = torch.utils.data.DataLoader(
test_dataset, batch_size=2, shuffle=False, num_workers=4,
collate_fn=utils.collate_fn)
assert os.path.isdir(os.path.join(test_im_dir, 'Normal'))

os.mkdir(os.path.join(test_im_dir, 'NormalEval'))
os.mkdir(os.path.join(test_im_dir, 'NormalEval', 'images'))

# images with package
if package_dir:

total_area = 0
total_conf = 0
# total_preds = len(os.listdir(os.path.join(test_im_dir, 'Normal')))
total_preds = 0
eval_data = {}

for f in sorted(os.listdir(os.path.join(test_im_dir, 'Normal'))):
f_im_path = os.path.join(test_im_dir, 'Normal', f)
f_im = PIL.Image.open(f_im_path).convert("RGB")
f_tensor = transforms.ToTensor()(f_im)

with torch.no_grad():
# making prediction
mask_pred = loaded_model([f_tensor.to(device)])

# package detected
try:
box_pred = mask_pred[0]['boxes'][0]
conf_score = mask_pred[0]['scores'][0]

# saving new image
f_im_eval = show_pred_box(f_im_path, box_pred)
new_f_name = f[:-4] + "_eval.png"
f_im_eval.save(os.path.join(test_im_dir, 'NormalEval', 'images', new_f_name))

# recording statistics
x0, y0, x1, y1 = [int(i) for i in box_pred.tolist()]
center_coords = ((x0 + x1) / 2, (y0 + y1) / 2)
area = (x1 - x0) * (y1 - y0)

stat_dict = {"box_coords": ((x0, y0), (x1, y1)),
"center_coords": center_coords,
"area": area,
"confidence": conf_score}

total_area += area
total_conf += conf_score
total_preds += 1
eval_data[new_f_name] = stat_dict

# package not detected
except IndexError:
pass

# writing statistics
avg_area = total_area / total_preds
avg_conf = conf_score / total_preds
with open(os.path.join(test_im_dir, 'NormalEval', 'preds.txt'), 'w') as out_file:
out_file.write("Summary\n")
out_file.write(f"\tTotal packages detected: {total_preds}\tAverage bbox area: {avg_area}\tAverage prediction confidence: {avg_conf}\n")
out_file.write(f"box_coords,center_coords,area,confidence")
for f in eval_data.keys():
out_file.write(f"{f["box_coords"]},{f["center_coords"]},{f["area"]},{f["confidence"]}")


# images without package
else:

pass


if output_iou:
# creating dataloader
test_dataset = PackageDataset(test_im_dir, get_transform(train=False))
data_loader = torch.utils.data.DataLoader(
test_dataset, batch_size=2, shuffle=False, num_workers=4,
collate_fn=utils.collate_fn)

evaluate(loaded_model, data_loader, device=device)
# outputing IoU metrics
evaluate(loaded_model, data_loader, device=device)


if __name__ == "__main__":
Expand Down

0 comments on commit 67c0a27

Please sign in to comment.