diff --git a/exit b/exit new file mode 100644 index 0000000..f468be3 --- /dev/null +++ b/exit @@ -0,0 +1,626 @@ +07e5431 (HEAD -> main, origin/main, origin/HEAD) changes to domain_adaptation.py +606214c implement rescaling for get_unsupervised_loader and mean_teacher_adaptation +fdc6b46 added new mito model and removed old mito models 2 and 3 - new model is called mitochondria2 (#156) +9a7a089 (tag: 0.4.1) Fix export of touching objects to IMOD (#155) +3459903 (tag: 0.4.0) Support checkpoint initialization for supervised training; fix issues… (#154) +8cb080c Fix supervised-training CLI and add test (#153) +21f9850 Add prepocess to pred func (#149) +e362fe0 added mitov3 model - average voxel_size in nm: 2.87 (#147) +3c31509 Change preprint link to publication link +8df45a1 Merge pull request #141 from computational-cell-analytics/140-update-voxel-size-for-mitochondria2-model +effdb93 Update inference.py +46b6300 Update pool visualization +8d7ee71 Merge pull request #139 from computational-cell-analytics/model-fixes +ba1e5e7 Fix mitochondria2 model and add proper model test +4dc4b21 Update vesicle pool visualization +fc7ca4d Merge pull request #138 from computational-cell-analytics/convenience +3b0fad1 Fix typo +6e3db6e Update size filter plugin to remove disconnected components +7046235 Add cli to visualize vesicle pool assignments +457d558 Save the pool color in the exported table +d0f284a match z alignment to AZ seg +c028aa9 Merge pull request #137 from computational-cell-analytics/2d_imod_export +acc3908 (origin/2d_imod_export) z plane +71ec138 testing +664f203 Implement size filtering widget +58d54c9 Merge branch 'main' of https://github.com/computational-cell-analytics/synapse-net into 2d_imod_export +10cbfd5 Update segmentation_widget.py +e8359ee Merge pull request #136 from computational-cell-analytics/minor-updates +ee8a0ce Fix bug in new membrane post-processing +efb5851 Update ribbon structure post-processing in segmentation GUI +eb2d3ed Support setting multiple devices in segmentation +d344d49 2D vesicle export to imod +9c252ed Implement scalable segmentation (#134) +43eff47 boundary mask for unsupervised training (#132) +bfccbf0 wrong format extension +e1996bd (tag: 0.3.0) Merge pull request #133 from computational-cell-analytics/constantinpape-patch-2 +0beffc8 Update __version__.py +1e25907 Merge pull request #131 from computational-cell-analytics/revision3 +0ec2a50 Update the voxel size for the new AZ model +75d60f9 Update inference CLI so that it works with checkpoints +f76dbf7 Add attempt at compartment evaluation code +d8d7c1b Update AZ model +928f330 Merge pull request #130 from computational-cell-analytics/revision2 +1f3ad5f Update CLI training info and add community submission info to doc +480d714 Update the SynapseNet trainign CLI +be0917a Fix issues in training CLI and add domain adaptation CLI +70628f6 Implement CLI for supervised training +0a8101e Update vesicle inference +408b807 Merge pull request #121 from computational-cell-analytics/revision +7967b4c Update environment.yaml +8ea5b66 more conservative SV filtering at z borders +fbe4a55 minor az postprocessing for data analysis +f516ae4 minor adjustment +34886d7 fixed a few mistakes +6e95660 make code more flexible +8763e49 make usage of surface dice more general +0c30cf0 updated SV seg h5 +48618d0 added compartment seg saved in h5 +f17c349 option to calc surface dice per component +ad03103 add dataset info +d966ea3 minor things for analysis; 1st implementation of surface dice for eval +7b83139 change presynaptic filtering +cab9f2b Merge branch 'main' of https://github.com/computational-cell-analytics/synapse-net into revision +e8788e5 Delete analysis_results directory +d6eaa66 add boundary threshold for compartment seg +1ae7401 Merge branch 'main' into revision +f465dbe Merge pull request #128 from computational-cell-analytics/comp_seg +78825be (origin/comp_seg) add option to change boundary threshold for compartment seg +db3e654 small change for compartment segmentation +a9b714f minor addition +ac4b9d7 store segID, use existing seg +17b9368 remove background from analysis +ba09eca Update to_imod.py (#127) +d7b9a1a exclude SV at boundary for automatic analysis +2fa3925 Merge branch 'main' of https://github.com/computational-cell-analytics/synapse-net into revision +61bd683 Update util.py (#126) +0785657 fix import +b842bca new data analysis, not tested +6525651 fix a few things +b1c8feb make training more usable +02da6a2 Merge branch 'main' of https://github.com/computational-cell-analytics/synapse-net into revision +59c3534 prepare for training AZ +d2aa6d8 empty crop without AZ annotation +b769a8e make az_thin more flexible +c41628c make prediction more flexible +5f66328 cropping stem data +f8d801f Clean up AZ scripts and summarize current state +d7f71a0 Update AZ training data +39d69d5 Implement AZ evaluation WIP +7dd6962 Work on revision of AZ model WIP +2f4564e Update start_page.md (#123) +ea5bbce Minor updates to documentation (#122) +f08ad89 Add SBD to CryoVesNet evaluation +d23d544 Add code for tomogram on-the-fly processing (#117) +964d6c3 Update environment.yaml (#120) +71a0343 Add mask intersection post-processing (#115) +ba00244 Add more outputs to debug performance (#108) +e2f4efa Update start_page.md +9e59086 added cristae model (#113) +07a6b61 (tag: 0.2.0) Update __version__.py +5a18127 Activate windows tests (#88) +20cc7b1 Merge pull request #107 from computational-cell-analytics/more-models-and-samples +38b1a2c Additional models +237cde7 Add new sample data +2afed55 added cristae model and single channel transfrom (#100) +92be799 Merge pull request #106 from computational-cell-analytics/adapt-mps-tiling +169ddd4 changed tiling for mps +8cc71f1 added foreground threshold to hmap (#101) +09aa4e0 Optimize scaler (#99) +4741402 added MPS default tiling (#104) +f7cca9a Portal things + few things for paper revision (#102) +38f7907 Add support for cryo-et-data (#89) +1f0387a Merge pull request #98 from computational-cell-analytics/cooper +af10206 use updated mito model +7dc2490 Add missing enviroment file for gpu installation on windows +b62ac1a added params to configure in_channels for model creation (#97) +51e28a4 Add new mito model (#96) +4245c5b Add boundary threshold to hmap (#95) +0483f53 added boundary threshold and area threshold to mitochondrai segmentation (#94) +07b4e7e added seed_distance to parameters of mitochondria segmentation (#93) +9847f1b Fix documentation typo +b56906c Update environment file in build_docs +3d70179 Merge pull request #92 from computational-cell-analytics/update-install +f57105c Update installation instructions +190a62f Fix typo in doc +d829715 Merge branch 'main' of https://github.com/computational-cell-analytics/synaptic-reconstruction into main +dee0935 Add video links to README and documentation +5c080f7 Save segmentations in analysis pipeline example +3ab9a8d Update documentation +79c6242 Add napari plugin images for doc +9cb8763 Implement segmentation post-processing widget +b75b968 Add link to analysis script to documentation +9221c9e Minor typo fixes in the doc (#87) +1810805 (tag: 0.1.0) Merge pull request #79 from computational-cell-analytics/first-release +03d42a8 Fix issues with inference code +9ceb256 Merge branch 'main' into first-release +458ec90 78 improvements to vesicle pool assignment widget (#81) +c0afedb Update the cooper scripts +c75abd9 Update the example scripts and add downloads for data from zenodo +ba56b57 Update tomo sample data and example analysis pipeline +fb1fcd0 Fix issues in active zone segmentation +a7ff94e Udpate example analysis pipeline +22ea714 Refactor segmentation functionality +8b5cf12 Update example analysis pipeline and sample data +4b3b539 Bump version; add outline for analysis pipeline +8a4d26d Update vesicle pool widget (#77) +9b94807 Fix typo +b5c9986 Add short headers for submodules +17557c9 Install pdoc in doc workflow +7793579 Add version to init +4415b18 Try to fix doc build +d0d307a Fix doc building +498eb5a Merge pull request #76 from computational-cell-analytics/clean-up +c4304c3 Add short rdme to scripts +c81d2b0 Rename library in all experiment scripts +004dca8 Rename synaptic_reconstruction package to synapse_net +c4c7aaf Add missing scripts and gitignores +c8cc5ea Update documentation +064ff2d Fix issues in ribbon postprocessing +57b7258 Merge pull request #75 from computational-cell-analytics/add-ribbon-inference +71f9b2c Add new sample data and test for file utils +1cc595e Add ribbon model and refactor IO functionality +0fe01c4 Fix issues in doc +ac90e63 Update start_page.md +409bcb6 Update README.md +efdee13 Merge pull request #73 from computational-cell-analytics/68-create-a-widget-for-vesicle-pool-assignments +223baee Refactor plugin functionality, update pool widget +ec1ff3e Merge branch 'main' into 68-create-a-widget-for-vesicle-pool-assignments +afeb8e7 Merge pull request #74 from computational-cell-analytics/software-updates +333b24a Fix CI +9b3e7b5 Add sample data, CI, and auto-scaling for segmenation CLI +90aa283 chagned display name of vesicle morphology +6989b45 added tooltip for query to see which filters can be applied +827ad03 cleaned up code +414b4cb Update doc and envs +bc39941 added vesicle pool widget +c4bbc00 morphology in dedicated layer +8e8c25d fix right argument for sampling +77a10a9 Merge branch 'main' of https://github.com/computational-cell-analytics/synaptic-reconstruction into main +8ef7559 Fix type annotation +5f15a66 Doc updates (#71) +1e2533f Update documentation and readme +e6971f5 fiexed morphology widgets +d883a07 Update all function doc strings +e2565f6 Merge branch 'main' into doc-updates +a65affb Update function doc strings +ae2a472 Merge pull request #70 from computational-cell-analytics/more-inner-ear-analysis +53e6c9b Add active zone inference +686a38a Merge branch 'main' into doc-updates +d723d5c Merge pull request #52 from computational-cell-analytics/more-inner-ear-analysis +2cc8f92 Merge branch 'main' into more-inner-ear-analysis +69d7a3d Remove lockfile +ff11f88 Add more evaluation scripts +e4c3040 Update docs WIP +1f392ef Add evaluation scripts for cooper data +2da1d37 changed distance and morphology widgets to properly add data to segmetation layer's properties attribute +8dbe7d3 Update build_docs.yaml +d8e37f5 Update build_docs.yaml +91e1237 Update build_docs.yaml +76525c9 Update build_docs.yaml +6b6d1b3 Update build_docs.yaml +e5b734c Update build_docs.yaml +b213c31 Update build_docs.yaml +e01b452 Update build_docs.yaml +06ebbec Update build_docs.yaml +681fc5d Update build_docs.yaml +cf1d722 Update build_docs.yaml +ba25763 Update build_docs.yaml +844b9b9 Update build_docs.yaml +0dd3b18 Update build_docs.yaml +37535b3 Update build_docs.yaml +65d6dd3 Update build_docs.yaml +bc53699 Update build_docs.yaml +f62b8cf Update doc_env.yaml +3844f77 Update doc_env.yaml +29f1768 Update doc_env.yaml +65eac11 Update doc_env.yaml +558fc4f Update doc_env.yaml +ee89f73 Update doc_env.yaml +1af04e3 Update doc_env.yaml +64c8780 Update build_docs.yaml +e34d6b8 Update build_docs.yaml +ea885b2 Update build_docs.yaml +3b27480 Update doc_env.yaml +a57a742 Update build_docs.yaml +b571218 Update doc_env.yaml +7e4844d Update build_docs.yaml +c15173d Update build_docs.yaml +f9d7310 Update build_docs.yaml +d44aced Update build_docs.yaml +e212583 Merge pull request #69 from computational-cell-analytics/59-build-documentation +efe2219 (origin/59-build-documentation) added github workflow yaml etc +d226a94 Merge pull request #67 from computational-cell-analytics/58-add-widget-for-morphology-analysis +0bacc48 refactored some code in to_imod.py +19d9217 added files for doc +4ec71a9 table data index and label_id mixup is corrected +be5808e implemented changes mentioned in PR 67 +e0262f9 Merge pull request #66 from computational-cell-analytics/65-add-mito-train-script +ca9b803 Merge branch 'main' into 58-add-widget-for-morphology-analysis +a59ed03 removed oboslete code +063b3b6 added morphology widget to measure morphology of structures in 2d and 3d volumetric data +70095c4 Update util.py +3281438 added default path for training data and removed obsolete code +b73b5be added train_mitochondria script and util for training +ce94ea8 Add CLI for IMOD export +3b4c6c4 Update analyses +45ef1c3 refactored morphology widge +fe7dfae Update weights_only handling +9fd4017 Add some TODOs for CLI update +4a33f11 Simplify model type selection logic +2ec0fd5 Add examples for network training (#61) +8df8eba some changed got lost +691f183 added napari.yaml +9728951 Update vesicle diameter analysis +59a38db Update all measurements for the inner ear analysis +e0dfda6 Merge branch 'main' into more-inner-ear-analysis +e95dd45 Merge pull request #60 from computational-cell-analytics/56-add-support-for-custom-model-paths-in-segmentation-plugin +3e11446 Update weights_only handling +0db92fc Add some TODOs for CLI update +f3be90a Simplify model type selection logic +008860f Merge remote-tracking branch 'origin/main' into 56-add-support-for-custom-model-paths-in-segmentation-plugin +a974e65 Add examples for network training (#61) +93a66c1 Update data summary +a0c31a8 Fix issue in data aggregation +cb693b1 Update data summaries +51165a5 Fix issues with the segmentation export to IMOD +eb21abd added morphology and first approach to measure vesicles in UI +7941db0 added custom model path with target selector to choose between desired segmentation functions +9b8c7a2 Add more inner ear analysis code +c396113 Merge pull request #55 from computational-cell-analytics/53-next-steps-plugin-implementation +5feff6a Update active zone analysis for SNAP/MUNC data +27c4d9d reworked mentioned parts +cf2898d Merge branch 'main' into 53-next-steps-plugin-implementation +6474587 voxel size in UI and scale is automatically calculated if there is enough information +2ccf340 Add script to extract vesicle diameters for inner ear data +3a13825 added description to volume_reader +3569d4f chagned training resolution to dictionaries +dead9c4 removed obsolete LOCs +cd54805 added scale calculation based on resolution and voxel size +d186ee8 Add download urls and train resolutions for most models +19b2701 Add functionality to extract training data info +5b778f7 Fix issues in filter updates and add script for active zone segmentation +10854de Add sato as better default edge filter +79478c5 added scale calculation based on resolution and voxel size +e3d3798 Add download urls and train resolutions for most models +186df5b Merge branch 'more-inner-ear-analysis' of https://github.com/computational-cell-analytics/synaptic-reconstruction into more-inner-ear-analysis +186c92d Update inner ear analysis scripts +0b7884d Merge branch 'main' of https://github.com/computational-cell-analytics/synaptic-reconstruction into more-inner-ear-analysis +ef44f99 Add functionality to extract training data info +bcc8863 added dummy traiing_voxel_size and dynamic calculation of sclae +16512a7 added voxel_size reading to reader plugin and stored it in napari.layer.metadata +903e59e Update inner ear analysis +30d4a08 added new tiling logic and napari progress bar +a2758f4 Fix issues in filter updates and add script for active zone segmentation +58f3e6b Add sato as better default edge filter +238b577 Merge branch 'main' of https://github.com/computational-cell-analytics/synaptic-reconstruction +305a80b Updates to inner ear training and eval +9472a47 Merge pull request #51 from computational-cell-analytics/more-plugin-updates +4e19b6e Update mito segmentation +ad4741b Update inner ear analysis +a051d85 Update CLI and model logic +55074db Implement CLI WIP +d991be3 Update distance measurement widget +f74e5cd Update plugins WIP +0f40d3c Update inner ear analysis +38d6edf Merge pull request #50 from computational-cell-analytics/49-plugin-updates +9980c2b added distance to object to DistanceMeasureWidget feature and saving to csv file. _compute_seg_object_distances needed to be adapted to accept 2D data (prev. only 3D data was accepted) +514e1f1 removed obsolete code +0ebd45c refactored layer_selector_widgets to BaseWidget and fixed loading image before plugin problem +ff5be12 Merge pull request #48 from computational-cell-analytics/plugin-updates +d5faa99 Updates to napari plugins and distance functionality +9581b2b Merge pull request #47 from computational-cell-analytics/45-user-interface +9622f88 pooch loading with state_dict works.. layer_selector_widgets should be refactored to be in BaseWidget +36d834f Implement inner ear analysis WIP +f56436a Update structure segmentation training +7af1511 added dynamic switch between segmentation functions absed on the model used +ec1f050 added segmentation widget +74c3e69 Merge pull request #46 from computational-cell-analytics/update-comp-seg +59e01dd (origin/update-comp-seg) Update result visualization +64d56e5 Update segmentation functionality and add more inference scripts +993a006 Merge branch 'update-comp-seg' of https://github.com/computational-cell-analytics/synaptic-reconstruction into update-comp-seg +aad578c Update compartment model training +6daf9de Update compartment segmentation +54204cc Add annotation script +bc981e2 Finalize CryoVesNet evaluation +6011472 Update compartment segmentation code +c35b04d Update cryovesnet evaluation +c5fceb7 Add prediction for compartments WIP +01804ec Update compartment segmentation logic +bbecced Merge branch 'main' of https://github.com/computational-cell-analytics/synaptic-reconstruction +9aba9f3 Implement cryo-ves-net evaluation for cooper data +7657af0 Merge pull request #43 from computational-cell-analytics/inner-ear-eval +f75091d Update inner ear analysis +4cdd288 More updates in the ribbon segmentation +85c5670 Update postprocessing code +13f2272 Update domain adaptation training logic for inner ear structures +be05252 Merge branch 'main' into inner-ear-eval +696f47d Update inner ear structure postprocessing +9cb25c2 Merge pull request #42 from computational-cell-analytics/fix-da-training +309d3c1 Merge pull request #39 from computational-cell-analytics/38-add-sampler-to-domain-adaptation +26e37eb Fix some issues in DA training +05d3555 Update supervised_training.py +062b956 Update for inner ear analysis +c61de04 Update postprocessing script +3cd5f0c Implement IMOD mask export WIP +eabd2aa Merge pull request #41 from computational-cell-analytics/sm-dev +b58f1d3 fix mask error +feadf6f Merge branch 'main' of https://github.com/computational-cell-analytics/synaptic-reconstruction into sm-dev +be7ba41 active zone training +fc7a3a1 Update radius extraction for IMOD +475a9b5 Merge pull request #40 from computational-cell-analytics/compartment-training +4363ce4 Update cryovesnet prediction +1b777bf Update compartment segmentation +febad29 Work on compartment segmentation +d47dc0a Implement compartment segmentation WIP +8a6026b Some updates to compartment training and cryo-ves-net +20a1390 Merge branch 'main' into compartment-training +e9c8cf9 Implement training for compartment model +967304e addded sampler to domain adaptation +d1b6bba Add readme for cryovesnet +cdf3315 Add initial code for cryovesnet +997d804 Update ground-truth annotation scripts +c590869 Update processing scripts for compartment annotations +760cd7c more about 2D +a18a992 Update correction table parsing for inner ear synapses +d59f0ff Update sam preprocessing for compartment annotation +dee589f Update gt generation +c412b47 Merge branch 'main' of https://github.com/computational-cell-analytics/synaptic-reconstruction into sm-dev +47428e5 masked segmentation +dd2a9ae Update README.md +83a5cfa Implement factor to increase the vesicle radius in the IMOD export +e6d189c Merge pull request #36 from computational-cell-analytics/vesicle-inference-updates +4696743 Merge pull request #37 from computational-cell-analytics/sm-dev +3edbacf segmentation also for other data than h5 +7b460b2 Unify all segmentation functions +fbfe91b Update scaling and masking functionality in mito and vesicle segmentation +3baea1c Add support for prediction with mask WIP +05bd064 Implement script for mask extraction from IMOD and mask interpolation +d09c245 Update border filtering +492a427 extract presynapse +3a8f223 DA for inner ear data, excluding testset data +a5c8965 DA for endbulb data +f977967 DA for cryo data +8ca79b3 Merge branch 'main' of https://github.com/computational-cell-analytics/synaptic-reconstruction into sm-dev +89ae892 Merge pull request #35 from computational-cell-analytics/update-distance-seg +11a668b Update vesicle segmentation +1bafacd fixing confusing print statement +63b9450 Merge branch 'main' of https://github.com/computational-cell-analytics/synaptic-reconstruction into sm-dev +4c20a77 Merge pull request #34 from computational-cell-analytics/ribbon-synapses +776e8e4 small things +60d2692 Update ribbon synapse evaluation +6802ad7 Add code for ribbon synapse eval +8cf1806 Merge pull request #33 from computational-cell-analytics/sm-dev +a2f351d correction of argparse +0cae179 option to split dataset in only train and val +0d418dc Merge pull request #32 from computational-cell-analytics/structure-training +5ccf1a6 (origin/structure-training) Add domain adaptation script +149d8a1 Merge pull request #31 from computational-cell-analytics/more-visualization +d67fe84 Merge branch 'main' into structure-training +9d02682 minor additions +46cfb54 Merge branch 'more-visualization' of https://github.com/computational-cell-analytics/synaptic-reconstruction into sm-dev +d204845 Minor update +db812d1 Finalize vesicle post-processing changes +9120dee Merge pull request #30 from computational-cell-analytics/sm-dev +7f1b058 matching object imod export to vesicle export +d324b28 few more prints +b7b4a31 Structure training updates +b52742b More updates WIP +796ec01 Update post-processing WIP +79f0125 Merge branch 'sm-dev' into more-visualization +482cddf Add new visualization script +5970812 push everything once more +0c830b6 can pass either model or model path for segmentation, functions have 2 options now, changed order of parameters +01b4cb1 exclude files not directory +f10a2a4 export vesicles for endbulb, not fully tested# +4fe8ef1 fix 2D seg tiling size, implement 2D evaluation +4bf312a 2D vesicle segmentation +f965616 evaluation script +85effc1 assisting file, finalising h5 segmentation +208ba17 fix tiling-halo issue; assisting files +d30537c Update segmentation functionality +8afad01 Merge pull request #29 from computational-cell-analytics/gt-extraction +9db47ed Finalize GT extraction for moser data +51166e7 Add inner ear vesicle extraction script +9841fdf Implement IMOD structure export +47945f4 Add script to check exported data +76654fb export AZ, endbulb, cristae, mito for wichmann data +196f5a2 add small things for training +7dec298 Add ground-truth export and analysis scripts for moser data +4e46957 Merge branch 'main' of https://github.com/computational-cell-analytics/synaptic-reconstruction into sm-dev +87edf43 finish training +80be9ba Update vesicle extraction for cryo data +b7f3f8b Update README.md +b50b829 Merge branch 'main' of https://github.com/computational-cell-analytics/synaptic-reconstruction into sm-dev +8f7085e Fix issues with loss masking code +82ee06b don't redo already postprocessed, small addition +de5b644 don't redo already postprocessed +305f81a Start code for cryo vesicle extraction +414677b Update default ignore table +f78c203 Merge pull request #24 from computational-cell-analytics/actin +6332bb5 Merge branch 'main' of https://github.com/computational-cell-analytics/synaptic-reconstruction into actin +3d42280 Merge pull request #26 from computational-cell-analytics/masked-loss +7ea303a Merge branch 'main' into masked-loss +20f5180 Merge pull request #22 from computational-cell-analytics/sm-dev +b582701 Add actin visualization scripts +56217b0 Add domain adaptation script for actin +f84f281 improve vesicle postprocessing +2d8db55 update right linux code for imod exportation +ab2ca1d Implement masked loss +c796077 Add more IMOD parsing functionality +b21a5ad Update frog synapse extraction +c297ed6 Update analysis for manual annotations in inner ear synapses +4ee4ca0 Use new column to parse inner ear synapses +54fb030 postprocess shape with new functions +ceda07f merging main and sm-dev +6c2669c Merge pull request #25 from computational-cell-analytics/shape-refinement +92cedcc (origin/shape-refinement) Update manual vesicle analysis +c693336 Minor fix +fc88c45 Fix issues in vesicle refinement +a058aef Implement vesicle shape refinement +488ccb7 trying out different wadershed functions +7096869 passing ndim to RawDataset +61242bd first implementation of postprocessing vesicle shape seg +e162f75 Implement first version of actin prediction +653b9c2 Merge branch 'main' into actin +a971a52 Merge pull request #23 from computational-cell-analytics/cooper-vesicle-extraction +a73c796 Implement vesicle correction for inner ear synapses +2def8e0 2D semisupervised, unsupervised, not fully tested +b7670ee add chages for 2D +f0a165a Add script for checking extracted vesicle ground-truth +799defc Implement actin segmentation training +468fa2e Finish vesicle extraction +ad62a9a Add first version of vesicle extraction for cooper data +f6021b0 Merge branch 'main' of https://github.com/computational-cell-analytics/synaptic-reconstruction into sm-dev +0a5d93d check for 2D +5a283ec Work on vesicle extraction WIP +f89fcd5 Improvements in inner ear vesicle analysis +ad68bd3 Update validation table handling +411850f Merge pull request #20 from computational-cell-analytics/training +3d8cb77 Implement training functionality +4832ea3 Implement training WIP +e101e0d Update vesicle GT processing +dd9ad05 Implement vesicle postprocessing +06d3307 Add scripts for ground-truth generation in cooper data +a0b8ee1 Merge pull request #19 from computational-cell-analytics/non-ax-aligned-bb +8ff548f Add ground-truth extraction code +04663e4 Enable import without napari-regiontable +da1a8a7 Update correction pool logic +c3362dd Update inner ear analysis to take into account pool correction +08d6852 Add new layer for vesicle pool correction +6233db5 Merge pull request #18 from computational-cell-analytics/tile-and-border +7d369d6 Make default tiling a function of VRAM and implement border clearing for vesicles +77d9ce6 Fix bug for prediction with nested input folders +6e83143 Enable different extensions for vesicle segmentation and update readme +21eaa65 Merge pull request #16 from computational-cell-analytics/inference_updates +8485053 Add documentation for the new inference code +c37babb Prepare data transfer documentation WIP +99c9abb Unify prediction scripts +4dd165f Update inference code in stucture segmentation script +adba18e Update default tiling logic +409997e Update visualization script +5bcbb1d Merge pull request #15 from computational-cell-analytics/14-add-visualization-script-to-check-segmentations +840cb08 added script that visualizes segmentation +5cfb164 Merge pull request #13 from computational-cell-analytics/imod-export +2317fff Enable passing serialized torch models to prediction scripts +3162438 Merge pull request #12 from computational-cell-analytics/inference-refactor +89da79e changed axis argument for standardize from list to tuple +0102ae8 added cristae segmentation script and fixed incorrect usage of np.stack in inference/util.py +2fce111 Implement imod vesicle export +95550f4 Fix issue in inference_helper and add vesicle segmentation +b87d759 Implement inference helper function +b21f24e Start refactoring inference functionality +a0dd0cf Merge pull request #10 from computational-cell-analytics/9-add-script-to-run-mitochondria-segmentation +5ddb555 (origin/9-add-script-to-run-mitochondria-segmentation) added compression to .tif export and slightly increased the min size filter for mitochondria +868fb99 removed unnecessary code +cd58951 added script to run mitochondria segmentation from directory or single path +0a17e75 Merge pull request #8 from computational-cell-analytics/7-add-segmentation-functionality-for-cristae +add56fb removed file and removed file path from .gitignore +72ac866 added cristae segmentation and simplified script +6c5ae24 added docstring to semgent_cristae +31d8a1e refactored predict_with_torch_em +63abf2d changed from instance segmentation to semantic segmentation +5097cfd added test file to gitignore +8ecbf67 added crstae segmentation +1c8a777 Merge pull request #6 from computational-cell-analytics/4-add-segmentation-functionality-for-mitochondria-cristae +1efa6a7 removed unnecessary imports +dfa5210 added util file to reuse torch_em predict function and tidied up code +3af2613 Stop tracking test.py +16856b0 added mito inference script as well as type annotations for it and for vesicles +cc401d8 added some more postprocessing steps +3e38632 added mito inference script and generic post processing +b39d6ba Misc updates +f7de8d4 Update binning +4a620a9 New skip column +60e4426 Fix issue in vesicle pool analysis +dbf9009 Update vesicle extraction +2c23a04 Update morphology computation +7e0c411 Fix issue in analysis code +c210002 Update morphology computation +777977b Update selected postprocessing +94c6a62 Implement selected postprocessing WIP +b7cdae2 Update validation table handling +58e5836 Update tomogram rescaling and IMOD file processing +4910da3 Update tomogram rescaling script +a6a5ad0 Add script for tomogram rescaling WIP +389be29 Add data extraction scripts for FrogEM +77af8f8 Update README.md +2f4ee2e Update README.md +c7afefb Update README.md +412970d Update README.md +0808f11 Update vesicle annotation +355692c Merge branch 'main' of https://github.com/computational-cell-analytics/synaptic-reconstruction into main +014809e Update manual segmentation handling for inner synapses +da6a470 Add vesicle annotation tool +4805cd7 Fix naming errors in manual analysis +681b48f Simplify manual processing +1ed1823 Update manual segmentation processing +f435f67 Update inner ear scripts +4f2c4c3 Update result checking +75c2128 Minor fixes in run analysis +1407896 Update analysis code +e1134c1 Add checking for manual analysis +835836f Treat corrected vesicle pools in run analysis +092a94d More robust analysis code +6fa7b60 Updates to morphology processing +5e54294 Add export for imod point annotations +7eba5fe Implement analysis of manual annotations WIP +e577cc9 Add script to check the correction table +ce5beb4 Update binning in inner ear result correction +6ed64e4 Fix color assignment without results +de44c5c Update vesicle correction script +6b13ea5 Minor updates to analysis scripts +c3635d5 Fix typo in inner ear processing script +aac4ef4 Update data processing scripts for inner ear tomogram +fd0aa81 Updates for correction inclusion +6ee2333 Handle tomograms with two ribbons in inner ear analysis +454bfda Update distance based post-processing +675b525 Update processing scripts +5be3930 Update structure post-processing for segmentation v2 +077846d Update processing scripts for inner ear synapses +c4bc452 Implement new training data export for inner ear synaptic structures +685a4bb Implement hole closing +d0afa16 Fix distance line creation for empty distances +2533f4c Merge branch 'main' of https://github.com/computational-cell-analytics/synaptic-reconstruction into main +acdb13b Add script for segmentation validation +f862d6b Catch issue in mesh volume computation +1d1bbe7 Minor updates +bcaa48f Update morphology analysis +68c7240 Update result visualization for inner ear +6b8bb88 Implement analysis +d523872 Update structure post-processing +c9907d0 Update structure postprocessing +6aeb3d9 Update segmentation +d8b87c4 Implement processing for new microscope +69627cf Fix ribbon post-processing +f7c1e97 Merge branch 'main' of https://github.com/computational-cell-analytics/synaptic-reconstruction +3e232b4 Update processing code +cc58a20 Update visualization scripts +cd76ab9 Update post-processing for ribbon and membrane +81e6a15 Add post-processing for structure segmentation +edffed3 Update segmentation functionality +fe78084 Add imod extraction functionality and script +e4cb154 Implement table parsing for inner ear experiments +e32f0f3 Update scripts for inner ear processing +eb3ba29 Update gitignore +d48445d Update scripts for inner ear processing +2d780d5 Fix bug in distance measurements +0ed51da Add simple script to check segmentations +eed7bab Add morphology computation for ribbon synapses +d5cf40e Fix issues in distance visualization for ribbon synapses +bf12559 Fix issue in distance sorting +599b427 Enable not filtering duplicates in distance measurement +725abed Fix issue in nearest neighbor filtering logic +cbc06ad (origin/tools-for-vesicles) Update liposome measurement script +80f7d43 Fix bug in distance computation +6a4d90d Add script for distance measurement in liposomes +2d45062 Add vesicle pool assignment for inner ear vesicles +5c0def7 Merge branch 'main' of https://github.com/computational-cell-analytics/synaptic-reconstruction into tools-for-vesicles +6fb683c Enable precomputing distances +64af817 Add preliminary boundary segmentations +42f740e Add more functionality for ribbon syn segmentation and distance measurements +bf10d4c Merge pull request #1 from computational-cell-analytics/tools-for-vesicles +937577d Update distance measurement tool +8f864f1 Implement functionality for filtering ribbons +aa97e08 Implement napari-based tools for segmentation correction and displaying distances +09075e5 Initial commit diff --git a/synapse_net/training/domain_adaptation.py b/synapse_net/training/domain_adaptation.py index 46927db..cc4d73c 100644 --- a/synapse_net/training/domain_adaptation.py +++ b/synapse_net/training/domain_adaptation.py @@ -8,6 +8,7 @@ import torch import torch_em import torch_em.self_training as self_training +from torch_em.self_training.logger import SelfTrainingTensorboardLogger from elf.io import open_file from sklearn.model_selection import train_test_split @@ -18,7 +19,152 @@ from ..inference.inference import get_model_path, compute_scale_from_voxel_size, get_available_models from ..inference.util import _Scaler +#TODO add background_mask functionality to torch_em `MeanTeacherTrainer` and `DefaultPseudoLabeler` +class PseudoLabelerWithBackgroundMask(self_training.DefaultPseudoLabeler): + """Subclass of DefaultPseudoLabeler, which can subtract background from the pseudo labels if a background mask is provided. + By default, assumes that the first channel contains the transformed raw data and the second channel contains the background mask. + Args: + confidence_mask_channel: A specific channel to use for computing the confidence mask. + By default the confidence mask is computed across all channels independently. + This is useful, if only one of the channels encodes a probability. + raw_channel: Channel index of the raw data, which will be used as input to the teacher model + background_mask_channel: Channel index of the background mask, which will be subtracted from the pseudo labels. + kwargs: Additional keyword arguments for `self_training.DefaultPseudoLabeler`. + """ + def __init__( + self, + confidence_mask_channel: Optional[int] = None, + raw_channel: Optional[int] = 0, + background_mask_channel: Optional[int] = 1, + **kwargs + ): + super().__init__(**kwargs) + self.confidence_mask_channel = confidence_mask_channel + self.raw_channel = raw_channel + self.background_mask_channel = background_mask_channel + + def _subtract_background(self, pseudo_labels: torch.Tensor, background_mask: torch.Tensor): + bool_mask = background_mask.bool() + return pseudo_labels.masked_fill(bool_mask, 0) + + def __call__(self, teacher: torch.nn.Module, input_: torch.Tensor) -> torch.Tensor: + """Compute pseudo-labels. + + Args: + teacher: The teacher model. + input_: The input for this batch. + + Returns: + The pseudo-labels. + """ + if input_.ndim != 5: + raise ValueError(f"Expect data with 5 dimensions (B, C, D, H, W), got shape {input_.shape}.") + + has_background_mask = input_.shape[1] > 1 + + if has_background_mask: + if self.background_mask_channel > input_.shape[1]: + raise ValueError(f"Channel index {self.background_mask_channel} is out of bounds for shape {input_.shape}.") + + background_mask = input_[:, self.background_mask_channel].unsqueeze(1) + input_ = input_[:, self.raw_channel].unsqueeze(1) + + pseudo_labels = teacher(input_) + + if self.activation is not None: + pseudo_labels = self.activation(pseudo_labels) + if self.confidence_threshold is None: + label_mask = None + else: + mask_input = pseudo_labels if self.confidence_mask_channel is None\ + else pseudo_labels[self.confidence_mask_channel:(self.confidence_mask_channel+1)] + label_mask = self._compute_label_mask_both_sides(mask_input) if self.threshold_from_both_sides\ + else self._compute_label_mask_one_side(mask_input) + if self.confidence_mask_channel is not None: + size = (pseudo_labels.shape[0], pseudo_labels.shape[1], *([-1] * (pseudo_labels.ndim - 2))) + label_mask = label_mask.expand(*size) + + if has_background_mask: + pseudo_labels = self._subtract_background(pseudo_labels, background_mask) + + return pseudo_labels, label_mask + +class MeanTeacherTrainerWithBackgroundMask(self_training.MeanTeacherTrainer): + """Subclass of MeanTeacherTrainer, updated to handle cases where the background mask is provided. + Once the pseudo labels are computed, the second channel of the teacher input is dropped, if it exists. + The second channel of the student input is also dropped, if it exists, since it is not needed for training. + + Args: + kwargs: Additional keyword arguments for `self_training.MeanTeacherTrainer`. + """ + def __init__(self, **kwargs): + super().__init__(**kwargs) + + def _train_epoch_unsupervised(self, progress, forward_context, backprop): + self.model.train() + + n_iter = 0 + t_per_iter = time.time() + + # Sample from both the supervised and unsupervised loader. + for xu1, xu2 in self.unsupervised_train_loader: + + # Keep only the first channel for xu2 (student input). + if xu2.ndim != 5: + raise ValueError(f"Expect xu2 to have 5 dimensions (B, C, D, H, W), got shape {xu2.shape}.") + if xu2.shape[1] > 1: + xu2 = xu2[:, :1].contiguous() + + xu1, xu2 = xu1.to(self.device, non_blocking=True), xu2.to(self.device, non_blocking=True) + + teacher_input, model_input = xu1, xu2 + + with forward_context(), torch.no_grad(): + # Compute the pseudo labels. + pseudo_labels, label_filter = self.pseudo_labeler(self.teacher, teacher_input) + + # Drop the second channel for xu1 (teacher input) after computing the pseudo labels. + if xu1.ndim != 5: + raise ValueError(f"Expect xu1 to have 5 dimensions (B, C, D, H, W), got shape {xu1.shape}.") + if xu1.shape[1] > 1: + xu1 = xu1[:, :1].contiguous() + + # If we have a sampler then check if the current batch matches the condition for inclusion in training. + if self.sampler is not None: + keep_batch = self.sampler(pseudo_labels, label_filter) + if not keep_batch: + continue + + self.optimizer.zero_grad() + # Perform unsupervised training + with forward_context(): + loss = self.unsupervised_loss(self.model, model_input, pseudo_labels, label_filter) + backprop(loss) + + if self.logger is not None: + with torch.no_grad(), forward_context(): + pred = self.model(model_input) if self._iteration % self.log_image_interval == 0 else None + self.logger.log_train_unsupervised( + self._iteration, loss, xu1, xu2, pred, pseudo_labels, label_filter + ) + lr = [pm["lr"] for pm in self.optimizer.param_groups][0] + self.logger.log_lr(self._iteration, lr) + if self.pseudo_labeler.confidence_threshold is not None: + self.logger.log_ct(self._iteration, self.pseudo_labeler.confidence_threshold) + + with torch.no_grad(): + self._momentum_update() + + self._iteration += 1 + n_iter += 1 + if self._iteration >= self.max_iteration: + break + progress.update(1) + + t_per_iter = (time.time() - t_per_iter) / n_iter + return t_per_iter + def mean_teacher_adaptation( name: str, unsupervised_train_paths: Tuple[str], @@ -37,14 +183,16 @@ def mean_teacher_adaptation( n_iterations: int = int(1e4), n_samples_train: Optional[int] = None, n_samples_val: Optional[int] = None, - train_mask_paths: Optional[Tuple[str]] = None, - val_mask_paths: Optional[Tuple[str]] = None, + train_sample_mask_paths: Optional[Tuple[str]] = None, + val_sample_mask_paths: Optional[Tuple[str]] = None, + train_background_mask_paths: Optional[Tuple[str]] = None, patch_sampler: Optional[callable] = None, pseudo_label_sampler: Optional[callable] = None, - device: int = 0, + target_vsize: Optional[float] = None, + device: Optional[torch.device] = None, check: bool = False, ) -> None: - """Run domain adaptation to transfer a network trained on a source domain for a supervised + """Run domain adapation to transfer a network trained on a source domain for a supervised segmentation task to perform this task on a different target domain. We support different domain adaptation settings: @@ -55,11 +203,9 @@ def mean_teacher_adaptation( Args: name: The name for the checkpoint to be trained. - unsupervsied_train_paths: Filepaths to the hdf5 files or similar file formats - for the training data in the target domain. + unsupervsied_train_paths: Filepaths to the hdf5 or mrc files for the training data in the target domain. This training data is used for unsupervised learning, so it does not require labels. - unsupervised_val_paths: Filepaths to the hdf5 files or similar file formats - for the validation data in the target domain. + unsupervised_val_paths: Filepaths to the hdf5 or mrc files for the validation data in the target domain. This validation data is used for unsupervised learning, so it does not require labels. patch_shape: The patch shape used for a training example. In order to run 2d training pass a patch shape with a singleton in the z-axis, @@ -87,61 +233,46 @@ def mean_teacher_adaptation( based on the patch_shape and size of the volumes used for training. n_samples_val: The number of val samples per epoch. By default this will be estimated based on the patch_shape and size of the volumes used for validation. - train_mask_paths: Sample masks used by the patch sampler to accept or reject patches for training. - val_mask_paths: Sample masks used by the patch sampler to accept or reject patches for validation. - patch_sampler: Accept or reject patches based on a condition. - pseudo_label_sampler: Mask out regions of the pseudo labels where the teacher is not confident before updating the gradients. - device: GPU ID for training. + train_sample_mask_paths: Filepaths to the sample masks used by the patch sampler to accept or reject + patches for training. + val_sample_mask_paths: Filepaths to the sample masks mrc files used by the patch sampler to accept or reject + patches for validation. + train_background_mask_paths: Filepaths to the background masks mrc files used for training. + Background masks are used to subtract background from the pseudo labels before the forward pass. + patch_sampler: A sampler for rejecting patches based on a defined conditon. + pseudo_label_sampler: A sampler for rejecting pseudo-labels based on a defined condition. + target_vsize: Target voxel size in Angstrom for rescaling inputs. + On-the-fly rescaling is currently only implemented for unsupervised training. + `unsupervised_train_paths` and `unsupervised_val_paths` must be mrc files to read input voxel size + and determine scale factor. + device: GPU device for training. check: Whether to check the training and validation loaders instead of running training. - """ # noqa + """ # noqa assert (supervised_train_paths is None) == (supervised_val_paths is None) is_2d, _ = _determine_ndim(patch_shape) - if source_checkpoint is None: - # training from scratch only makes sense if we have supervised training data - # that's why we have the assertion here. - assert supervised_train_paths is not None - print("Mean teacher training from scratch (AdaMT)") - if is_2d: - model = get_2d_model(out_channels=2) - else: - model = get_3d_model(out_channels=2) - reinit_teacher = True - else: - print("Mean teacher training initialized from source model:", source_checkpoint) - if os.path.isdir(source_checkpoint): - model = torch_em.util.load_model(source_checkpoint) - else: - model = torch.load(source_checkpoint, weights_only=False) - reinit_teacher = False - - optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) - scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", factor=0.5, patience=5) - - # self training functionality - pseudo_labeler = self_training.DefaultPseudoLabeler(confidence_threshold=confidence_threshold) - loss = self_training.DefaultSelfTrainingLoss() - loss_and_metric = self_training.DefaultSelfTrainingLossAndMetric() - unsupervised_train_loader = get_unsupervised_loader( - data_paths=unsupervised_train_paths, - raw_key=raw_key, - patch_shape=patch_shape, - batch_size=batch_size, - n_samples=n_samples_train, - sample_mask_paths=train_mask_paths, - sampler=patch_sampler + data_paths=unsupervised_train_paths, + raw_key=raw_key, + patch_shape=patch_shape, + batch_size=batch_size, + n_samples=n_samples_train, + sample_mask_paths=train_sample_mask_paths, + background_mask_paths=train_background_mask_paths, + sampler=patch_sampler, + target_vsize=target_vsize, ) unsupervised_val_loader = get_unsupervised_loader( - data_paths=unsupervised_val_paths, - raw_key=raw_key, - patch_shape=patch_shape, - batch_size=batch_size, - n_samples=n_samples_val, - sample_mask_paths=val_mask_paths, - sampler=patch_sampler + data_paths=unsupervised_val_paths, + raw_key=raw_key, + patch_shape=patch_shape, + batch_size=batch_size, + n_samples=n_samples_val, + sample_mask_paths=val_sample_mask_paths, + background_mask_paths=None, + sampler=patch_sampler, + target_vsize=target_vsize, ) - if supervised_train_paths is not None: assert label_key is not None supervised_train_loader = get_supervised_loader( @@ -164,9 +295,41 @@ def mean_teacher_adaptation( check_loader(supervised_train_loader, n_samples=4) check_loader(supervised_val_loader, n_samples=4) return + + if source_checkpoint is None: + # training from scratch only makes sense if we have supervised training data + # that's why we have the assertion here. + assert supervised_train_paths is not None + print("Mean teacher training from scratch (AdaMT)") + if is_2d: + model = get_2d_model(out_channels=2) + else: + model = get_3d_model(out_channels=2) + reinit_teacher = True + else: + print("Mean teacher training initialized from source model:", source_checkpoint) + if os.path.isdir(source_checkpoint): + model = torch_em.util.load_model(source_checkpoint) + else: + model = torch.load(source_checkpoint, weights_only=False) + reinit_teacher = False + + optimizer = torch.optim.Adam(model.parameters(), lr=lr) + scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", factor=0.5, patience=5) + + # self training functionality + if train_background_mask_paths is not None: + pseudo_labeler = PseudoLabelerWithBackgroundMask(confidence_threshold=confidence_threshold, background_mask_channel=1) + trainer_class = MeanTeacherTrainerWithBackgroundMask + else: + pseudo_labeler = self_training.DefaultPseudoLabeler(confidence_threshold=confidence_threshold) + trainer_class = self_training.MeanTeacherTrainer + + loss = self_training.DefaultSelfTrainingLoss() + loss_and_metric = self_training.DefaultSelfTrainingLossAndMetric() - device = torch.device(f"cuda:{device}") if torch.cuda.is_available() else torch.device("cpu") - trainer = self_training.MeanTeacherTrainer( + device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") + trainer = trainer_class( name=name, model=model, optimizer=optimizer, diff --git a/synapse_net/training/semisupervised_training.py b/synapse_net/training/semisupervised_training.py index 084d798..25778ff 100644 --- a/synapse_net/training/semisupervised_training.py +++ b/synapse_net/training/semisupervised_training.py @@ -1,16 +1,17 @@ -from typing import Optional, Tuple - +from typing import Optional, Tuple, List +import os import numpy as np import uuid import h5py import torch import torch_em import torch_em.self_training as self_training +import torch_em.transform from torchvision import transforms from synapse_net.file_utils import read_mrc from .supervised_training import get_2d_model, get_3d_model, get_supervised_loader, _determine_ndim - +from pathlib import Path def weak_augmentations(p: float = 0.75) -> callable: """The weak augmentations used in the unsupervised data loader. @@ -31,10 +32,44 @@ def weak_augmentations(p: float = 0.75) -> callable: ]) return torch_em.transform.raw.get_raw_transform(normalizer=norm, augmentation1=aug) -def drop_mask_channel(x): - x = x[:1] - return x +#TODO add new arguments for channel-wise transforms to torch_em to minimize the need for helper functions + +# Helper functions: +# DropChannel - drops the unneeded sample_mask channel after sampling patches +# ComposedTransform - combine transforms +# ChannelWiseRawTransform - applies raw_transform to only the raw channel (channel 0) +# ChannelWiseAugmentations - applies augmentations to only the raw channel +# ChannelSplitterSampler - torch_em `MinForegroundSampler` expects raw and mask to be passed as x and y, +# but in `get_unsupervised_loader` they are stacked + +class DropChannel: + def __init__(self, channel: int): + self.channel = channel + + def __call__(self, data): + if data.ndim != 4: + raise ValueError("Expect data with 4 dimensions (C, D, H, W).") + if self.channel >= data.shape[0]: + raise ValueError(f"Drop channel index {self.channel} is out of bounds for shape {data.shape}.") + return np.delete(data, self.channel, axis=0) + +class ChannelWiseRawTransform: + def __init__(self, base_transform: callable, transform_channel: int = 0): + self.base_transform = base_transform + self.transform_channel = transform_channel + + def __call__(self, data): + if data.ndim != 4: + raise ValueError("Expect data with 4 dimensions (C, D, H, W).") + + if self.transform_channel >= data.shape[0]: + raise ValueError(f"Transform channel index {self.transform_channel} is out of bounds for shape {data.shape}.") + + output = data.copy() + output[self.transform_channel] = self.base_transform(data[self.transform_channel]) + return output + class ComposedTransform: def __init__(self, *funcs): self.funcs = funcs @@ -44,6 +79,23 @@ def __call__(self, x): x = f(x) return x +class ChannelWiseAugmentations: + def __init__(self, transform_channel: int = 0, base_augmentations: Optional[callable] = None): + self.transform_channel = transform_channel + + self.base_augmentations = weak_augmentations() if base_augmentations is None else base_augmentations + + def __call__(self, data): + if data.ndim != 4: + raise ValueError("Expect data with 4 dimensions (C, D, H, W).") + + if self.transform_channel >= data.shape[0]: + raise ValueError(f"Augmentations channel index {self.transform_channel} is out of bounds for shape {data.shape}.") + + output = data.clone() + output[self.transform_channel] = self.base_augmentations(data[self.transform_channel]) + return output + class ChannelSplitterSampler: def __init__(self, sampler): self.sampler = sampler @@ -52,20 +104,53 @@ def __call__(self, x): raw, mask = x[0], x[1] return self.sampler(raw, mask) +def get_stacked_path(inputs: List[np.ndarray]): + """ + Helper function to get_unsupervised_loader(). Stacks inputs along the channel axis then writes + to temporary h5 files for use by RawDataset() + + Args: + inputs: List of numpy arrays to be stacked along axis 0. + + Returns: Path to temporary h5 file containing stacked inputs. + """ + TMP_ROOT = os.environ.get("TMPDIR", "/tmp") + tmp_path = f"{TMP_ROOT}/stacked_{uuid.uuid4().hex}.h5" + + c = len(inputs) + ref_shape = inputs[0].shape + for i, arr in enumerate(inputs): + if arr.shape != ref_shape: + raise ValueError(f"Shape mistmatch for input {i}: {arr.shape} != {ref_shape}") + + with h5py.File(tmp_path, "w") as f: + ds = f.create_dataset("raw", shape = (c, *ref_shape), dtype=inputs[0].dtype, + compression=None, chunks=(1, 32, 256, 256)) + + for i, arr in enumerate(inputs): + # cast all inputs to the same dtype + if arr.dtype != ds.dtype: + arr = arr.astype(ds.dtype, copy=False) + ds[i] = arr + + return tmp_path + def get_unsupervised_loader( data_paths: Tuple[str], raw_key: str, patch_shape: Tuple[int, int, int], batch_size: int, - n_samples: Optional[int], + n_samples: Optional[int] = None, sample_mask_paths: Optional[Tuple[str]] = None, + background_mask_paths: Tuple[str] = None, sampler: Optional[callable] = None, exclude_top_and_bottom: bool = False, + target_vsize: Optional[float] = None, ) -> torch.utils.data.DataLoader: """Get a dataloader for unsupervised segmentation training. Args: - data_paths: The filepaths to the hdf5 files containing the training data. + data_paths: The filepaths to the hdf5 or mrc files containing the training data. raw_key: The key that holds the raw data inside of the hdf5. patch_shape: The patch shape used for a training example. In order to run 2d training pass a patch shape with a singleton in the z-axis, @@ -75,55 +160,173 @@ def get_unsupervised_loader( based on the patch_shape and size of the volumes used for training. exclude_top_and_bottom: Whether to exluce the five top and bottom slices to avoid artifacts at the border of tomograms. - sample_mask_paths: The filepaths to the corresponding sample masks for each tomogram. + sample_mask_paths: The mrc filepaths to the corresponding sample masks for each tomogram. + background_mask_paths: The mrc filepaths to the corresponding background masks for each tomogram. sampler: Accept or reject patches based on a condition. + target_vsize: Target voxel size in Angstrom for rescaling. + Source voxel size is read from mrc `data_paths` to determine the scale factor. Returns: The PyTorch dataloader. """ + # We exclude the top and bottom slices where the tomogram reconstruction is bad. - # TODO this seems unneccesary if we have a boundary mask - remove? if exclude_top_and_bottom: roi = np.s_[5:-5, :, :] else: roi = None + + # get configurations + has_sample_mask = sample_mask_paths is not None + has_background_mask = background_mask_paths is not None + apply_rescale = target_vsize is not None + + # initialize class instances + base_transform = torch_em.transform.get_raw_transform() + channelwise_raw_transform = ChannelWiseRawTransform(base_transform) + drop_channel = DropChannel(channel = 1) + + if apply_rescale: + + # read voxel size from mrc to determine the rescale factor + mrc_path = next((p for p in data_paths if Path(p).suffix == ".mrc"), None) + + if mrc_path is None: + raise ValueError("No mrc file found in data_paths to read voxel size.") + + source_vsize = read_mrc(mrc_path)[1] + + scale = ( + source_vsize["z"] / (target_vsize / 10), + source_vsize["y"] / (target_vsize / 10), + source_vsize["x"] / (target_vsize / 10), + ) + # rescaling is performed differently for float and int data + rescale_raw = torch_em.transform.generic.Rescale(scale) + rescale_mask = torch_em.transform.generic.Rescale(scale, is_label=True) + # stack tomograms and masks and write to temp files to use as input to RawDataset() - if sample_mask_paths is not None: + stacked_paths = [] + + # Case 1: both sample masks and background masks are provided, e.g., for the train data loader + if has_sample_mask and has_background_mask: + assert len(data_paths) == len(sample_mask_paths) == len(background_mask_paths), \ + f"Expected equal number of paths, got {len(data_paths)} data paths, {len(sample_mask_paths)} sample mask paths \ + and {len(background_mask_paths)} background mask paths." + + for i, (data_path, sample_mask_path, background_mask_path) in enumerate(zip(data_paths, sample_mask_paths, background_mask_paths)): + if Path(data_path).suffix == ".h5": + with h5py.File(data_path, "r") as f: + raw = f[raw_key][:] + else: + raw = read_mrc(data_path)[0] + sample_mask = read_mrc(sample_mask_path)[0] + background_mask = read_mrc(background_mask_path)[0] + + if apply_rescale: + raw = rescale_raw(raw) + sample_mask = rescale_mask(sample_mask) + background_mask = rescale_mask(background_mask) + print(f"{Path(data_path).stem}: rescaled inputs to {target_vsize}A with shape {raw.shape}") + + stacked_path = get_stacked_path([raw, sample_mask, background_mask]) + stacked_paths.append(stacked_path) + + # update variables for RawDataset() + data_paths = tuple(stacked_paths) + raw_transform = ComposedTransform(channelwise_raw_transform, drop_channel) + augmentations = (ChannelWiseAugmentations(), ChannelWiseAugmentations()) + sampler = ChannelSplitterSampler(sampler) + + # Case 2: only sample masks are provided, e.g., for the validation data loader + elif has_sample_mask: assert len(data_paths) == len(sample_mask_paths), \ - f"Expected equal number of data_paths and and sample_masks_paths, got {len(data_paths)} data paths and {len(sample_mask_paths)} mask paths." - - stacked_paths = [] - for i, (data_path, mask_path) in enumerate(zip(data_paths, sample_mask_paths)): - raw = read_mrc(data_path)[0] - mask = read_mrc(mask_path)[0] - stacked = np.stack([raw, mask], axis=0) + f"Expected equal number of paths, got {len(data_paths)} data paths and {len(sample_mask_paths)} sample mask paths." - tmp_path = f"/tmp/stacked{i}_{uuid.uuid4().hex}.h5" - with h5py.File(tmp_path, "w") as f: - f.create_dataset("raw", data=stacked, compression="gzip") - stacked_paths.append(tmp_path) + for i, (data_path, sample_mask_path) in enumerate(zip(data_paths, sample_mask_paths)): + if Path(data_path).suffix == ".h5": + with h5py.File(data_path, "r") as f: + raw = f[raw_key][:] + else: + raw = read_mrc(data_path)[0] + sample_mask = read_mrc(sample_mask_path)[0] + + if apply_rescale: + raw = rescale_raw(raw) + sample_mask = rescale_mask(sample_mask) + print(f"{Path(data_path).stem}: rescaled inputs to {target_vsize}A with shape {raw.shape}") + + stacked_path = get_stacked_path([raw, sample_mask]) + + + stacked_paths.append(stacked_path) # update variables for RawDataset() data_paths = tuple(stacked_paths) - base_transform = torch_em.transform.get_raw_transform() - raw_transform = ComposedTransform(base_transform, drop_mask_channel) + raw_transform = ComposedTransform(channelwise_raw_transform, drop_channel) + augmentations = (weak_augmentations(), weak_augmentations()) sampler = ChannelSplitterSampler(sampler) - with_channels = True - else: - raw_transform = torch_em.transform.get_raw_transform() - with_channels = False + + # Case 3: only background masks are provided + elif has_background_mask: + assert len(data_paths) == len(background_mask_paths), \ + f"Expected equal number of paths, got {len(data_paths)} data paths and {len(background_mask_paths)} background mask paths." + + for i, (data_path, background_mask_path) in enumerate(zip(data_paths, background_mask_paths)): + if Path(data_path).suffix == ".h5": + with h5py.File(data_path, "r") as f: + raw = f[raw_key][:] + else: + raw = read_mrc(data_path)[0] + background_mask = read_mrc(background_mask_path)[0] + + if apply_rescale: + raw = rescale_raw(raw) + background_mask = rescale_mask(background_mask) + print(f"{Path(data_path).stem}: rescaled inputs to {target_vsize}A with shape {raw.shape}") + + stacked_path = get_stacked_path([raw, background_mask]) + stacked_paths.append(stacked_path) + + # update variables for RawDataset() + data_paths = tuple(stacked_paths) + raw_transform = base_transform + augmentations = (ChannelWiseAugmentations(), ChannelWiseAugmentations()) sampler = None + # Case 4: neither mask is present, use default behavior + else: + for i, data_path in enumerate(data_paths): + + if Path(data_path).suffix == ".h5": + with h5py.File(data_path, "r") as f: + raw = f[raw_key][:] + else: + raw = read_mrc(data_path)[0] + + if apply_rescale: + raw = rescale_raw(raw) + print(f"{Path(data_path).stem}: rescaled inputs to {target_vsize}A with shape {raw.shape}") + + stacked_path = get_stacked_path([raw]) + stacked_paths.append(stacked_path) + + # update variables for RawDataset() + data_paths = tuple(stacked_paths) + raw_transform = base_transform + augmentations = (weak_augmentations(), weak_augmentations()) + sampler = None + + raw_key = "raw" _, ndim = _determine_ndim(patch_shape) transform = torch_em.transform.get_augmentations(ndim=ndim) + with_channels = True if n_samples is None: n_samples_per_ds = None else: n_samples_per_ds = int(n_samples / len(data_paths)) - augmentations = (weak_augmentations(), weak_augmentations()) - datasets = [ torch_em.data.RawDataset(path, raw_key, patch_shape, raw_transform, transform, roi=roi, n_samples=n_samples_per_ds, sampler=sampler, ndim=ndim, with_channels=with_channels, augmentations=augmentations) @@ -131,12 +334,12 @@ def get_unsupervised_loader( ] ds = torch.utils.data.ConcatDataset(datasets) - num_workers = 4 * batch_size + num_workers = 4 * batch_size loader = torch_em.segmentation.get_data_loader(ds, batch_size=batch_size, num_workers=num_workers, shuffle=True) + return loader - # TODO: use different paths for supervised and unsupervised training # (We are currently not using this functionality directly, so this is not a high priority) def semisupervised_training( diff --git a/synapse_net/training/supervised_training.py b/synapse_net/training/supervised_training.py index 47858c3..fbea897 100644 --- a/synapse_net/training/supervised_training.py +++ b/synapse_net/training/supervised_training.py @@ -6,6 +6,7 @@ import torch_em from sklearn.model_selection import train_test_split from torch_em.model import AnisotropicUNet, UNet2d +from torch_em.transform.raw import get_raw_transform, normalize_percentile from synapse_net.inference.inference import get_model_path, get_available_models @@ -101,6 +102,7 @@ def get_supervised_loader( ignore_label: Optional[int] = None, label_transform: Optional[callable] = None, label_paths: Optional[Tuple[str]] = None, + percentile_norm: bool = False, **loader_kwargs, ) -> torch.utils.data.DataLoader: """Get a dataloader for supervised segmentation training. @@ -168,11 +170,13 @@ def get_supervised_loader( elif len(label_paths) != len(data_paths): raise ValueError(f"Data paths and label paths don't match: {len(data_paths)} != {len(label_paths)}") + raw_transform = get_raw_transform(normalizer=normalize_percentile) if percentile_norm else None loader = torch_em.default_segmentation_loader( data_paths, raw_key, label_paths, label_key, sampler=sampler, batch_size=batch_size, patch_shape=patch_shape, ndim=ndim, is_seg_dataset=True, label_transform=label_transform, transform=transform, + raw_transform=raw_transform, num_workers=num_workers, shuffle=shuffle, n_samples=n_samples, label_dtype=label_dtype, rois=rois, **loader_kwargs, ) @@ -205,6 +209,8 @@ def supervised_training( out_channels: int = 2, mask_channel: bool = False, checkpoint_path: Optional[str] = None, + save_every_kth_epoch: Optional[int] = None, + percentile_norm: bool = False, **loader_kwargs, ): """Run supervised segmentation training. @@ -249,16 +255,20 @@ def supervised_training( mask_channel: Whether the last channels in the labels should be used for masking the loss. This can be used to implement more complex masking operations and is not compatible with `ignore_label`. checkpoint_path: Path to the directory where 'best.pt' resides; continue training this model. + save_every_kth_epoch: Save checkpoints after every kth epoch in a separate file. + The corresponding checkpoints will be saved with the naming scheme 'epoch-{epoch}.pt'. loader_kwargs: Additional keyword arguments for the dataloader. """ train_loader = get_supervised_loader(train_paths, raw_key, label_key, patch_shape, batch_size, n_samples=n_samples_train, rois=train_rois, sampler=sampler, ignore_label=ignore_label, label_transform=label_transform, - label_paths=train_label_paths, **loader_kwargs) + label_paths=train_label_paths, percentile_norm=percentile_norm, + **loader_kwargs) val_loader = get_supervised_loader(val_paths, raw_key, label_key, patch_shape, batch_size, n_samples=n_samples_val, rois=val_rois, sampler=sampler, ignore_label=ignore_label, label_transform=label_transform, - label_paths=val_label_paths, **loader_kwargs) + label_paths=val_label_paths, percentile_norm=percentile_norm, + **loader_kwargs) if check: from torch_em.util.debug import check_loader @@ -317,7 +327,7 @@ def supervised_training( loss=loss, metric=metric, ) - trainer.fit(n_iterations) + trainer.fit(n_iterations, save_every_kth_epoch=save_every_kth_epoch) def _derive_key_from_files(files, key):