File size: 3,953 Bytes
8bacc03 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | import tempfile
import unittest
from contextlib import redirect_stdout
from io import StringIO
from pathlib import Path
from summarize_logs import main, summarize_log
class SummarizeLogsTest(unittest.TestCase):
def test_summarize_log_extracts_metadata_best_and_final_metrics(self):
content = """\
[00:00:00.000] python train.py --dataset fundus --lb_domain 1 --lb_num 40 --save_name fundus_dm1_lb40_full --gpu 0 --model MedSAM
[00:00:00.001] Namespace(dataset='fundus', save_name='fundus_dm1_lb40_full', max_iterations=30000, num_eval_iter=500, domain_num=4, lb_domain=1, lb_num=40, model='MedSAM')
[00:10:00.000] iteration 500 : loss : 0.377457, sam_sup_loss : 0.100733, sam_unsup_loss : 0.257611, unet_sup_loss : 0.269560, unet_unsup_loss : 0.582658, cons_w : 0.007820, mask_ratio : 0.893574, sd:0.672295,0.863125,ud:0.487320,0.651552,d:0.629366,0.868507,s_m_r:0.827062,0.613741,0.523415
[00:10:00.001] sam_ulb_cup_dice:0.672295, sam_ulb_disc_dice:0.863125, unet_ulb_cup_dice:0.487320, unet_ulb_disc_dice:0.651552, ulb_cup_dice:0.629366, ulb_disc_dice:0.868507
[00:10:00.002] test unet model
[00:10:01.000] epoch 1 : loss : 0.000000
\tval_cup_dice: 0.595408, val_disc_dice: 0.803349,
[00:10:01.001] val_cup_best_dice: 0.595408 at 500 iter, val_disc_best_dice: 0.803349 at 500 iter, val_best_avg_dice: 0.699379 at 500 iter, cup_dice: 0.595408, disc_dice: 0.803349
[00:10:01.002] test sam model
[00:10:01.500] domain1 epoch 1 : loss : 0.000000
\tval_cup_dice: 0.900000, val_disc_dice: 0.950000,
[00:10:02.000] epoch 1 : loss : 0.000000
\tval_cup_dice: 0.826944, val_disc_dice: 0.929697,
[00:10:02.001] stu_val_cup_best_dice: 0.826944 at 500 iter, stu_val_disc_best_dice: 0.929697 at 500 iter, val_best_avg_dice: 0.878321 at 500 iter, cup_dice: 0.826944, disc_dice: 0.929697
"""
with tempfile.TemporaryDirectory() as tmpdir:
path = Path(tmpdir) / "log.txt"
path.write_text(content)
row = summarize_log(path, Path(tmpdir))
self.assertEqual(row["dataset"], "fundus")
self.assertEqual(row["save_name"], "fundus_dm1_lb40_full")
self.assertEqual(row["last_iteration"], 500)
self.assertAlmostEqual(row["unet_best_avg_dice"], 0.699379)
self.assertEqual(row["unet_best_avg_iter"], 500)
self.assertAlmostEqual(row["sam_best_avg_dice"], 0.878321)
self.assertEqual(row["sam_best_avg_iter"], 500)
self.assertAlmostEqual(row["unet_final_avg_dice"], 0.6993785)
self.assertAlmostEqual(row["sam_final_avg_dice"], 0.8783205)
self.assertAlmostEqual(row["last_mask_ratio"], 0.893574)
def test_main_writes_summary_to_output_file_without_printing_rows(self):
content = """\
[00:00:00.000] python train.py --dataset BUSI --lb_domain 1 --lb_num 40 --save_name BUSI_dm1_lb40_full --model MedSAM
[00:00:01.000] iteration 500 : loss : 0.377457, mask_ratio : 0.893574
[00:00:02.000] test unet model
[00:00:03.000] epoch 1 : loss : 0.000000
\tval_base_dice: 0.496102,
[00:00:04.000] val_best_avg_dice: 0.496102 at 500 iter, base_dice: 0.496102
"""
with tempfile.TemporaryDirectory() as tmpdir:
model_dir = Path(tmpdir) / "model"
log_path = model_dir / "BUSI" / "train" / "BUSI_dm1_lb40_full" / "log.txt"
output_path = Path(tmpdir) / "summary.json"
log_path.parent.mkdir(parents=True)
log_path.write_text(content)
stdout = StringIO()
with redirect_stdout(stdout):
exit_code = main([
"--model-dir",
str(model_dir),
"--format",
"json",
"--output",
str(output_path),
])
self.assertEqual(exit_code, 0)
self.assertEqual(stdout.getvalue(), "")
self.assertIn('"dataset": "BUSI"', output_path.read_text())
if __name__ == "__main__":
unittest.main()
|