s473558/si23traktor
2
2
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

make code more flexible

Browse Source
This commit is contained in:
s473558 2023-06-05 04:42:53 +02:00
parent 28bf53c037
commit e27acbacaf
1 changed files with 43 additions and 54 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
neural_network/inference.py
View File

@ -10,6 +10,7 @@ parser.add_argument('-i', '--input',
help='path to the input image') help='path to the input image')
args = vars(parser.parse_args()) args = vars(parser.parse_args())
def main(path):
# the computation device # the computation device
device = ('cuda' if torch.cuda.is_available() else 'cpu') device = ('cuda' if torch.cuda.is_available() else 'cpu')
# list containing all the class labels # list containing all the class labels
@ -38,9 +39,9 @@ transform = transforms.Compose([
# read and preprocess the image # read and preprocess the image
image = cv2.imread(args['input']) image = cv2.imread(path)
# get the ground truth class # get the ground truth class
gt_class = args['input'].split('/')[-2] gt_class = path.split('/')[-2]
orig_image = image.copy() orig_image = image.copy()
# convert to RGB format # convert to RGB format
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
@ -51,20 +52,8 @@ with torch.no_grad():
outputs = model(image.to(device)) outputs = model(image.to(device))
output_label = torch.topk(outputs, 1) output_label = torch.topk(outputs, 1)
pred_class = labels[int(output_label.indices)] pred_class = labels[int(output_label.indices)]
cv2.putText(orig_image,
f"GT: {gt_class}", return pred_class
(10, 25),
cv2.FONT_HERSHEY_SIMPLEX, if __name__ == "__main__":
0.6, (0, 255, 0), 2, cv2.LINE_AA main(args['input'])
)
cv2.putText(orig_image,
f"Pred: {pred_class}",
(10, 55),
cv2.FONT_HERSHEY_SIMPLEX,
0.6, (0, 0, 255), 2, cv2.LINE_AA
)
print(f"GT: {gt_class}, pred: {pred_class}")
cv2.imshow('Result', orig_image)
cv2.waitKey(0)
cv2.imwrite(f"outputs/{gt_class}{args['input'].split('/')[-1].split('.')[0]}.png",
orig_image)