Update visualization.py

visualization.py

@@ -66,9 +66,9 @@ def print_top_classes(original_image, **kwargs):
         if len(CLS2IDX[cls_idx]) > max_str_len:
             max_str_len = len(CLS2IDX[cls_idx])
 
-    print('Top 5 classes:')
     for cls_idx in class_indices:
-        output_string = '\t{} : {}'.format(cls_idx, CLS2IDX[cls_idx])
-        output_string += ' ' * (max_str_len - len(CLS2IDX[cls_idx])) + '\t\t'
-        output_string += 'value = {:.3f}\t prob = {:.1f}%'.format(predictions[0, cls_idx], 100 * prob[0, cls_idx])
-        print(output_string)
+        output_string = '{} : {}'.format(cls_idx, CLS2IDX[cls_idx])
+        # output_string += ' ' * (max_str_len - len(CLS2IDX[cls_idx])) + '\t\t'
+        output_string += 'value = {:.3f}  prob = {:.1f}%'.format(predictions[0, cls_idx], 100 * prob[0, cls_idx])
+        output.append(output_string)
+    return output