From 90721f24c4caa2d40bd3877b4d9d5a1c8eef9190 Mon Sep 17 00:00:00 2001 From: Emmanuel Schmidbauer Date: Tue, 21 Mar 2023 12:55:09 -0400 Subject: [PATCH] init commit --- .gitignore | 67 + LICENSE | 201 +++ MANIFEST.in | 1 + README.md | 109 +- dlas/__init__.py | 0 dlas/data/README.md | 81 + dlas/data/__init__.py | 137 ++ dlas/data/audio/__init__.py | 0 dlas/data/audio/audio_with_noise_dataset.py | 246 +++ dlas/data/audio/fast_paired_dataset.py | 336 +++++ .../fast_paired_dataset_with_phonemes.py | 348 +++++ dlas/data/audio/gpt_tts_dataset.py | 113 ++ dlas/data/audio/gpt_tts_tokenizer.json | 1 + dlas/data/audio/grand_conjoined_dataset.py | 246 +++ dlas/data/audio/nv_tacotron_dataset.py | 260 ++++ dlas/data/audio/paired_voice_audio_dataset.py | 350 +++++ dlas/data/audio/preprocessed_mel_dataset.py | 76 + dlas/data/audio/unsupervised_audio_dataset.py | 228 +++ dlas/data/audio/voice_tokenizer.py | 187 +++ dlas/data/audio/wav_aug.py | 64 + dlas/data/combined_dataset.py | 35 + dlas/data/data_sampler.py | 69 + dlas/data/images/__init__.py | 0 .../images/base_unsupervised_image_dataset.py | 143 ++ dlas/data/images/byol_attachment.py | 392 +++++ dlas/data/images/chunk_with_reference.py | 58 + dlas/data/images/cifar.py | 179 +++ dlas/data/images/full_image_dataset.py | 416 ++++++ dlas/data/images/image_corruptor.py | 216 +++ dlas/data/images/image_folder_dataset.py | 306 ++++ dlas/data/images/image_label_parser.py | 130 ++ ..._pair_with_corresponding_points_dataset.py | 82 + dlas/data/images/multi_frame_dataset.py | 120 ++ dlas/data/images/multiscale_dataset.py | 180 +++ dlas/data/images/paired_frame_dataset.py | 104 ++ dlas/data/images/random_dataset.py | 17 + dlas/data/images/single_image_dataset.py | 86 ++ dlas/data/images/stylegan2_dataset.py | 107 ++ dlas/data/images/zip_file_dataset.py | 71 + dlas/data/text/__init__.py | 0 dlas/data/text/hf_datasets_wrapper.py | 39 + dlas/data/torch_dataset.py | 72 + dlas/data/util.py | 655 ++++++++ dlas/data/zero_pad_dict_collate.py | 46 + dlas/models/__init__.py | 0 dlas/models/arch_util.py | 1167 +++++++++++++++ dlas/models/audio/__init__.py | 0 dlas/models/audio/asr/__init__.py | 0 dlas/models/audio/asr/w2v_wrapper.py | 214 +++ dlas/models/audio/audio_resnet.py | 395 +++++ dlas/models/audio/mel2vec.py | 770 ++++++++++ dlas/models/audio/music/__init__.py | 0 dlas/models/audio/music/cheater_gen_ar.py | 130 ++ dlas/models/audio/music/diffwave.py | 189 +++ dlas/models/audio/music/encoders.py | 53 + dlas/models/audio/music/flat_diffusion.py | 389 +++++ dlas/models/audio/music/gpt_music.py | 383 +++++ dlas/models/audio/music/gpt_music2.py | 190 +++ .../audio/music/instrument_quantizer.py | 180 +++ dlas/models/audio/music/m2v_code_to_mel.py | 63 + dlas/models/audio/music/mel2vec_codes_gpt.py | 55 + dlas/models/audio/music/music_quantizer.py | 283 ++++ dlas/models/audio/music/music_quantizer2.py | 292 ++++ dlas/models/audio/music/tfdpc_v5.py | 457 ++++++ .../audio/music/transformer_diffusion12.py | 697 +++++++++ .../audio/music/transformer_diffusion13.py | 354 +++++ .../audio/music/transformer_diffusion14.py | 257 ++++ .../audio/music/unet_diffusion_music_codes.py | 848 +++++++++++ .../music/unet_diffusion_waveform_gen.py | 520 +++++++ .../music/unet_diffusion_waveform_gen3.py | 372 +++++ .../unet_diffusion_waveform_gen_simple.py | 416 ++++++ dlas/models/audio/tts/__init__.py | 0 .../audio/tts/autoregressive_codegen.py | 304 ++++ .../audio/tts/autoregressive_codegen2.py | 282 ++++ dlas/models/audio/tts/ctc_code_generator.py | 132 ++ dlas/models/audio/tts/diffusion_encoder.py | 319 ++++ dlas/models/audio/tts/lucidrains_dvae.py | 287 ++++ dlas/models/audio/tts/mini_encoder.py | 261 ++++ .../audio/tts/random_latent_converter.py | 64 + dlas/models/audio/tts/tacotron2/LICENSE | 32 + dlas/models/audio/tts/tacotron2/__init__.py | 6 + .../audio/tts/tacotron2/audio_processing.py | 94 ++ dlas/models/audio/tts/tacotron2/hparams.py | 90 ++ dlas/models/audio/tts/tacotron2/layers.py | 83 + dlas/models/audio/tts/tacotron2/loss.py | 65 + dlas/models/audio/tts/tacotron2/stft.py | 145 ++ dlas/models/audio/tts/tacotron2/taco_utils.py | 57 + dlas/models/audio/tts/tacotron2/tacotron2.py | 544 +++++++ dlas/models/audio/tts/tacotron2/text/LICENSE | 19 + .../audio/tts/tacotron2/text/__init__.py | 87 ++ .../audio/tts/tacotron2/text/cleaners.py | 90 ++ .../audio/tts/tacotron2/text/cmudict.py | 63 + .../audio/tts/tacotron2/text/numbers.py | 71 + .../audio/tts/tacotron2/text/symbols.py | 20 + .../audio/tts/tacotron2/wave_tacotron.py | 270 ++++ dlas/models/audio/tts/transformer_builders.py | 128 ++ .../audio/tts/transformer_diffusion_tts.py | 293 ++++ .../audio/tts/transformer_diffusion_tts2.py | 281 ++++ dlas/models/audio/tts/unet_diffusion_tts7.py | 572 +++++++ dlas/models/audio/tts/unet_diffusion_tts9.py | 573 +++++++ .../audio/tts/unet_diffusion_tts_flat.py | 433 ++++++ .../audio/tts/unet_diffusion_vocoder.py | 323 ++++ .../tts/unet_diffusion_vocoder_with_ref.py | 398 +++++ dlas/models/audio/tts/unified_voice2.py | 648 ++++++++ dlas/models/audio/tts/unified_voice3.py | 491 ++++++ dlas/models/audio/tts/unified_voice4.py | 482 ++++++ dlas/models/audio/tts/voice_voice_clip.py | 122 ++ dlas/models/audio/tts/w2v_matcher.py | 190 +++ dlas/models/audio/vocoders/__init__.py | 0 .../models/audio/vocoders/univnet/__init__.py | 0 .../audio/vocoders/univnet/generator.py | 121 ++ dlas/models/audio/vocoders/univnet/lvcnet.py | 235 +++ .../audio/vocoders/waveglow/__init__.py | 0 .../audio/vocoders/waveglow/denoiser.py | 43 + .../audio/vocoders/waveglow/waveglow.py | 325 ++++ dlas/models/classifiers/__init__.py | 0 dlas/models/classifiers/cifar_resnet.py | 179 +++ .../classifiers/resnet_with_checkpointing.py | 198 +++ dlas/models/classifiers/torch_models.py | 11 + dlas/models/classifiers/twin_cifar_resnet.py | 255 ++++ .../classifiers/weighted_conv_resnet.py | 450 ++++++ dlas/models/classifiers/wide_kernel_vgg.py | 88 ++ dlas/models/clip/__init__.py | 0 dlas/models/clip/clip.py | 67 + dlas/models/clip/clvp.py | 210 +++ dlas/models/clip/contrastive_audio.py | 279 ++++ dlas/models/clip/cvvp.py | 148 ++ dlas/models/clip/mel_text_clip.py | 161 ++ dlas/models/clip/text_cond_clip.py | 117 ++ dlas/models/clip/text_voice_clip.py | 211 +++ dlas/models/composable/README.md | 61 + dlas/models/composable/__init__.py | 0 dlas/models/diffusion/__init__.py | 0 dlas/models/diffusion/fp16_util.py | 235 +++ dlas/models/diffusion/gaussian_diffusion.py | 1252 ++++++++++++++++ dlas/models/diffusion/losses.py | 77 + dlas/models/diffusion/nn.py | 186 +++ dlas/models/diffusion/resample.py | 195 +++ dlas/models/diffusion/respace.py | 150 ++ dlas/models/diffusion/rrdb_diffusion.py | 233 +++ dlas/models/diffusion/unet_diffusion.py | 936 ++++++++++++ dlas/models/diffusion/unet_latent_guide.py | 851 +++++++++++ dlas/models/image_generation/RRDBNet_arch.py | 355 +++++ dlas/models/image_generation/ResGen_arch.py | 231 +++ dlas/models/image_generation/__init__.py | 0 .../discriminator_vgg_arch.py | 267 ++++ .../models/image_generation/glean/__init__.py | 0 dlas/models/image_generation/glean/glean.py | 138 ++ .../glean/stylegan2_latent_bank.py | 71 + .../image_generation/srflow/FlowActNorms.py | 126 ++ .../srflow/FlowAffineCouplingsAblation.py | 128 ++ .../image_generation/srflow/FlowStep.py | 125 ++ .../srflow/FlowUpsamplerNet.py | 317 ++++ .../image_generation/srflow/Permutations.py | 43 + .../image_generation/srflow/RRDBNet_arch.py | 277 ++++ .../image_generation/srflow/SRFlowNet_arch.py | 194 +++ dlas/models/image_generation/srflow/Split.py | 70 + .../image_generation/srflow/__init__.py | 0 dlas/models/image_generation/srflow/flow.py | 151 ++ .../image_generation/srflow/glow_arch.py | 12 + .../image_generation/srflow/module_util.py | 82 + dlas/models/image_generation/srflow/thops.py | 52 + .../stylegan/Discriminator_StyleGAN.py | 395 +++++ .../image_generation/stylegan/__init__.py | 11 + .../stylegan/stylegan2_lucidrains.py | 983 ++++++++++++ .../stylegan/stylegan2_rosinality.py | 766 ++++++++++ dlas/models/image_latents/__init__.py | 0 dlas/models/image_latents/byol/__init__.py | 0 .../image_latents/byol/byol_model_wrapper.py | 294 ++++ .../image_latents/byol/byol_structural.py | 202 +++ .../fixup_resnet/DiscriminatorResnet_arch.py | 209 +++ .../image_latents/fixup_resnet/__init__.py | 0 dlas/models/image_latents/spinenet_arch.py | 394 +++++ dlas/models/image_latents/vit_latent.py | 95 ++ dlas/models/lucidrains/__init__.py | 0 dlas/models/lucidrains/dalle/__init__.py | 1 + dlas/models/lucidrains/dalle/attention.py | 423 ++++++ dlas/models/lucidrains/dalle/reversible.py | 173 +++ dlas/models/lucidrains/dalle/transformer.py | 261 ++++ dlas/models/lucidrains/performer/__init__.py | 0 .../performer/autoregressive_wrapper.py | 117 ++ .../lucidrains/performer/performer_enc_dec.py | 102 ++ .../lucidrains/performer/performer_pytorch.py | 741 +++++++++ .../models/lucidrains/performer/reversible.py | 169 +++ dlas/models/lucidrains/vq.py | 471 ++++++ dlas/models/lucidrains/x_transformers.py | 1330 +++++++++++++++++ dlas/models/optical_flow/PWCNet.py | 328 ++++ dlas/models/vqvae/__init__.py | 0 dlas/models/vqvae/dvae.py | 245 +++ dlas/models/vqvae/gumbel_quantizer.py | 63 + dlas/models/vqvae/scaled_weight_conv.py | 178 +++ dlas/models/vqvae/vector_quantizer.py | 257 ++++ dlas/models/vqvae/vqvae.py | 320 ++++ dlas/multi_modal_train.py | 69 + dlas/process_video.py | 235 +++ dlas/requirements.txt | 48 + dlas/scripts/__init__.py | 0 dlas/scripts/audio/__init__.py | 0 dlas/scripts/audio/gen/__init__.py | 0 dlas/scripts/audio/gen/ctc_codes.py | 120 ++ dlas/scripts/audio/gen/music_joiner.py | 132 ++ .../audio/gen/speech_synthesis_utils.py | 126 ++ dlas/scripts/audio/gen/use_diffuse_tts.py | 139 ++ .../gen/use_diffuse_voice_translation.py | 141 ++ .../scripts/audio/gen/use_discrete_vocoder.py | 64 + .../audio/gen/use_discrete_vocoder_one_way.py | 64 + dlas/scripts/audio/gen/use_gpt_tts.py | 213 +++ dlas/scripts/audio/gen/use_mel2vec_codes.py | 51 + dlas/scripts/audio/gen/w2v_patcher.py | 16 + dlas/scripts/audio/gen_mel.py | 21 + dlas/scripts/audio/mel_bin_norm_compute.py | 52 + .../play_with_spectral_representations.py | 11 + .../scripts/audio/prep_music/demucs_notes.txt | 18 + .../prep_music/generate_long_cheaters.py | 93 ++ .../audio/prep_music/generate_long_mels.py | 108 ++ .../audio/prep_music/phase_1_split_files.py | 77 + dlas/scripts/audio/preparation/__init__.py | 0 .../preparation/combine_phonetic_and_text.py | 33 + .../filter_clips_with_no_hifreq_data.py | 25 + .../audio/preparation/gen_dvae_codes.py | 44 + .../audio/preparation/phase_1_split_files.py | 76 + .../preparation/phase_2_sample_and_filter.py | 145 ++ .../phase_3_generate_similarities.py | 133 ++ dlas/scripts/audio/preparation/pipeline.py | 25 + .../process_spleeter_filter_outputs.py | 25 + .../audio/preparation/save_mels_to_disk.py | 43 + .../spleeter_filter_noisy_clips.py | 53 + .../preparation/spleeter_utils/__init__.py | 0 .../spleeter_utils/spleeter_dataset.py | 65 + .../audio/preparation/split_on_silence.py | 45 + dlas/scripts/audio/random_mp3_splitter.py | 78 + .../spleeter_split_voice_and_background.py | 74 + dlas/scripts/audio/test_audio_gen.py | 99 ++ dlas/scripts/audio/test_audio_segmentor.py | 108 ++ dlas/scripts/audio/test_audio_similarity.py | 41 + .../audio/test_audio_speech_recognition.py | 78 + dlas/scripts/audio/use_vocoder.py | 29 + dlas/scripts/audio/word_error_rate.py | 44 + .../byol/byol_extract_wrapped_model.py | 24 + dlas/scripts/byol/byol_resnet_playground.py | 213 +++ .../scripts/byol/byol_segformer_playground.py | 258 ++++ dlas/scripts/byol/byol_spinenet_playground.py | 257 ++++ dlas/scripts/byol/tsne_torch.py | 429 ++++++ dlas/scripts/classify_into_folders.py | 83 + .../diffusion/diffusion_correction_surfer.py | 103 ++ dlas/scripts/diffusion/diffusion_inference.py | 73 + .../diffusion/diffusion_noise_surfer.py | 123 ++ .../diffusion/diffusion_recursive_sampler.py | 97 ++ .../diffusion/diffusion_spacing_surfer.py | 98 ++ dlas/scripts/do_to_files.py | 43 + dlas/scripts/extract_square_images.py | 117 ++ dlas/scripts/extract_subimages.py | 186 +++ dlas/scripts/extract_subimages_with_ref.py | 298 ++++ dlas/scripts/extract_temporal_squares.py | 201 +++ dlas/scripts/find_faulty_files.py | 90 ++ dlas/scripts/folderize_imagenet_val.py | 26 + dlas/scripts/gen_kmeans_clusters.py | 91 ++ dlas/scripts/hugging_face_hub_upload.py | 15 + .../scripts/srflow_latent_space_playground.py | 267 ++++ dlas/scripts/stitch_images.py | 20 + .../stylegan2/convert_weights_rosinality.py | 310 ++++ .../scripts/stylegan2/dnnlib/tflib/network.py | 18 + .../ui/image_labeler/image_labeler_ui.py | 197 +++ dlas/scripts/ui/image_labeler/label_editor.py | 34 + .../pretrained_image_patch_classifier.py | 51 + .../test_image_patch_classifier.py | 30 + .../ui/image_pair_labeler/image_pair_ui.py | 107 ++ dlas/scripts/use_generator_as_filter.py | 49 + dlas/scripts/validate_data.py | 69 + dlas/sweep.py | 65 + dlas/test.py | 113 ++ dlas/torch_intermediary/__init__.py | 68 + dlas/train.py | 485 ++++++ dlas/trainer/ExtensibleTrainer.py | 597 ++++++++ dlas/trainer/README.md | 68 + dlas/trainer/__init__.py | 0 dlas/trainer/base_model.py | 204 +++ dlas/trainer/batch_size_optimizer.py | 145 ++ .../custom_training_components/__init__.py | 0 .../progressive_zoom.py | 149 ++ .../stereoscopic.py | 49 + .../tecogan_losses.py | 449 ++++++ dlas/trainer/eval/__init__.py | 0 dlas/trainer/eval/audio_diffusion_fid.py | 427 ++++++ dlas/trainer/eval/eval_wer.py | 145 ++ dlas/trainer/eval/evaluator.py | 59 + dlas/trainer/eval/fid.py | 54 + dlas/trainer/eval/flow_gaussian_nll.py | 34 + dlas/trainer/eval/mel_evaluator.py | 43 + dlas/trainer/eval/music_diffusion_fid.py | 389 +++++ .../single_point_pair_contrastive_eval.py | 70 + dlas/trainer/eval/sr_diffusion_fid.py | 69 + dlas/trainer/eval/sr_fid.py | 93 ++ dlas/trainer/eval/sr_style.py | 63 + dlas/trainer/experiments/__init__.py | 0 dlas/trainer/experiments/experiments.py | 83 + dlas/trainer/feature_model.py | 108 ++ dlas/trainer/inject.py | 71 + dlas/trainer/injectors/__init__.py | 0 dlas/trainer/injectors/audio_injectors.py | 506 +++++++ dlas/trainer/injectors/base_injectors.py | 582 ++++++++ .../injectors/gaussian_diffusion_injector.py | 215 +++ dlas/trainer/injectors/spec_augment.py | 131 ++ dlas/trainer/loss.py | 56 + dlas/trainer/losses.py | 614 ++++++++ dlas/trainer/lr_scheduler.py | 222 +++ dlas/trainer/networks.py | 83 + dlas/trainer/optimizers/lamb.py | 131 ++ dlas/trainer/optimizers/larc.py | 110 ++ dlas/trainer/optimizers/sgd.py | 72 + dlas/trainer/steps.py | 406 +++++ dlas/use_discriminator_as_filter.py | 74 + dlas/utils/__init__.py | 0 dlas/utils/audio.py | 14 + dlas/utils/audio_resampler.py | 256 ++++ dlas/utils/colors.py | 409 +++++ dlas/utils/convert_model.py | 77 + dlas/utils/distributed_checkpont.py | 55 + dlas/utils/gpu_mem_track.py | 84 ++ dlas/utils/kmeans.py | 174 +++ dlas/utils/loss_accumulator.py | 83 + dlas/utils/music_utils.py | 96 ++ dlas/utils/numeric_stability.py | 142 ++ dlas/utils/options.py | 137 ++ dlas/utils/util.py | 683 +++++++++ dlas/utils/weight_scheduler.py | 69 + requirements.txt | 52 + setup.py | 30 + 328 files changed, 58708 insertions(+), 1 deletion(-) create mode 100644 .gitignore create mode 100644 LICENSE create mode 100644 MANIFEST.in create mode 100644 dlas/__init__.py create mode 100644 dlas/data/README.md create mode 100644 dlas/data/__init__.py create mode 100644 dlas/data/audio/__init__.py create mode 100644 dlas/data/audio/audio_with_noise_dataset.py create mode 100644 dlas/data/audio/fast_paired_dataset.py create mode 100644 dlas/data/audio/fast_paired_dataset_with_phonemes.py create mode 100644 dlas/data/audio/gpt_tts_dataset.py create mode 100644 dlas/data/audio/gpt_tts_tokenizer.json create mode 100644 dlas/data/audio/grand_conjoined_dataset.py create mode 100644 dlas/data/audio/nv_tacotron_dataset.py create mode 100644 dlas/data/audio/paired_voice_audio_dataset.py create mode 100644 dlas/data/audio/preprocessed_mel_dataset.py create mode 100644 dlas/data/audio/unsupervised_audio_dataset.py create mode 100644 dlas/data/audio/voice_tokenizer.py create mode 100644 dlas/data/audio/wav_aug.py create mode 100644 dlas/data/combined_dataset.py create mode 100644 dlas/data/data_sampler.py create mode 100644 dlas/data/images/__init__.py create mode 100644 dlas/data/images/base_unsupervised_image_dataset.py create mode 100644 dlas/data/images/byol_attachment.py create mode 100644 dlas/data/images/chunk_with_reference.py create mode 100644 dlas/data/images/cifar.py create mode 100644 dlas/data/images/full_image_dataset.py create mode 100644 dlas/data/images/image_corruptor.py create mode 100644 dlas/data/images/image_folder_dataset.py create mode 100644 dlas/data/images/image_label_parser.py create mode 100644 dlas/data/images/image_pair_with_corresponding_points_dataset.py create mode 100644 dlas/data/images/multi_frame_dataset.py create mode 100644 dlas/data/images/multiscale_dataset.py create mode 100644 dlas/data/images/paired_frame_dataset.py create mode 100644 dlas/data/images/random_dataset.py create mode 100644 dlas/data/images/single_image_dataset.py create mode 100644 dlas/data/images/stylegan2_dataset.py create mode 100644 dlas/data/images/zip_file_dataset.py create mode 100644 dlas/data/text/__init__.py create mode 100644 dlas/data/text/hf_datasets_wrapper.py create mode 100644 dlas/data/torch_dataset.py create mode 100644 dlas/data/util.py create mode 100644 dlas/data/zero_pad_dict_collate.py create mode 100644 dlas/models/__init__.py create mode 100644 dlas/models/arch_util.py create mode 100644 dlas/models/audio/__init__.py create mode 100644 dlas/models/audio/asr/__init__.py create mode 100644 dlas/models/audio/asr/w2v_wrapper.py create mode 100644 dlas/models/audio/audio_resnet.py create mode 100644 dlas/models/audio/mel2vec.py create mode 100644 dlas/models/audio/music/__init__.py create mode 100644 dlas/models/audio/music/cheater_gen_ar.py create mode 100644 dlas/models/audio/music/diffwave.py create mode 100644 dlas/models/audio/music/encoders.py create mode 100644 dlas/models/audio/music/flat_diffusion.py create mode 100644 dlas/models/audio/music/gpt_music.py create mode 100644 dlas/models/audio/music/gpt_music2.py create mode 100644 dlas/models/audio/music/instrument_quantizer.py create mode 100644 dlas/models/audio/music/m2v_code_to_mel.py create mode 100644 dlas/models/audio/music/mel2vec_codes_gpt.py create mode 100644 dlas/models/audio/music/music_quantizer.py create mode 100644 dlas/models/audio/music/music_quantizer2.py create mode 100644 dlas/models/audio/music/tfdpc_v5.py create mode 100644 dlas/models/audio/music/transformer_diffusion12.py create mode 100644 dlas/models/audio/music/transformer_diffusion13.py create mode 100644 dlas/models/audio/music/transformer_diffusion14.py create mode 100644 dlas/models/audio/music/unet_diffusion_music_codes.py create mode 100644 dlas/models/audio/music/unet_diffusion_waveform_gen.py create mode 100644 dlas/models/audio/music/unet_diffusion_waveform_gen3.py create mode 100644 dlas/models/audio/music/unet_diffusion_waveform_gen_simple.py create mode 100644 dlas/models/audio/tts/__init__.py create mode 100644 dlas/models/audio/tts/autoregressive_codegen.py create mode 100644 dlas/models/audio/tts/autoregressive_codegen2.py create mode 100644 dlas/models/audio/tts/ctc_code_generator.py create mode 100644 dlas/models/audio/tts/diffusion_encoder.py create mode 100644 dlas/models/audio/tts/lucidrains_dvae.py create mode 100644 dlas/models/audio/tts/mini_encoder.py create mode 100644 dlas/models/audio/tts/random_latent_converter.py create mode 100644 dlas/models/audio/tts/tacotron2/LICENSE create mode 100644 dlas/models/audio/tts/tacotron2/__init__.py create mode 100644 dlas/models/audio/tts/tacotron2/audio_processing.py create mode 100644 dlas/models/audio/tts/tacotron2/hparams.py create mode 100644 dlas/models/audio/tts/tacotron2/layers.py create mode 100644 dlas/models/audio/tts/tacotron2/loss.py create mode 100644 dlas/models/audio/tts/tacotron2/stft.py create mode 100644 dlas/models/audio/tts/tacotron2/taco_utils.py create mode 100644 dlas/models/audio/tts/tacotron2/tacotron2.py create mode 100644 dlas/models/audio/tts/tacotron2/text/LICENSE create mode 100644 dlas/models/audio/tts/tacotron2/text/__init__.py create mode 100644 dlas/models/audio/tts/tacotron2/text/cleaners.py create mode 100644 dlas/models/audio/tts/tacotron2/text/cmudict.py create mode 100644 dlas/models/audio/tts/tacotron2/text/numbers.py create mode 100644 dlas/models/audio/tts/tacotron2/text/symbols.py create mode 100644 dlas/models/audio/tts/tacotron2/wave_tacotron.py create mode 100644 dlas/models/audio/tts/transformer_builders.py create mode 100644 dlas/models/audio/tts/transformer_diffusion_tts.py create mode 100644 dlas/models/audio/tts/transformer_diffusion_tts2.py create mode 100644 dlas/models/audio/tts/unet_diffusion_tts7.py create mode 100644 dlas/models/audio/tts/unet_diffusion_tts9.py create mode 100644 dlas/models/audio/tts/unet_diffusion_tts_flat.py create mode 100644 dlas/models/audio/tts/unet_diffusion_vocoder.py create mode 100644 dlas/models/audio/tts/unet_diffusion_vocoder_with_ref.py create mode 100644 dlas/models/audio/tts/unified_voice2.py create mode 100644 dlas/models/audio/tts/unified_voice3.py create mode 100644 dlas/models/audio/tts/unified_voice4.py create mode 100644 dlas/models/audio/tts/voice_voice_clip.py create mode 100644 dlas/models/audio/tts/w2v_matcher.py create mode 100644 dlas/models/audio/vocoders/__init__.py create mode 100644 dlas/models/audio/vocoders/univnet/__init__.py create mode 100644 dlas/models/audio/vocoders/univnet/generator.py create mode 100644 dlas/models/audio/vocoders/univnet/lvcnet.py create mode 100644 dlas/models/audio/vocoders/waveglow/__init__.py create mode 100644 dlas/models/audio/vocoders/waveglow/denoiser.py create mode 100644 dlas/models/audio/vocoders/waveglow/waveglow.py create mode 100644 dlas/models/classifiers/__init__.py create mode 100644 dlas/models/classifiers/cifar_resnet.py create mode 100644 dlas/models/classifiers/resnet_with_checkpointing.py create mode 100644 dlas/models/classifiers/torch_models.py create mode 100644 dlas/models/classifiers/twin_cifar_resnet.py create mode 100644 dlas/models/classifiers/weighted_conv_resnet.py create mode 100644 dlas/models/classifiers/wide_kernel_vgg.py create mode 100644 dlas/models/clip/__init__.py create mode 100644 dlas/models/clip/clip.py create mode 100644 dlas/models/clip/clvp.py create mode 100644 dlas/models/clip/contrastive_audio.py create mode 100644 dlas/models/clip/cvvp.py create mode 100644 dlas/models/clip/mel_text_clip.py create mode 100644 dlas/models/clip/text_cond_clip.py create mode 100644 dlas/models/clip/text_voice_clip.py create mode 100644 dlas/models/composable/README.md create mode 100644 dlas/models/composable/__init__.py create mode 100644 dlas/models/diffusion/__init__.py create mode 100644 dlas/models/diffusion/fp16_util.py create mode 100644 dlas/models/diffusion/gaussian_diffusion.py create mode 100644 dlas/models/diffusion/losses.py create mode 100644 dlas/models/diffusion/nn.py create mode 100644 dlas/models/diffusion/resample.py create mode 100644 dlas/models/diffusion/respace.py create mode 100644 dlas/models/diffusion/rrdb_diffusion.py create mode 100644 dlas/models/diffusion/unet_diffusion.py create mode 100644 dlas/models/diffusion/unet_latent_guide.py create mode 100644 dlas/models/image_generation/RRDBNet_arch.py create mode 100644 dlas/models/image_generation/ResGen_arch.py create mode 100644 dlas/models/image_generation/__init__.py create mode 100644 dlas/models/image_generation/discriminator_vgg_arch.py create mode 100644 dlas/models/image_generation/glean/__init__.py create mode 100644 dlas/models/image_generation/glean/glean.py create mode 100644 dlas/models/image_generation/glean/stylegan2_latent_bank.py create mode 100644 dlas/models/image_generation/srflow/FlowActNorms.py create mode 100644 dlas/models/image_generation/srflow/FlowAffineCouplingsAblation.py create mode 100644 dlas/models/image_generation/srflow/FlowStep.py create mode 100644 dlas/models/image_generation/srflow/FlowUpsamplerNet.py create mode 100644 dlas/models/image_generation/srflow/Permutations.py create mode 100644 dlas/models/image_generation/srflow/RRDBNet_arch.py create mode 100644 dlas/models/image_generation/srflow/SRFlowNet_arch.py create mode 100644 dlas/models/image_generation/srflow/Split.py create mode 100644 dlas/models/image_generation/srflow/__init__.py create mode 100644 dlas/models/image_generation/srflow/flow.py create mode 100644 dlas/models/image_generation/srflow/glow_arch.py create mode 100644 dlas/models/image_generation/srflow/module_util.py create mode 100644 dlas/models/image_generation/srflow/thops.py create mode 100644 dlas/models/image_generation/stylegan/Discriminator_StyleGAN.py create mode 100644 dlas/models/image_generation/stylegan/__init__.py create mode 100644 dlas/models/image_generation/stylegan/stylegan2_lucidrains.py create mode 100644 dlas/models/image_generation/stylegan/stylegan2_rosinality.py create mode 100644 dlas/models/image_latents/__init__.py create mode 100644 dlas/models/image_latents/byol/__init__.py create mode 100644 dlas/models/image_latents/byol/byol_model_wrapper.py create mode 100644 dlas/models/image_latents/byol/byol_structural.py create mode 100644 dlas/models/image_latents/fixup_resnet/DiscriminatorResnet_arch.py create mode 100644 dlas/models/image_latents/fixup_resnet/__init__.py create mode 100644 dlas/models/image_latents/spinenet_arch.py create mode 100644 dlas/models/image_latents/vit_latent.py create mode 100644 dlas/models/lucidrains/__init__.py create mode 100644 dlas/models/lucidrains/dalle/__init__.py create mode 100644 dlas/models/lucidrains/dalle/attention.py create mode 100644 dlas/models/lucidrains/dalle/reversible.py create mode 100644 dlas/models/lucidrains/dalle/transformer.py create mode 100644 dlas/models/lucidrains/performer/__init__.py create mode 100644 dlas/models/lucidrains/performer/autoregressive_wrapper.py create mode 100644 dlas/models/lucidrains/performer/performer_enc_dec.py create mode 100644 dlas/models/lucidrains/performer/performer_pytorch.py create mode 100644 dlas/models/lucidrains/performer/reversible.py create mode 100644 dlas/models/lucidrains/vq.py create mode 100644 dlas/models/lucidrains/x_transformers.py create mode 100644 dlas/models/optical_flow/PWCNet.py create mode 100644 dlas/models/vqvae/__init__.py create mode 100644 dlas/models/vqvae/dvae.py create mode 100644 dlas/models/vqvae/gumbel_quantizer.py create mode 100644 dlas/models/vqvae/scaled_weight_conv.py create mode 100644 dlas/models/vqvae/vector_quantizer.py create mode 100644 dlas/models/vqvae/vqvae.py create mode 100644 dlas/multi_modal_train.py create mode 100644 dlas/process_video.py create mode 100644 dlas/requirements.txt create mode 100644 dlas/scripts/__init__.py create mode 100644 dlas/scripts/audio/__init__.py create mode 100644 dlas/scripts/audio/gen/__init__.py create mode 100644 dlas/scripts/audio/gen/ctc_codes.py create mode 100644 dlas/scripts/audio/gen/music_joiner.py create mode 100644 dlas/scripts/audio/gen/speech_synthesis_utils.py create mode 100644 dlas/scripts/audio/gen/use_diffuse_tts.py create mode 100644 dlas/scripts/audio/gen/use_diffuse_voice_translation.py create mode 100644 dlas/scripts/audio/gen/use_discrete_vocoder.py create mode 100644 dlas/scripts/audio/gen/use_discrete_vocoder_one_way.py create mode 100644 dlas/scripts/audio/gen/use_gpt_tts.py create mode 100644 dlas/scripts/audio/gen/use_mel2vec_codes.py create mode 100644 dlas/scripts/audio/gen/w2v_patcher.py create mode 100644 dlas/scripts/audio/gen_mel.py create mode 100644 dlas/scripts/audio/mel_bin_norm_compute.py create mode 100644 dlas/scripts/audio/play_with_spectral_representations.py create mode 100644 dlas/scripts/audio/prep_music/demucs_notes.txt create mode 100644 dlas/scripts/audio/prep_music/generate_long_cheaters.py create mode 100644 dlas/scripts/audio/prep_music/generate_long_mels.py create mode 100644 dlas/scripts/audio/prep_music/phase_1_split_files.py create mode 100644 dlas/scripts/audio/preparation/__init__.py create mode 100644 dlas/scripts/audio/preparation/combine_phonetic_and_text.py create mode 100644 dlas/scripts/audio/preparation/filter_clips_with_no_hifreq_data.py create mode 100644 dlas/scripts/audio/preparation/gen_dvae_codes.py create mode 100644 dlas/scripts/audio/preparation/phase_1_split_files.py create mode 100644 dlas/scripts/audio/preparation/phase_2_sample_and_filter.py create mode 100644 dlas/scripts/audio/preparation/phase_3_generate_similarities.py create mode 100644 dlas/scripts/audio/preparation/pipeline.py create mode 100644 dlas/scripts/audio/preparation/process_spleeter_filter_outputs.py create mode 100644 dlas/scripts/audio/preparation/save_mels_to_disk.py create mode 100644 dlas/scripts/audio/preparation/spleeter_filter_noisy_clips.py create mode 100644 dlas/scripts/audio/preparation/spleeter_utils/__init__.py create mode 100644 dlas/scripts/audio/preparation/spleeter_utils/spleeter_dataset.py create mode 100644 dlas/scripts/audio/preparation/split_on_silence.py create mode 100644 dlas/scripts/audio/random_mp3_splitter.py create mode 100644 dlas/scripts/audio/spleeter_split_voice_and_background.py create mode 100644 dlas/scripts/audio/test_audio_gen.py create mode 100644 dlas/scripts/audio/test_audio_segmentor.py create mode 100644 dlas/scripts/audio/test_audio_similarity.py create mode 100644 dlas/scripts/audio/test_audio_speech_recognition.py create mode 100644 dlas/scripts/audio/use_vocoder.py create mode 100644 dlas/scripts/audio/word_error_rate.py create mode 100644 dlas/scripts/byol/byol_extract_wrapped_model.py create mode 100644 dlas/scripts/byol/byol_resnet_playground.py create mode 100644 dlas/scripts/byol/byol_segformer_playground.py create mode 100644 dlas/scripts/byol/byol_spinenet_playground.py create mode 100644 dlas/scripts/byol/tsne_torch.py create mode 100644 dlas/scripts/classify_into_folders.py create mode 100644 dlas/scripts/diffusion/diffusion_correction_surfer.py create mode 100644 dlas/scripts/diffusion/diffusion_inference.py create mode 100644 dlas/scripts/diffusion/diffusion_noise_surfer.py create mode 100644 dlas/scripts/diffusion/diffusion_recursive_sampler.py create mode 100644 dlas/scripts/diffusion/diffusion_spacing_surfer.py create mode 100644 dlas/scripts/do_to_files.py create mode 100644 dlas/scripts/extract_square_images.py create mode 100644 dlas/scripts/extract_subimages.py create mode 100644 dlas/scripts/extract_subimages_with_ref.py create mode 100644 dlas/scripts/extract_temporal_squares.py create mode 100644 dlas/scripts/find_faulty_files.py create mode 100644 dlas/scripts/folderize_imagenet_val.py create mode 100644 dlas/scripts/gen_kmeans_clusters.py create mode 100644 dlas/scripts/hugging_face_hub_upload.py create mode 100644 dlas/scripts/srflow_latent_space_playground.py create mode 100644 dlas/scripts/stitch_images.py create mode 100644 dlas/scripts/stylegan2/convert_weights_rosinality.py create mode 100644 dlas/scripts/stylegan2/dnnlib/tflib/network.py create mode 100644 dlas/scripts/ui/image_labeler/image_labeler_ui.py create mode 100644 dlas/scripts/ui/image_labeler/label_editor.py create mode 100644 dlas/scripts/ui/image_labeler/pretrained_image_patch_classifier.py create mode 100644 dlas/scripts/ui/image_labeler/test_image_patch_classifier.py create mode 100644 dlas/scripts/ui/image_pair_labeler/image_pair_ui.py create mode 100644 dlas/scripts/use_generator_as_filter.py create mode 100644 dlas/scripts/validate_data.py create mode 100644 dlas/sweep.py create mode 100644 dlas/test.py create mode 100644 dlas/torch_intermediary/__init__.py create mode 100644 dlas/train.py create mode 100644 dlas/trainer/ExtensibleTrainer.py create mode 100644 dlas/trainer/README.md create mode 100644 dlas/trainer/__init__.py create mode 100644 dlas/trainer/base_model.py create mode 100644 dlas/trainer/batch_size_optimizer.py create mode 100644 dlas/trainer/custom_training_components/__init__.py create mode 100644 dlas/trainer/custom_training_components/progressive_zoom.py create mode 100644 dlas/trainer/custom_training_components/stereoscopic.py create mode 100644 dlas/trainer/custom_training_components/tecogan_losses.py create mode 100644 dlas/trainer/eval/__init__.py create mode 100644 dlas/trainer/eval/audio_diffusion_fid.py create mode 100644 dlas/trainer/eval/eval_wer.py create mode 100644 dlas/trainer/eval/evaluator.py create mode 100644 dlas/trainer/eval/fid.py create mode 100644 dlas/trainer/eval/flow_gaussian_nll.py create mode 100644 dlas/trainer/eval/mel_evaluator.py create mode 100644 dlas/trainer/eval/music_diffusion_fid.py create mode 100644 dlas/trainer/eval/single_point_pair_contrastive_eval.py create mode 100644 dlas/trainer/eval/sr_diffusion_fid.py create mode 100644 dlas/trainer/eval/sr_fid.py create mode 100644 dlas/trainer/eval/sr_style.py create mode 100644 dlas/trainer/experiments/__init__.py create mode 100644 dlas/trainer/experiments/experiments.py create mode 100644 dlas/trainer/feature_model.py create mode 100644 dlas/trainer/inject.py create mode 100644 dlas/trainer/injectors/__init__.py create mode 100644 dlas/trainer/injectors/audio_injectors.py create mode 100644 dlas/trainer/injectors/base_injectors.py create mode 100644 dlas/trainer/injectors/gaussian_diffusion_injector.py create mode 100644 dlas/trainer/injectors/spec_augment.py create mode 100644 dlas/trainer/loss.py create mode 100644 dlas/trainer/losses.py create mode 100644 dlas/trainer/lr_scheduler.py create mode 100644 dlas/trainer/networks.py create mode 100644 dlas/trainer/optimizers/lamb.py create mode 100644 dlas/trainer/optimizers/larc.py create mode 100644 dlas/trainer/optimizers/sgd.py create mode 100644 dlas/trainer/steps.py create mode 100644 dlas/use_discriminator_as_filter.py create mode 100644 dlas/utils/__init__.py create mode 100644 dlas/utils/audio.py create mode 100644 dlas/utils/audio_resampler.py create mode 100644 dlas/utils/colors.py create mode 100644 dlas/utils/convert_model.py create mode 100644 dlas/utils/distributed_checkpont.py create mode 100644 dlas/utils/gpu_mem_track.py create mode 100644 dlas/utils/kmeans.py create mode 100644 dlas/utils/loss_accumulator.py create mode 100644 dlas/utils/music_utils.py create mode 100644 dlas/utils/numeric_stability.py create mode 100644 dlas/utils/options.py create mode 100644 dlas/utils/util.py create mode 100644 dlas/utils/weight_scheduler.py create mode 100644 requirements.txt create mode 100644 setup.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..69b3218 --- /dev/null +++ b/.gitignore @@ -0,0 +1,67 @@ + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +env/ +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +*.egg-info/ +.installed.cfg +*.egg + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*,cover +.hypothesis/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +#Ipython Notebook +.ipynb_checkpoints + +# pyenv +.python-version diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..261eeb9 --- /dev/null +++ b/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..72a4f93 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1 @@ +recursive-include dlas/* diff --git a/README.md b/README.md index 23191fd..610c636 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,109 @@ -# dlas +# (QoL improvements for) Deep Learning Art School +This fork of [neonbjb/DL-Art-School](https://github.com/neonbjb/DL-Art-School/) contains a few fixes and QoL improvements, including but not limited to: +* sanity tidying, like: + - not outputing to `./DL-Art-School/experiments/` + - the custom module loader for networks/injectors getting fixed + - BitsAndBytes integration: + + working but output untested: Adam/AdamW + + toggles available in `./codes/torch_indermediary/__init__.py` + +--- +# Deep Learning Art School + +Send your Pytorch model to art class! + +This repository is both a framework and a set of tools for training deep neural networks that create images. It started +as a branch of the [open-mmlab](https://github.com/open-mmlab) project developed by [Multimedia Laboratory, CUHK](http://mmlab.ie.cuhk.edu.hk) +but has been almost completely re-written at every level. + +## Why do we need another training framework + +These are a dime a dozen, no doubt. DL Art School (*DLAS*) differentiates itself by being configuration driven. You write +the model code (specifically, a torch.nn.Module) and (possibly) some losses, then you cobble together a config file written +in yaml that tells DLAS how to train it. Swapping model architectures and tuning hyper-parameters is simple and often +requires no changes to actual code. You also don't need to remember complex command line incantations. This effectively +enables you to run multiple concurrent experiments that use the same codebase, as well as retain backwards compatibility +for past experiments. + +Training effective generators often means juggling multiple loss functions. As a result, DLAS' configuration language is +specifically designed to make it easy to support large number of losses and networks that interact with each other. As an +example: some GANs I have trained in this framework consist of more than 15 losses and use 2 separate discriminators and +require no bespoke code. + +Generators are also notorious GPU memory hogs. I have spent substantial time streamlining the training framework to support +gradient checkpointing and FP16. DLAS also supports "mega batching", where multiple forward passes contribute to a single +backward pass. Most models can be trained on midrange GPUs with 8-11GB of memory. + +The final value-added feature is interpretability. Tensorboard logging operates out of the box with no custom code. +Intermediate images from within the training pipeline can be intermittently surfaced as normal PNG files so you can +see what your network is up to. Validation passes are also cached as images so you can view how your network improves +over time. + +## Modeling Capabilities + +DLAS was built with extensibility in mind. One of the reasons I'm putting in the effort to better document this code is the +incredible ease with which I have been able to train entirely new model types with no changes to the core training code. + +I intend to fill out the sections below with sample configurations which can be used to train different architectures. +You will need to bring your own data. + +### Super-resolution +- [GAN-based SR (ESRGAN)](https://github.com/neonbjb/DL-Art-School/tree/gan_lab/recipes/esrgan) +- [SRFlow](https://github.com/neonbjb/DL-Art-School/tree/gan_lab/recipes/srflow) +- [GLEAN](https://github.com/neonbjb/DL-Art-School/tree/gan_lab/recipes/glean) +- Video SR (TecoGAN) (*documentation TBC*) + +### Style Transfer +* Stylegan2 (*documentation TBC*) + +### Latent development +* [BYOL](https://github.com/neonbjb/DL-Art-School/tree/gan_lab/recipes/byol) +* iGPT (*documentation TBC*) + +## Dependencies and Installation + +- Python 3 +- [PyTorch >= 1.6](https://pytorch.org) +- NVIDIA GPU + [CUDA](https://developer.nvidia.com/cuda-downloads) +- Python packages: `pip install -r requirements.txt` +- Some video utilities require [FFMPEG](https://ffmpeg.org/) + +## User Guide +TBC + +### Development Environment +If you aren't already using [Pycharm](https://www.jetbrains.com/pycharm/) - now is the time to try it out. This project was built in Pycharm and comes with +an IDEA project for you to get started with. I've done all of my development on this repo in this IDE and lean heavily +on its incredible debugger. It's free. Try it out. You won't be sorry. + +### Dataset Preparation +DLAS comes with some Dataset instances that I have created for my own use. Unless you want to use one of the recipes above, you'll need to provide your own. Here is how to add your own Dataset: + +1. Create a Dataset in codes/data/ which takes a single Python dict as a constructor and extracts options from that dict. +2. Register your Dataset in codes/data/__init__.py +3. Your Dataset should return a dict of tensors. The keys of the dict are injected directly into the training state, which you can interact within your configuration file. + +### Training and Testing +There are currently 3 base scripts for interacting with models. They all take a single parameter, `-opt` which specifies the configuration file which controls how they work. Configs (will be) documented above in the user guide. + +#### train.py +Start (or continue) a training session: +`python train.py -opt ` + +Start a distributed training session: +`python -m torch.distributed.launch --nproc_per_node= --master_port=1234 train.py -o --launcher=pytorch` + +#### test.py +Runs a model against a validation or test set of data and reports metrics (for now, just PSNR and a custom perceptual metric) +`python test.py -opt ` + +#### process_video.py +Breaks a video into individual frames and uses a network to do processing on it, then reassembles the output back into video form. +`python process_video -opt ` + +## Contributing +At this time I am not taking feature requests or bug reports, but I appreciate all contributions. + +## License +This project is released under the Apache 2.0 license. diff --git a/dlas/__init__.py b/dlas/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/dlas/data/README.md b/dlas/data/README.md new file mode 100644 index 0000000..be9af64 --- /dev/null +++ b/dlas/data/README.md @@ -0,0 +1,81 @@ +# DLAS Datasets + +## Quick Overview + +DLAS uses the standard Torch Dataset infrastructure. Datasets are expected to be constructed using an "options" dict, +which is fed directly from the configuration file. They are also expected to output a dict, where the keys are injected +directly into the trainer state. + +Datasets conforming to the above expectations must be registered in `__init__.py` to be used by a configuration. + +## Reference Datasets + +This directory contains several reference datasets which I have used in building DLAS. They include: + +1. Stylegan2Dataset - Reads a set of images from a directory, performs some basic augmentations on them and injects + them directly into the state. LQ = HQ in this dataset. +1. SingleImageDataset - Reads image patches from a 'chunked' format along with the reference image and metadata about + how the patch was originally computed. The 'chunked' format is described below. Includes built-in ImageCorruption + features actuated by `image_corruptor.py`. +1. MultiframeDataset - Similar to SingleImageDataset, but infers a temporal relationship between images based on their + filenames: the last 12 characters before the file extension are assumed to be a frame counter. Images from this + dataset are grouped together with a temporal dimension for working with video data. +1. ImageFolderDataset - Reads raw images from a folder and feeds them into the model. Capable of performing corruptions + on those images like the above. +1. MultiscaleDataset - Reads full images from a directory and builds a tree of images constructed by cropping squares + from the source image and resizing them to the target size recursively until the native resolution is hit. Each + recursive step decreases the crop size by a factor of 2. +1. TorchDataset - A wrapper for miscellaneous pytorch datasets (e.g. MNIST, CIFAR, etc) which extracts the images + and reformats them in a way that the DLAS trainer understands. +1. FullImageDataset - An image patch dataset where the patches are dynamically extracted from full-size images. I have + generally stopped using this for performance reasons and it should be considered deprecated. + +## Information about the "chunked" format + +This is the main format I have used in my experiments with image super resolution. It is fast to read and provides +rich metadata on the images that the patches are derived from, including a downsized "reference" fullsize image and +information on where the crop was taken from in the original image. + +### Creating a chunked dataset + +The file format for 'chunked' datasets is very particular. I recommend using `scripts/extract_subimages_with_ref.py` +to build these datasets from raw images. Here is how you would do that: + +1. Edit `scripts/extract_subimages_with_ref.py` to set these configuration options: + ``` + opt['input_folder'] = + opt['save_folder'] = + opt['crop_sz'] = [256, 512] # A list, the size of each sub-image that will be extracted and turned into patches. + opt['step'] = [128, 256] # The pixel distance the algorithm will step for each sub-image. If this is < crop_sz, patches will share image content. + opt['thres_sz'] = 128 # Amount of space that must be present on the edges of an image for it to be included in the image patch. Generally should be equal to the lowest step size. + opt['resize_final_img'] = [1, .5] # Reduction factor that will be applied to image patches at this crop_sz level. TODO: infer this. + opt['only_resize'] = False # If true, disables the patch-removal algorithm and just resizes the input images. + opt['vertical_split'] = False # Used for stereoscopic images. Not documented. + ``` + Note: the defaults should work fine for many applications. +1. Execute the script: `python scripts/extract_subimages_with_ref.py`. If you are having issues with imports, make sure + you set `PYTHONPATH` to the repo root. + +### Chunked cache + +To make trainer startup fast, the chunked datasets perform some preprocessing the first time they are loaded. The entire +dataset is scanned and a cache is built up and saved in cache.pth. Future invocations only need to load cache.pth on +startup, which greatly speeds up trainer startup when you are debugging issues. + +There is an important caveat here: this cache will not be recomputed unless you delete it. This means if you add new +images to your dataset, you must delete the cache for them to be picked up! Likewise, if you copy your dataset to a +new file path or a different computer, cache.pth must be deleted for it to work. In the latter case, you'll likely run +into some weird errors. + +### Details about the dataset format + +If you look inside of a dataset folder output by above, you'll see a list of folders. Each folder represents a single +image that was found by the script. + +Inside of that folder, you will see 3 different types of files: + +1. Image patches, each of which have a unique ID within the given set. These IDs do not necessarily need to be unique + across the entire dataset. +1. `centers.pt` A pytorch pickle which is just a dict that describes some metadata about the patches, like: where they + were located in the source image and their original width/height. +1. `ref.jpg` Is a square version of the original image that is downsampled to the patch size. \ No newline at end of file diff --git a/dlas/data/__init__.py b/dlas/data/__init__.py new file mode 100644 index 0000000..a65d453 --- /dev/null +++ b/dlas/data/__init__.py @@ -0,0 +1,137 @@ +"""create dataset and dataloader""" +import torch +import torch.utils.data +from munch import munchify + +from dlas.utils.util import opt_get + + +def create_dataloader(dataset, dataset_opt, opt=None, sampler=None, collate_fn=None, shuffle=True): + phase = dataset_opt['phase'] + pin_memory = opt_get(dataset_opt, ['pin_memory'], True) + if phase == 'train': + if opt_get(opt, ['dist'], False): + world_size = torch.distributed.get_world_size() + num_workers = dataset_opt['n_workers'] + assert dataset_opt['batch_size'] % world_size == 0 + batch_size = dataset_opt['batch_size'] // world_size + else: + num_workers = dataset_opt['n_workers'] + batch_size = dataset_opt['batch_size'] + return torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, + num_workers=num_workers, sampler=sampler, drop_last=True, + pin_memory=pin_memory, collate_fn=collate_fn, persistent_workers=True) + else: + batch_size = dataset_opt['batch_size'] or 1 + return torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=0, + pin_memory=pin_memory, collate_fn=collate_fn) + + +def create_dataset(dataset_opt, return_collate=False): + mode = dataset_opt['mode'] + collate = None + + # datasets for image restoration + if mode == 'fullimage': + from dlas.data.images.full_image_dataset import FullImageDataset as D + elif mode == 'single_image_extensible': + from dlas.data.images.single_image_dataset import \ + SingleImageDataset as D + elif mode == 'multi_frame_extensible': + from dlas.data.images.multi_frame_dataset import MultiFrameDataset as D + elif mode == 'combined': + from dlas.data.combined_dataset import CombinedDataset as D + elif mode == 'multiscale': + from dlas.data.images.multiscale_dataset import MultiScaleDataset as D + elif mode == 'paired_frame': + from dlas.data.images.paired_frame_dataset import \ + PairedFrameDataset as D + elif mode == 'stylegan2': + from dlas.data.images.stylegan2_dataset import Stylegan2Dataset as D + elif mode == 'imagefolder': + from dlas.data.images.image_folder_dataset import \ + ImageFolderDataset as D + elif mode == 'torch_dataset': + from data.torch_dataset import TorchDataset as D + elif mode == 'byol_dataset': + from dlas.data.images.byol_attachment import ByolDatasetWrapper as D + elif mode == 'byol_structured_dataset': + from dlas.data.images.byol_attachment import \ + StructuredCropDatasetWrapper as D + elif mode == 'random_aug_wrapper': + from dlas.data.images.byol_attachment import \ + DatasetRandomAugWrapper as D + elif mode == 'random_dataset': + from dlas.data.images.random_dataset import RandomDataset as D + elif mode == 'zipfile': + from dlas.data.images.zip_file_dataset import ZipFileDataset as D + elif mode == 'nv_tacotron': + from dlas.data.audio.nv_tacotron_dataset import TextMelCollate as C + from dlas.data.audio.nv_tacotron_dataset import TextWavLoader as D + from dlas.models.audio.tts.tacotron2 import create_hparams + default_params = create_hparams() + default_params.update(dataset_opt) + dataset_opt = munchify(default_params) + if opt_get(dataset_opt, ['needs_collate'], True): + collate = C() + elif mode == 'paired_voice_audio': + from dlas.data.audio.paired_voice_audio_dataset import \ + TextWavLoader as D + from dlas.models.audio.tts.tacotron2 import create_hparams + default_params = create_hparams() + default_params.update(dataset_opt) + dataset_opt = munchify(default_params) + elif mode == 'fast_paired_voice_audio': + from dlas.data.audio.fast_paired_dataset import \ + FastPairedVoiceDataset as D + from dlas.models.audio.tts.tacotron2 import create_hparams + default_params = create_hparams() + default_params.update(dataset_opt) + dataset_opt = munchify(default_params) + elif mode == 'fast_paired_voice_audio_with_phonemes': + from dlas.data.audio.fast_paired_dataset_with_phonemes import \ + FastPairedVoiceDataset as D + from dlas.models.audio.tts.tacotron2 import create_hparams + default_params = create_hparams() + default_params.update(dataset_opt) + dataset_opt = munchify(default_params) + elif mode == 'gpt_tts': + from dlas.data.audio.gpt_tts_dataset import GptTtsCollater as C + from dlas.data.audio.gpt_tts_dataset import GptTtsDataset as D + collate = C(dataset_opt) + elif mode == 'unsupervised_audio': + from dlas.data.audio.unsupervised_audio_dataset import \ + UnsupervisedAudioDataset as D + elif mode == 'unsupervised_audio_with_noise': + from dlas.data.audio.audio_with_noise_dataset import \ + AudioWithNoiseDataset as D + elif mode == 'preprocessed_mel': + from dlas.data.audio.preprocessed_mel_dataset import \ + PreprocessedMelDataset as D + elif mode == 'grand_conjoined_voice': + from dlas.data.audio.grand_conjoined_dataset import \ + GrandConjoinedDataset as D + from dlas.data.zero_pad_dict_collate import ZeroPadDictCollate as C + if opt_get(dataset_opt, ['needs_collate'], False): + collate = C() + else: + raise NotImplementedError( + 'Dataset [{:s}] is not recognized.'.format(mode)) + dataset = D(dataset_opt) + + if return_collate: + return dataset, collate + else: + return dataset + + +def get_dataset_debugger(dataset_opt): + mode = dataset_opt['mode'] + if mode == 'paired_voice_audio': + from dlas.data.audio.paired_voice_audio_dataset import \ + PairedVoiceDebugger + return PairedVoiceDebugger() + elif mode == 'fast_paired_voice_audio': + from dlas.data.audio.fast_paired_dataset import FastPairedVoiceDebugger + return FastPairedVoiceDebugger() + return None diff --git a/dlas/data/audio/__init__.py b/dlas/data/audio/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/dlas/data/audio/audio_with_noise_dataset.py b/dlas/data/audio/audio_with_noise_dataset.py new file mode 100644 index 0000000..16eb69e --- /dev/null +++ b/dlas/data/audio/audio_with_noise_dataset.py @@ -0,0 +1,246 @@ +import random +import sys +from math import pi + +import torch +import torch.nn.functional as F +import torchaudio +from torch.utils.data import Dataset +from tqdm import tqdm + +from dlas.data.audio.unsupervised_audio_dataset import ( + UnsupervisedAudioDataset, load_audio) +from dlas.data.util import (find_files_of_type, is_audio_file, + load_paths_from_cache) +from dlas.utils.util import opt_get + + +def _integration_fn_fully_enabled(n): + return torch.ones((n,)) + + +# Randomly assigns up to 5 blocks of the output tensor the value '1'. Rest is zero +def _integration_fn_spiky(n): + fn = torch.zeros((n,)) + spikes = random.randint(1, 5) + for _ in range(spikes): + sz = random.randint(n//8, n//2) + pos = random.randint(0, n) + extent = min(n, sz+pos) + fn[pos:extent] = 1 + return fn + + +# Uses a sinusoidal ramp up and down (of random length) to a peak which is held for a random duration. +def _integration_fn_smooth(n): + center = random.randint(1, n-2) + max_duration = n-center-1 + duration = random.randint(max_duration//4, max_duration) + end = center+duration + + ramp_up_sz = random.randint(n//16, n//4) + ramp_up = torch.sin(pi*torch.arange(0, ramp_up_sz)/(2*ramp_up_sz)) + if ramp_up_sz > center: + ramp_up = ramp_up[(ramp_up_sz-center):] + ramp_up_sz = center + + ramp_down_sz = random.randint(n//16, n//4) + ramp_down = torch.flip( + torch.sin(pi*torch.arange(0, ramp_down_sz)/(2*ramp_down_sz)), dims=[0]) + if ramp_down_sz > (n-end): + ramp_down = ramp_down[:(n-end)] + ramp_down_sz = n-end + + fn = torch.zeros((n,)) + fn[(center-ramp_up_sz):center] = ramp_up + fn[center:end] = 1 + fn[end:(end+ramp_down_sz)] = ramp_down + + return fn + + +def load_rir(path, sr, max_sz): + rir = load_audio(path, sr).abs() + if rir.shape[-1] > max_sz: + rir = rir[:, :max_sz] + rir = (rir / torch.norm(rir, p=2)).flip([1]) + return rir + + +''' +Wraps a unsupervised_audio_dataset and applies noise to the output clips, then provides labels depending on what +noise was added. +''' + + +class AudioWithNoiseDataset(Dataset): + def __init__(self, opt): + self.underlying_dataset = UnsupervisedAudioDataset(opt) + self.env_noise_paths = load_paths_from_cache( + opt['env_noise_paths'], opt['env_noise_cache']) + self.music_paths = load_paths_from_cache( + opt['music_paths'], opt['music_cache']) + self.openair_paths = find_files_of_type( + 'img', opt['openair_path'], qualifier=is_audio_file)[0] + self.min_volume = opt_get(opt, ['min_noise_volume'], .2) + self.max_volume = opt_get(opt, ['max_noise_volume'], .5) + self.sampling_rate = self.underlying_dataset.sampling_rate + self.use_gpu_for_reverb_compute = opt_get( + opt, ['use_gpu_for_reverb_compute'], True) + self.openair_kernels = None + self.current_item_fetch = 0 + self.fetch_error_count = 0 + + def load_openair_kernels(self): + if self.use_gpu_for_reverb_compute and self.openair_kernels is None: + # Load the openair reverbs as CUDA tensors. + self.openair_kernels = [] + for oa in self.openair_paths: + self.openair_kernels.append(load_rir( + oa, self.underlying_dataset.sampling_rate, self.underlying_dataset.sampling_rate*2).cuda()) + + def __getitem__(self, item): + if self.current_item_fetch != item: + self.current_item_fetch = item + self.fetch_error_count = 0 + + # Load on the fly to prevent GPU memory sharing across process errors. + self.load_openair_kernels() + + out = self.underlying_dataset[item] + clip = out['clip'] + dlen = clip.shape[-1] + clip = clip[:, :out['clip_lengths']] + padding_room = dlen - clip.shape[-1] + augpath = '' + augvol = 0 + try: + # Randomly adjust clip volume, regardless of the selection, between + clipvol = (random.random() * (.8-.5) + .5) + clip = clip * clipvol + + label = random.randint(0, 4) # Current excludes GSM corruption. + # label = 3 + if label > 0 and label < 4: # 0 is basically "leave it alone" + aug_needed = True + augvol = (random.random() * (self.max_volume - + self.min_volume) + self.min_volume) + if label == 1: + # Add environmental noise. + augpath = random.choice(self.env_noise_paths) + intg_fns = [_integration_fn_fully_enabled] + elif label == 2: + # Add music + augpath = random.choice(self.music_paths) + intg_fns = [_integration_fn_fully_enabled] + augvol *= .5 # Music is often severely in the background. + elif label == 3: + augpath = random.choice(self.underlying_dataset.audiopaths) + # This can take two forms: + if padding_room < 22000 or random.random() < .5: + # (1) The voices talk over one another. If there is no padding room, we always take this choice. + intg_fns = [_integration_fn_smooth, + _integration_fn_fully_enabled] + else: + # (2) There are simply two voices in the clip, separated from one another. + # This is a special case that does not use the same logic as the rest of the augmentations. + aug = load_audio( + augpath, self.underlying_dataset.sampling_rate) + # Pad with some random silence + aug = F.pad(aug, (random.randint(20, 4000), 0)) + # Fit what we can given the padding room we have. + aug = aug[:, :padding_room] + clip = torch.cat([clip, aug], dim=1) + # Restore some meta-parameters. + padding_room = dlen - clip.shape[-1] + out['clip_lengths'] = torch.tensor(clip.shape[-1]) + aug_needed = False + if aug_needed: + aug = load_audio( + augpath, self.underlying_dataset.sampling_rate) + if aug.shape[1] > clip.shape[1]: + n, cn = aug.shape[1], clip.shape[1] + gap = n-cn + placement = random.randint(0, gap) + aug = aug[:, placement:placement+cn] + aug = random.choice(intg_fns)(aug.shape[1]) * aug + aug = aug * augvol + if aug.shape[1] < clip.shape[1]: + gap = clip.shape[1] - aug.shape[1] + placement = random.randint(0, gap-1) + aug = torch.nn.functional.pad( + aug, (placement, gap-placement)) + clip = clip + aug + elif label == 4: + # Perform reverb (to simulate being in a large room with an omni-mic). This is performed by convolving + # impulse recordings from openair over the input clip. + if self.use_gpu_for_reverb_compute: + rir = random.choice(self.openair_kernels) + else: + augpath = random.choice(self.openair_paths) + rir = load_rir( + augpath, self.underlying_dataset.sampling_rate, clip.shape[-1]) + clip = torch.nn.functional.pad(clip, (rir.shape[1]-1, 0)) + if self.use_gpu_for_reverb_compute: + clip = clip.cuda() + clip = torch.nn.functional.conv1d( + clip.unsqueeze(0), rir.unsqueeze(0)).squeeze(0).cpu() + elif label == 5: + # Apply the GSM codec to simulate cellular phone audio. + clip = torchaudio.functional.apply_codec( + clip, self.underlying_dataset.sampling_rate, format="gsm") + except: + if self.fetch_error_count > 10: + print( + f"Exception encountered processing {item}, re-trying because this is often just a failed aug.") + print(sys.exc_info()) + # raise # Uncomment to surface exceptions. + self.fetch_error_count += 1 + return self[item] + + clip.clip_(-1, 1) + # Restore padding. + clip = F.pad(clip, (0, padding_room)) + out['clip'] = clip + out['label'] = label + # out['aug'] = aug + out['augpath'] = augpath + out['augvol'] = augvol + out['clipvol'] = clipvol + return out + + def __len__(self): + return len(self.underlying_dataset) + + +if __name__ == '__main__': + params = { + 'mode': 'unsupervised_audio_with_noise', + 'path': ['y:/clips/books1'], + 'cache_path': 'D:\\data\\clips_for_noise_classifier.pth', + 'sampling_rate': 22050, + 'pad_to_samples': 400000, + 'do_augmentation': False, + 'phase': 'train', + 'n_workers': 4, + 'batch_size': 256, + 'extra_samples': 0, + 'env_noise_paths': ['E:\\audio\\UrbanSound\\filtered', 'E:\\audio\\UrbanSound\\MSSND'], + 'env_noise_cache': 'E:\\audio\\UrbanSound\\cache.pth', + 'music_paths': ['E:\\audio\\music\\FMA\\fma_large', 'E:\\audio\\music\\maestro\\maestro-v3.0.0'], + 'music_cache': 'E:\\audio\\music\\cache.pth', + 'openair_path': 'D:\\data\\audio\\openair\\resampled', + 'use_gpu_for_reverb_compute': False, + } + from data import create_dataloader, create_dataset, util + + ds = create_dataset(params) + dl = create_dataloader(ds, params, pin_memory=False) + i = 0 + for b in tqdm(dl): + for b_ in range(b['clip'].shape[0]): + # torchaudio.save(f'{i}_clip_{b_}_{b["label"][b_].item()}.wav', b['clip'][b_][:, :b['clip_lengths'][b_]], ds.sampling_rate) + # torchaudio.save(f'{i}_clip_{b_}_aug.wav', b['aug'][b_], ds.sampling_rate) + print( + f'{i} aug path: {b["augpath"][b_]} aug volume: {b["augvol"][b_]} clip volume: {b["clipvol"][b_]}') + i += 1 diff --git a/dlas/data/audio/fast_paired_dataset.py b/dlas/data/audio/fast_paired_dataset.py new file mode 100644 index 0000000..45b3085 --- /dev/null +++ b/dlas/data/audio/fast_paired_dataset.py @@ -0,0 +1,336 @@ +import hashlib +import os +import random +import sys +import time +from itertools import groupby + +import torch +import torch.nn.functional as F +import torch.utils.data +import torchaudio +from tqdm import tqdm +from transformers import Wav2Vec2CTCTokenizer + +from dlas.data.audio.paired_voice_audio_dataset import CharacterTokenizer +from dlas.data.audio.unsupervised_audio_dataset import (load_audio, + load_similar_clips) +from dlas.utils.util import opt_get + + +def parse_tsv_aligned_codes(line, base_path): + fpt = line.strip().split('\t') + + def convert_string_list_to_tensor(strlist): + if strlist.startswith('['): + strlist = strlist[1:] + if strlist.endswith(']'): + strlist = strlist[:-1] + as_ints = [int(s) for s in strlist.split(', ')] + return torch.tensor(as_ints) + return os.path.join(base_path, f'{fpt[1]}'), fpt[0], convert_string_list_to_tensor(fpt[2]) + + +class FastPairedVoiceDataset(torch.utils.data.Dataset): + """ + This dataset is derived from paired_voice_audio, but it only supports loading from TSV files generated from the + ocotillo transcription engine, which includes alignment codes. To support the vastly larger TSV files, this dataset + uses an indexing mechanism which randomly selects offsets within the translation file to seek to. The data returned + is relative to these offsets. + + In practice, this means two things: + 1) Index {i} of this dataset means nothing: fetching from the same index will almost always return different data. + As a result, this dataset should not be used for validation or test runs. Use PairedVoiceAudio dataset instead. + 2) This dataset has a slight bias for items with longer text or longer filenames. + + The upshot is that this dataset loads extremely quickly and consumes almost no system memory. + """ + + def __init__(self, hparams): + self.paths = hparams['path'] + if not isinstance(self.paths, list): + self.paths = [self.paths] + self.paths_size_bytes = [os.path.getsize(p) for p in self.paths] + self.total_size_bytes = sum(self.paths_size_bytes) + self.types = opt_get(hparams, ['types'], [0 for _ in self.paths]) + + self.load_conditioning = opt_get(hparams, ['load_conditioning'], False) + self.conditioning_candidates = opt_get( + hparams, ['num_conditioning_candidates'], 1) + self.conditioning_length = opt_get( + hparams, ['conditioning_length'], 44100) + self.produce_ctc_metadata = opt_get( + hparams, ['produce_ctc_metadata'], False) + self.debug_failures = opt_get( + hparams, ['debug_loading_failures'], False) + self.text_cleaners = hparams.text_cleaners + self.sample_rate = hparams.sample_rate + self.aligned_codes_to_audio_ratio = 443 * self.sample_rate // 22050 + self.max_wav_len = opt_get(hparams, ['max_wav_length'], None) + self.load_aligned_codes = opt_get( + hparams, ['load_aligned_codes'], False) + if self.max_wav_len is not None: + self.max_aligned_codes = self.max_wav_len // self.aligned_codes_to_audio_ratio + self.max_text_len = opt_get(hparams, ['max_text_length'], None) + assert self.max_wav_len is not None and self.max_text_len is not None + self.use_bpe_tokenizer = opt_get(hparams, ['use_bpe_tokenizer'], False) + if self.use_bpe_tokenizer: + from dlas.data.audio.voice_tokenizer import VoiceBpeTokenizer + self.tokenizer = VoiceBpeTokenizer(opt_get( + hparams, ['tokenizer_vocab'], '../experiments/bpe_lowercase_asr_256.json')) + else: + self.tokenizer = CharacterTokenizer() + # records how many items are skipped when accessing an index. + self.skipped_items = 0 + + self.load_times = torch.zeros((256,)) + self.load_ind = 0 + + def get_wav_text_pair(self, audiopath_and_text): + # separate filename and text + audiopath, text = audiopath_and_text[0], audiopath_and_text[1] + text_seq = self.get_text(text) + wav = load_audio(audiopath, self.sample_rate) + return (text_seq, wav, text, audiopath_and_text[0]) + + def get_text(self, text): + tokens = self.tokenizer.encode(text) + tokens = torch.IntTensor(tokens) + if self.use_bpe_tokenizer: + # Assert if any UNK,start tokens encountered. + assert not torch.any(tokens == 1) + # The stop token should always be sacred. + assert not torch.any(tokens == 0) + return tokens + + def load_random_line(self, depth=0): + assert depth < 10 + + rand_offset = random.randint(0, self.total_size_bytes) + for i in range(len(self.paths)): + if rand_offset < self.paths_size_bytes[i]: + break + else: + rand_offset -= self.paths_size_bytes[i] + path = self.paths[i] + type = self.types[i] + with open(path, 'r', encoding='utf-8') as f: + f.seek(rand_offset) + # Read the rest of the line we seeked to, then the line after that. + try: # This can fail when seeking to a UTF-8 escape byte. + f.readline() + except: + # On failure, just recurse and try again. + return self.load_random_line(depth=depth + 1), type + l2 = f.readline() + + if l2: + try: + base_path = os.path.dirname(path) + return parse_tsv_aligned_codes(l2, base_path), type + except: + print(f"error parsing random offset: {sys.exc_info()}") + # On failure, just recurse and try again. + return self.load_random_line(depth=depth+1), type + + def get_ctc_metadata(self, codes): + grouped = groupby(codes.tolist()) + rcodes, repeats, seps = [], [], [0] + for val, group in grouped: + if val == 0: + # This is a very important distinction! It means the padding belongs to the character proceeding it. + seps[-1] = len(list(group)) + else: + rcodes.append(val) + repeats.append(len(list(group))) + seps.append(0) + + rcodes = torch.tensor(rcodes) + # These clip values are sane maximum values which I did not see in the datasets I have access to. + repeats = torch.clip(torch.tensor(repeats), min=1, max=30) + seps = torch.clip(torch.tensor(seps[:-1]), max=120) + + # Pad or clip the codes to get them to exactly self.max_text_len + orig_lens = rcodes.shape[0] + if rcodes.shape[0] < self.max_text_len: + gap = self.max_text_len - rcodes.shape[0] + rcodes = F.pad(rcodes, (0, gap)) + # The minimum value for repeats is 1, hence this is the pad value too. + repeats = F.pad(repeats, (0, gap), value=1) + seps = F.pad(seps, (0, gap)) + elif rcodes.shape[0] > self.max_text_len: + rcodes = rcodes[:self.max_text_len] + repeats = rcodes[:self.max_text_len] + seps = seps[:self.max_text_len] + + return { + 'ctc_raw_codes': rcodes, + 'ctc_separators': seps, + 'ctc_repeats': repeats, + 'ctc_raw_lengths': orig_lens, + } + + def __getitem__(self, index): + start = time.time() + self.skipped_items += 1 + apt, type = self.load_random_line() + try: + tseq, wav, text, path = self.get_wav_text_pair(apt) + if text is None or len(text.strip()) == 0: + raise ValueError + cond, cond_is_self = load_similar_clips(apt[0], self.conditioning_length, self.sample_rate, + n=self.conditioning_candidates) if self.load_conditioning else (None, False) + except: + if self.skipped_items > 100: + raise # Rethrow if we have nested too far. + if self.debug_failures: + print(f"error loading {apt[0]} {sys.exc_info()}") + return self[(index+1) % len(self)] + raw_codes = apt[2] + aligned_codes = raw_codes + + actually_skipped_items = self.skipped_items + self.skipped_items = 0 + if wav is None or \ + (self.max_wav_len is not None and wav.shape[-1] > self.max_wav_len) or \ + (self.max_text_len is not None and tseq.shape[0] > self.max_text_len): + # Basically, this audio file is nonexistent or too long to be supported by the dataset. + # It's hard to handle this situation properly. Best bet is to return the a random valid token and skew the dataset somewhat as a result. + if self.debug_failures: + print( + f"error loading {path}: ranges are out of bounds; {wav.shape[-1]}, {tseq.shape[0]}") + rv = random.randint(0, len(self)-1) + return self[rv] + orig_output = wav.shape[-1] + orig_text_len = tseq.shape[0] + orig_aligned_code_length = aligned_codes.shape[0] + if wav.shape[-1] != self.max_wav_len: + wav = F.pad(wav, (0, self.max_wav_len - wav.shape[-1])) + # These codes are aligned to audio inputs, so make sure to pad them as well. + aligned_codes = F.pad( + aligned_codes, (0, self.max_aligned_codes-aligned_codes.shape[0])) + if tseq.shape[0] != self.max_text_len: + tseq = F.pad(tseq, (0, self.max_text_len - tseq.shape[0])) + + elapsed = time.time() - start + self.load_times[self.load_ind] = elapsed + self.load_ind = (self.load_ind + 1) % len(self.load_times) + + res = { + 'real_text': text, + 'padded_text': tseq, + 'text_lengths': torch.tensor(orig_text_len, dtype=torch.long), + 'wav': wav, + 'wav_lengths': torch.tensor(orig_output, dtype=torch.long), + 'filenames': path, + 'skipped_items': actually_skipped_items, + 'load_time': self.load_times.mean(), + 'type': type, + } + if self.load_conditioning: + res['conditioning'] = cond + res['conditioning_contains_self'] = cond_is_self + if self.load_aligned_codes: + res['aligned_codes'] = aligned_codes + res['aligned_codes_lengths'] = orig_aligned_code_length + if self.produce_ctc_metadata: + res.update(self.get_ctc_metadata(raw_codes)) + + return res + + def __len__(self): + # 1000 cuts down a TSV file to the actual length pretty well. + return self.total_size_bytes // 1000 + + +class FastPairedVoiceDebugger: + def __init__(self): + self.total_items = 0 + self.loaded_items = 0 + self.self_conditioning_items = 0 + self.unique_files = set() + self.load_time = 0 + + def get_state(self): + return {'total_items': self.total_items, + 'loaded_items': self.loaded_items, + 'self_conditioning_items': self.self_conditioning_items} + + def load_state(self, state): + if isinstance(state, dict): + self.total_items = opt_get(state, ['total_items'], 0) + self.loaded_items = opt_get(state, ['loaded_items'], 0) + self.self_conditioning_items = opt_get( + state, ['self_conditioning_items'], 0) + + def update(self, batch): + self.total_items += batch['wav'].shape[0] + self.loaded_items += batch['skipped_items'].sum().item() + self.load_time = batch['load_time'].mean().item() + for filename in batch['filenames']: + self.unique_files.add(hashlib.sha256(filename.encode('utf-8'))) + if 'conditioning' in batch.keys(): + self.self_conditioning_items += batch['conditioning_contains_self'].sum( + ).item() + + def get_debugging_map(self): + return { + 'total_samples_loaded': self.total_items, + 'percent_skipped_samples': (self.loaded_items - self.total_items) / self.loaded_items, + 'percent_conditioning_is_self': self.self_conditioning_items / self.loaded_items, + 'unique_files_loaded': len(self.unique_files), + 'load_time': self.load_time, + } + + +if __name__ == '__main__': + batch_sz = 16 + params = { + 'mode': 'fast_paired_voice_audio', + 'path': ['y:/libritts/train-other-500/transcribed-oco.tsv', + 'y:/libritts/train-clean-100/transcribed-oco.tsv', + 'y:/libritts/train-clean-360/transcribed-oco.tsv', + 'y:/clips/books1/transcribed-oco.tsv', + 'y:/clips/books2/transcribed-oco.tsv', + 'y:/bigasr_dataset/hifi_tts/transcribed-oco.tsv', + 'y:/clips/podcasts-1/transcribed-oco.tsv',], + 'types': [0, 1, 1, 1, 2, 2, 0], + 'phase': 'train', + 'n_workers': 0, + 'batch_size': batch_sz, + 'max_wav_length': 220500, + 'max_text_length': 500, + 'sample_rate': 22050, + 'load_conditioning': True, + 'num_conditioning_candidates': 2, + 'conditioning_length': 102400, + 'use_bpe_tokenizer': True, + 'load_aligned_codes': True, + 'produce_ctc_metadata': True, + } + from data import create_dataloader, create_dataset + + def save(b, i, ib, key, c=None): + if c is not None: + torchaudio.save(f'{i}_clip_{ib}_{key}_{c}.wav', + b[key][ib][c], 22050) + else: + torchaudio.save(f'{i}_clip_{ib}_{key}.wav', b[key][ib], 22050) + + ds, c = create_dataset(params, return_collate=True) + dl = create_dataloader(ds, params, collate_fn=c) + i = 0 + m = None + max_pads, max_repeats = 0, 0 + for i, b in tqdm(enumerate(dl)): + for ib in range(batch_sz): + # max_pads = max(max_pads, b['ctc_pads'].max()) + # max_repeats = max(max_repeats, b['ctc_repeats'].max()) + print(f'{i} {ib} {b["real_text"][ib]}') + save(b, i, ib, 'wav') + save(b, i, ib, 'conditioning', 0) + save(b, i, ib, 'conditioning', 1) + pass + if i > 15: + break + print(max_pads, max_repeats) diff --git a/dlas/data/audio/fast_paired_dataset_with_phonemes.py b/dlas/data/audio/fast_paired_dataset_with_phonemes.py new file mode 100644 index 0000000..4a2bec2 --- /dev/null +++ b/dlas/data/audio/fast_paired_dataset_with_phonemes.py @@ -0,0 +1,348 @@ +import hashlib +import os +import random +import sys +import time +from itertools import groupby + +import torch +import torch.nn.functional as F +import torch.utils.data +import torchaudio +from tqdm import tqdm +from transformers import Wav2Vec2Processor + +from dlas.data.audio.paired_voice_audio_dataset import CharacterTokenizer +from dlas.data.audio.unsupervised_audio_dataset import (load_audio, + load_similar_clips) +from dlas.utils.util import opt_get + + +def parse_tsv_aligned_codes(line, base_path): + fpt = line.strip().split('\t') + + def convert_string_list_to_tensor(strlist): + if strlist.startswith('['): + strlist = strlist[1:] + if strlist.endswith(']'): + strlist = strlist[:-1] + as_ints = [int(s) for s in strlist.split(', ')] + return torch.tensor(as_ints) + return os.path.join(base_path, f'{fpt[1]}'), fpt[0], convert_string_list_to_tensor(fpt[2]) + + +class FastPairedVoiceDataset(torch.utils.data.Dataset): + """ + This dataset is derived from paired_voice_audio, but it only supports loading from TSV files generated from the + ocotillo transcription engine, which includes alignment codes. To support the vastly larger TSV files, this dataset + uses an indexing mechanism which randomly selects offsets within the translation file to seek to. The data returned + is relative to these offsets. + + In practice, this means two things: + 1) Index {i} of this dataset means nothing: fetching from the same index will almost always return different data. + As a result, this dataset should not be used for validation or test runs. Use PairedVoiceAudio dataset instead. + 2) This dataset has a slight bias for items with longer text or longer filenames. + + The upshot is that this dataset loads extremely quickly and consumes almost no system memory. + """ + + def __init__(self, hparams): + self.paths = hparams['path'] + phoneme_paths = hparams['phoneme_paths'] + self.paths = [(p, False) for p in self.paths] + [(p, True) + for p in phoneme_paths] + + self.paths_size_bytes = [os.path.getsize(p) for p, _ in self.paths] + self.total_size_bytes = sum(self.paths_size_bytes) + self.types = opt_get(hparams, ['types'], [0 for _ in self.paths]) + + self.normal_text_end_token = hparams['normal_text_end_token'] + self.load_conditioning = opt_get(hparams, ['load_conditioning'], False) + self.conditioning_candidates = opt_get( + hparams, ['num_conditioning_candidates'], 1) + self.conditioning_length = opt_get( + hparams, ['conditioning_length'], 44100) + self.produce_ctc_metadata = opt_get( + hparams, ['produce_ctc_metadata'], False) + self.debug_failures = opt_get( + hparams, ['debug_loading_failures'], False) + self.text_cleaners = hparams.text_cleaners + self.sample_rate = hparams.sample_rate + self.aligned_codes_to_audio_ratio = 443 * self.sample_rate // 22050 + self.max_wav_len = opt_get(hparams, ['max_wav_length'], None) + self.load_aligned_codes = opt_get( + hparams, ['load_aligned_codes'], False) + if self.max_wav_len is not None: + self.max_aligned_codes = self.max_wav_len // self.aligned_codes_to_audio_ratio + self.max_text_len = opt_get(hparams, ['max_text_length'], None) + assert self.max_wav_len is not None and self.max_text_len is not None + self.use_bpe_tokenizer = opt_get(hparams, ['use_bpe_tokenizer'], False) + if self.use_bpe_tokenizer: + from dlas.data.audio.voice_tokenizer import VoiceBpeTokenizer + self.tokenizer = VoiceBpeTokenizer(opt_get( + hparams, ['tokenizer_vocab'], '../experiments/bpe_lowercase_asr_256.json')) + else: + self.tokenizer = CharacterTokenizer() + self.ipa_phoneme_tokenizer = Wav2Vec2Processor.from_pretrained( + "facebook/wav2vec2-lv-60-espeak-cv-ft").tokenizer + self.ipa_phoneme_tokenizer.do_phonemize = False + # records how many items are skipped when accessing an index. + self.skipped_items = 0 + + self.load_times = torch.zeros((256,)) + self.load_ind = 0 + + def get_wav_text_pair(self, audiopath_and_text, is_phonetic): + # separate filename and text + audiopath, text = audiopath_and_text[0], audiopath_and_text[1] + text_seq = self.get_text(text, is_phonetic) + wav = load_audio(audiopath, self.sample_rate) + return (text_seq, wav, text, audiopath_and_text[0]) + + def get_text(self, text, is_phonetic): + if is_phonetic: + tokens = self.ipa_phoneme_tokenizer.encode(text) + else: + tokens = self.tokenizer.encode(text) + tokens = torch.IntTensor(tokens) + if self.use_bpe_tokenizer: + # Assert if any UNK,start tokens encountered. + assert not torch.any(tokens == 1) + # The stop token should always be sacred. + assert not torch.any(tokens == 0) + return tokens + + def load_random_line(self, depth=0): + assert depth < 10 + + rand_offset = random.randint(0, self.total_size_bytes) + for i in range(len(self.paths)): + if rand_offset < self.paths_size_bytes[i]: + break + else: + rand_offset -= self.paths_size_bytes[i] + path, is_phonetic = self.paths[i] + type = self.types[i] + with open(path, 'r', encoding='utf-8') as f: + f.seek(rand_offset) + # Read the rest of the line we seeked to, then the line after that. + try: # This can fail when seeking to a UTF-8 escape byte. + f.readline() + except: + # On failure, just recurse and try again. + return self.load_random_line(depth=depth + 1) + l2 = f.readline() + + if l2: + try: + base_path = os.path.dirname(path) + return parse_tsv_aligned_codes(l2, base_path), type, is_phonetic + except: + print(f"error parsing random offset: {sys.exc_info()}") + # On failure, just recurse and try again. + return self.load_random_line(depth=depth+1) + + def get_ctc_metadata(self, codes): + grouped = groupby(codes.tolist()) + rcodes, repeats, seps = [], [], [0] + for val, group in grouped: + if val == 0: + # This is a very important distinction! It means the padding belongs to the character proceeding it. + seps[-1] = len(list(group)) + else: + rcodes.append(val) + repeats.append(len(list(group))) + seps.append(0) + + rcodes = torch.tensor(rcodes) + # These clip values are sane maximum values which I did not see in the datasets I have access to. + repeats = torch.clip(torch.tensor(repeats), min=1, max=30) + seps = torch.clip(torch.tensor(seps[:-1]), max=120) + + # Pad or clip the codes to get them to exactly self.max_text_len + orig_lens = rcodes.shape[0] + if rcodes.shape[0] < self.max_text_len: + gap = self.max_text_len - rcodes.shape[0] + rcodes = F.pad(rcodes, (0, gap)) + # The minimum value for repeats is 1, hence this is the pad value too. + repeats = F.pad(repeats, (0, gap), value=1) + seps = F.pad(seps, (0, gap)) + elif rcodes.shape[0] > self.max_text_len: + rcodes = rcodes[:self.max_text_len] + repeats = rcodes[:self.max_text_len] + seps = seps[:self.max_text_len] + return { + 'ctc_raw_codes': rcodes, + 'ctc_separators': seps, + 'ctc_repeats': repeats, + 'ctc_raw_lengths': orig_lens, + } + + def __getitem__(self, index): + start = time.time() + self.skipped_items += 1 + apt, type, is_phonetic = self.load_random_line() + try: + tseq, wav, text, path = self.get_wav_text_pair(apt, is_phonetic) + if text is None or len(text.strip()) == 0: + raise ValueError + cond, cond_is_self = load_similar_clips(apt[0], self.conditioning_length, self.sample_rate, + n=self.conditioning_candidates) if self.load_conditioning else (None, False) + except: + if self.skipped_items > 100: + raise # Rethrow if we have nested too far. + if self.debug_failures: + print(f"error loading {apt[0]} {sys.exc_info()}") + return self[(index+1) % len(self)] + raw_codes = apt[2] + aligned_codes = raw_codes + + actually_skipped_items = self.skipped_items + self.skipped_items = 0 + if wav is None or \ + (self.max_wav_len is not None and wav.shape[-1] > self.max_wav_len) or \ + (self.max_text_len is not None and tseq.shape[0] > self.max_text_len): + # Basically, this audio file is nonexistent or too long to be supported by the dataset. + # It's hard to handle this situation properly. Best bet is to return the a random valid token and skew the dataset somewhat as a result. + if self.debug_failures: + print( + f"error loading {path}: ranges are out of bounds; {wav.shape[-1]}, {tseq.shape[0]}") + rv = random.randint(0, len(self)-1) + return self[rv] + + # Shift phonetic token and aligned_code tokens over. + if is_phonetic: + tseq = tseq + self.normal_text_end_token + # But keep the padding/stop tokens. + if self.load_aligned_codes: + aligned_codes = aligned_codes + self.normal_text_end_token + + orig_output = wav.shape[-1] + orig_text_len = tseq.shape[0] + orig_aligned_code_length = aligned_codes.shape[0] + if wav.shape[-1] != self.max_wav_len: + wav = F.pad(wav, (0, self.max_wav_len - wav.shape[-1])) + # These codes are aligned to audio inputs, so make sure to pad them as well. + aligned_codes = F.pad( + aligned_codes, (0, self.max_aligned_codes-aligned_codes.shape[0])) + if tseq.shape[0] != self.max_text_len: + tseq = F.pad(tseq, (0, self.max_text_len - tseq.shape[0])) + + elapsed = time.time() - start + self.load_times[self.load_ind] = elapsed + self.load_ind = (self.load_ind + 1) % len(self.load_times) + + res = { + 'real_text': text, + 'padded_text': tseq, + 'text_lengths': torch.tensor(orig_text_len, dtype=torch.long), + 'wav': wav, + 'wav_lengths': torch.tensor(orig_output, dtype=torch.long), + 'filenames': path, + 'skipped_items': actually_skipped_items, + 'load_time': self.load_times.mean(), + 'type': type, + } + if self.load_conditioning: + res['conditioning'] = cond + res['conditioning_contains_self'] = cond_is_self + if self.load_aligned_codes: + res['aligned_codes']: aligned_codes + res['aligned_codes_lengths']: orig_aligned_code_length + if self.produce_ctc_metadata: + res.update(self.get_ctc_metadata(raw_codes)) + + return res + + def __len__(self): + # 1000 cuts down a TSV file to the actual length pretty well. + return self.total_size_bytes // 1000 + + +class FastPairedVoiceDebugger: + def __init__(self): + self.total_items = 0 + self.loaded_items = 0 + self.self_conditioning_items = 0 + self.unique_files = set() + self.load_time = 0 + + def get_state(self): + return {'total_items': self.total_items, + 'loaded_items': self.loaded_items, + 'self_conditioning_items': self.self_conditioning_items} + + def load_state(self, state): + if isinstance(state, dict): + self.total_items = opt_get(state, ['total_items'], 0) + self.loaded_items = opt_get(state, ['loaded_items'], 0) + self.self_conditioning_items = opt_get( + state, ['self_conditioning_items'], 0) + + def update(self, batch): + self.total_items += batch['wav'].shape[0] + self.loaded_items += batch['skipped_items'].sum().item() + self.load_time = batch['load_time'].mean().item() + for filename in batch['filenames']: + self.unique_files.add(hashlib.sha256(filename.encode('utf-8'))) + if 'conditioning' in batch.keys(): + self.self_conditioning_items += batch['conditioning_contains_self'].sum( + ).item() + + def get_debugging_map(self): + return { + 'total_samples_loaded': self.total_items, + 'percent_skipped_samples': (self.loaded_items - self.total_items) / self.loaded_items, + 'percent_conditioning_is_self': self.self_conditioning_items / self.loaded_items, + 'unique_files_loaded': len(self.unique_files), + 'load_time': self.load_time, + } + + +if __name__ == '__main__': + batch_sz = 16 + params = { + 'mode': 'fast_paired_voice_audio_with_phonemes', + 'path': ['y:/libritts/train-clean-100/transcribed-oco.tsv',], + 'phoneme_paths': ['y:/libritts/train-other-500/transcribed-phoneme-oco.tsv'], + 'types': [0, 0], + 'normal_text_end_token': 256, + 'phase': 'train', + 'n_workers': 0, + 'batch_size': batch_sz, + 'max_wav_length': 220500, + 'max_text_length': 500, + 'sample_rate': 22050, + 'load_conditioning': True, + 'num_conditioning_candidates': 2, + 'conditioning_length': 102400, + 'use_bpe_tokenizer': True, + 'load_aligned_codes': False, + 'debug_loading_failures': True, + } + from data import create_dataloader, create_dataset + + def save(b, i, ib, key, c=None): + if c is not None: + torchaudio.save(f'{i}_clip_{ib}_{key}_{c}.wav', + b[key][ib][c], 22050) + else: + torchaudio.save(f'{i}_clip_{ib}_{key}.wav', b[key][ib], 22050) + + ds, c = create_dataset(params, return_collate=True) + dl = create_dataloader(ds, params, collate_fn=c) + i = 0 + m = None + max_pads, max_repeats = 0, 0 + for i, b in tqdm(enumerate(dl)): + for ib in range(batch_sz): + # max_pads = max(max_pads, b['ctc_pads'].max()) + # max_repeats = max(max_repeats, b['ctc_repeats'].max()) + print(f'{i} {ib} {b["real_text"][ib]}') + # save(b, i, ib, 'wav') + # save(b, i, ib, 'conditioning', 0) + # save(b, i, ib, 'conditioning', 1) + pass + if i > 15: + break + print(max_pads, max_repeats) diff --git a/dlas/data/audio/gpt_tts_dataset.py b/dlas/data/audio/gpt_tts_dataset.py new file mode 100644 index 0000000..35c83f0 --- /dev/null +++ b/dlas/data/audio/gpt_tts_dataset.py @@ -0,0 +1,113 @@ +import os + +import torch +import torch.nn.functional as F +import torch.utils.data +from torch import LongTensor +from tqdm import tqdm + +from dlas.models.audio.tts.tacotron2 import (load_filepaths_and_text, symbols, + text_to_sequence) + + +class GptTtsDataset(torch.utils.data.Dataset): + MAX_SYMBOLS_PER_PHRASE = 200 + NUMBER_SYMBOLS = len(symbols) + NUMBER_TEXT_TOKENS = NUMBER_SYMBOLS + MAX_SYMBOLS_PER_PHRASE + 2 + TEXT_START_TOKEN = LongTensor([NUMBER_TEXT_TOKENS-1]) + TEXT_STOP_TOKEN = LongTensor([NUMBER_TEXT_TOKENS-2]) + + def __init__(self, opt): + self.path = os.path.dirname(opt['path']) + self.audiopaths_and_text = load_filepaths_and_text(opt['path']) + self.text_cleaners = ['english_cleaners'] + + self.MEL_DICTIONARY_SIZE = opt['mel_vocab_size']+3 + self.MEL_START_TOKEN = LongTensor([self.MEL_DICTIONARY_SIZE-3]) + self.MEL_STOP_TOKEN = LongTensor([self.MEL_DICTIONARY_SIZE-2]) + + def __getitem__(self, index): + # Fetch text and add start/stop tokens. + audiopath_and_text = self.audiopaths_and_text[index] + audiopath, text = audiopath_and_text[0], audiopath_and_text[1] + text = torch.IntTensor(text_to_sequence(text, self.text_cleaners)) + text = torch.cat([self.TEXT_START_TOKEN, text, + self.TEXT_STOP_TOKEN], dim=0) + + # Fetch quantized MELs + quant_path = audiopath.replace('wavs/', 'quantized_mels/') + '.pth' + filename = os.path.join(self.path, quant_path) + qmel = torch.load(filename) + qmel = torch.cat([self.MEL_START_TOKEN, qmel, self.MEL_STOP_TOKEN]) + + return text, qmel, audiopath + + def __len__(self): + return len(self.audiopaths_and_text) + + +class GptTtsCollater(): + MAX_SYMBOLS_PER_PHRASE = 200 + NUMBER_SYMBOLS = len(symbols) + NUMBER_TEXT_TOKENS = NUMBER_SYMBOLS + MAX_SYMBOLS_PER_PHRASE + 2 + + def __init__(self, opt): + self.MEL_DICTIONARY_SIZE = opt['mel_vocab_size']+3 + self.MEL_PAD_TOKEN = self.MEL_DICTIONARY_SIZE-1 + + def __call__(self, batch): + text_lens = [len(x[0]) for x in batch] + # max_text_len = max(text_lens) + # This forces all outputs to have the full 200 characters. Testing if this makes a difference. + max_text_len = self.MAX_SYMBOLS_PER_PHRASE + mel_lens = [len(x[1]) for x in batch] + max_mel_len = max(mel_lens) + texts = [] + qmels = [] + # This is the sequential "background" tokens that are used as padding for text tokens, as specified in the DALLE paper. + text_range_embedding = torch.arange(max_text_len) + self.NUMBER_SYMBOLS + for b in batch: + text, qmel, _ = b + text = F.pad(text, (0, max_text_len-len(text)), value=0) + text = torch.where(text == 0, text_range_embedding, text) + texts.append(text) + qmels.append(F.pad(qmel, (0, max_mel_len-len(qmel)), + value=self.MEL_PAD_TOKEN)) + + filenames = [j[2] for j in batch] + + padded_qmel_gt = torch.stack(qmels)[:, 1:-1] + padded_qmel_gt = padded_qmel_gt * (padded_qmel_gt < 512) + + return { + 'padded_text': torch.stack(texts), + 'input_lengths': LongTensor(text_lens), + 'padded_qmel': torch.stack(qmels), + 'padded_qmel_gt': padded_qmel_gt, + 'output_lengths': LongTensor(mel_lens), + 'filenames': filenames + } + + +if __name__ == '__main__': + params = { + 'mode': 'gpt_tts', + 'path': 'E:\\audio\\LJSpeech-1.1\\ljs_audio_text_train_filelist.txt', + 'phase': 'train', + 'n_workers': 0, + 'batch_size': 16, + 'mel_vocab_size': 512, + } + from data import create_dataloader, create_dataset + + ds, c = create_dataset(params, return_collate=True) + dl = create_dataloader(ds, params, collate_fn=c) + i = 0 + m = [] + max_text = 0 + max_mel = 0 + for b in tqdm(dl): + max_mel = max(max_mel, b['padded_qmel'].shape[2]) + max_text = max(max_text, b['padded_text'].shape[1]) + m = torch.stack(m) + print(m.mean(), m.std()) diff --git a/dlas/data/audio/gpt_tts_tokenizer.json b/dlas/data/audio/gpt_tts_tokenizer.json new file mode 100644 index 0000000..a128f27 --- /dev/null +++ b/dlas/data/audio/gpt_tts_tokenizer.json @@ -0,0 +1 @@ +{"version":"1.0","truncation":null,"padding":null,"added_tokens":[{"id":0,"special":true,"content":"[STOP]","single_word":false,"lstrip":false,"rstrip":false,"normalized":false},{"id":1,"special":true,"content":"[UNK]","single_word":false,"lstrip":false,"rstrip":false,"normalized":false},{"id":2,"special":true,"content":"[SPACE]","single_word":false,"lstrip":false,"rstrip":false,"normalized":false}],"normalizer":null,"pre_tokenizer":{"type":"Whitespace"},"post_processor":null,"decoder":null,"model":{"type":"BPE","dropout":null,"unk_token":"[UNK]","continuing_subword_prefix":null,"end_of_word_suffix":null,"fuse_unk":false,"vocab":{"[STOP]":0,"[UNK]":1,"[SPACE]":2,"!":3,"'":4,"(":5,")":6,",":7,"-":8,".":9,"/":10,":":11,";":12,"?":13,"a":14,"b":15,"c":16,"d":17,"e":18,"f":19,"g":20,"h":21,"i":22,"j":23,"k":24,"l":25,"m":26,"n":27,"o":28,"p":29,"q":30,"r":31,"s":32,"t":33,"u":34,"v":35,"w":36,"x":37,"y":38,"z":39,"th":40,"in":41,"the":42,"an":43,"er":44,"ou":45,"re":46,"on":47,"at":48,"ed":49,"en":50,"to":51,"ing":52,"and":53,"is":54,"as":55,"al":56,"or":57,"of":58,"ar":59,"it":60,"es":61,"he":62,"st":63,"le":64,"om":65,"se":66,"be":67,"ad":68,"ow":69,"ly":70,"ch":71,"wh":72,"that":73,"you":74,"li":75,"ve":76,"ac":77,"ti":78,"ld":79,"me":80,"was":81,"gh":82,"id":83,"ll":84,"wi":85,"ent":86,"for":87,"ay":88,"ro":89,"ver":90,"ic":91,"her":92,"ke":93,"his":94,"no":95,"ut":96,"un":97,"ir":98,"lo":99,"we":100,"ri":101,"ha":102,"with":103,"ght":104,"out":105,"im":106,"ion":107,"all":108,"ab":109,"one":110,"ne":111,"ge":112,"ould":113,"ter":114,"mo":115,"had":116,"ce":117,"she":118,"go":119,"sh":120,"ur":121,"am":122,"so":123,"pe":124,"my":125,"de":126,"are":127,"but":128,"ome":129,"fr":130,"ther":131,"fe":132,"su":133,"do":134,"con":135,"te":136,"ain":137,"ere":138,"po":139,"if":140,"they":141,"us":142,"ag":143,"tr":144,"now":145,"oun":146,"this":147,"have":148,"not":149,"sa":150,"il":151,"up":152,"thing":153,"from":154,"ap":155,"him":156,"ack":157,"ation":158,"ant":159,"our":160,"op":161,"like":162,"ust":163,"ess":164,"bo":165,"ok":166,"ul":167,"ind":168,"ex":169,"com":170,"some":171,"there":172,"ers":173,"co":174,"res":175,"man":176,"ard":177,"pl":178,"wor":179,"way":180,"tion":181,"fo":182,"ca":183,"were":184,"by":185,"ate":186,"pro":187,"ted":188,"ound":189,"own":190,"would":191,"ts":192,"what":193,"qu":194,"ally":195,"ight":196,"ck":197,"gr":198,"when":199,"ven":200,"can":201,"ough":202,"ine":203,"end":204,"per":205,"ous":206,"od":207,"ide":208,"know":209,"ty":210,"very":211,"si":212,"ak":213,"who":214,"about":215,"ill":216,"them":217,"est":218,"red":219,"ye":220,"could":221,"ong":222,"your":223,"their":224,"em":225,"just":226,"other":227,"into":228,"any":229,"whi":230,"um":231,"tw":232,"ast":233,"der":234,"did":235,"ie":236,"been":237,"ace":238,"ink":239,"ity":240,"back":241,"ting":242,"br":243,"more":244,"ake":245,"pp":246,"then":247,"sp":248,"el":249,"use":250,"bl":251,"said":252,"over":253,"get":254},"merges":["t h","i n","th e","a n","e r","o u","r e","o n","a t","e d","e n","t o","in g","an d","i s","a s","a l","o r","o f","a r","i t","e s","h e","s t","l e","o m","s e","b e","a d","o w","l y","c h","w h","th at","y ou","l i","v e","a c","t i","l d","m e","w as","g h","i d","l l","w i","en t","f or","a y","r o","v er","i c","h er","k e","h is","n o","u t","u n","i r","l o","w e","r i","h a","wi th","gh t","ou t","i m","i on","al l","a b","on e","n e","g e","ou ld","t er","m o","h ad","c e","s he","g o","s h","u r","a m","s o","p e","m y","d e","a re","b ut","om e","f r","the r","f e","s u","d o","c on","t e","a in","er e","p o","i f","the y","u s","a g","t r","n ow","ou n","th is","ha ve","no t","s a","i l","u p","th ing","fr om","a p","h im","ac k","at ion","an t","ou r","o p","li ke","u st","es s","b o","o k","u l","in d","e x","c om","s ome","the re","er s","c o","re s","m an","ar d","p l","w or","w ay","ti on","f o","c a","w ere","b y","at e","p ro","t ed","oun d","ow n","w ould","t s","wh at","q u","al ly","i ght","c k","g r","wh en","v en","c an","ou gh","in e","en d","p er","ou s","o d","id e","k now","t y","ver y","s i","a k","wh o","ab out","i ll","the m","es t","re d","y e","c ould","on g","you r","the ir","e m","j ust","o ther","in to","an y","wh i","u m","t w","as t","d er","d id","i e","be en","ac e","in k","it y","b ack","t ing","b r","mo re","a ke","p p","the n","s p","e l","u se","b l","sa id","o ver","ge t"]}} \ No newline at end of file diff --git a/dlas/data/audio/grand_conjoined_dataset.py b/dlas/data/audio/grand_conjoined_dataset.py new file mode 100644 index 0000000..7c24f6c --- /dev/null +++ b/dlas/data/audio/grand_conjoined_dataset.py @@ -0,0 +1,246 @@ +import os + +import torch +import torch.nn.functional as F +import torch.utils.data +import torchaudio +from munch import munchify +from tqdm import tqdm + +from dlas.data.audio.unsupervised_audio_dataset import UnsupervisedAudioDataset +from dlas.data.text.hf_datasets_wrapper import HfDataset +from dlas.utils.util import opt_get + + +def build_paired_voice_dataset(args): + from models.audio.tts.tacotron2 import create_hparams + + from dlas.data.audio.paired_voice_audio_dataset import TextWavLoader as D + default_params = create_hparams() + default_params.update(args) + dataset_opt = munchify(default_params) + return D(dataset_opt) + + +class GrandConjoinedDataset(torch.utils.data.Dataset): + """ + A joint text & speech dataset that joins three separate datasets into a single batch: + 1. Unpaired text + 2. Unpaired speech + 3. Paired speech & text + + Supports situations where the underlying data sources for these three elements are differently sized, e.g. you can + have a massive text corpus of 1B elements, a smaller unpaired speech corpus, and a small paired speech<->text corpus. + + Performs tokenization at this level, ignoring any tokenization performed by upstream datasets. + """ + + def __init__(self, opt): + sample_rate = 22050 # Fixed. + paired_dataset_args = opt['paired_dataset_args'] + self.only_paired = opt_get(opt, ['only_paired'], False) + if not self.only_paired: + unsupervised_audio_args = opt['unsupervised_audio_args'] + text_corpus_args = opt['text_corpus_args'] + + self.max_paired_audio_length = opt['max_paired_audio_length'] + self.max_paired_text_length = opt['max_paired_text_length'] + self.max_solo_audio_length = opt['max_solo_audio_length'] + self.max_solo_text_length = opt['max_solo_text_length'] + self.collate = opt_get(opt, ['needs_collate'], False) + self.sample_rate = sample_rate + self.num_conditioning_candidates = opt_get( + opt, ['num_conditioning_candidates'], 0) + self.conditioning_length = opt_get(opt, ['conditioning_length'], 44000) + load_conditioning = self.num_conditioning_candidates > 0 + + # Set some sane arguments for all three datasets. + paired_dataset_args['needs_collate'] = self.collate + paired_dataset_args['load_conditioning'] = load_conditioning + paired_dataset_args['num_conditioning_candidates'] = self.num_conditioning_candidates + paired_dataset_args['conditioning_length'] = self.conditioning_length + paired_dataset_args['sample_rate'] = sample_rate + paired_dataset_args['max_wav_length'] = self.max_paired_audio_length + paired_dataset_args['max_text_length'] = self.max_paired_text_length + self.speech_and_text = build_paired_voice_dataset(paired_dataset_args) + + if not self.only_paired: + unsupervised_audio_args['sampling_rate'] = sample_rate + unsupervised_audio_args['do_augmentation'] = False + unsupervised_audio_args['resample_clip'] = False + unsupervised_audio_args['extra_samples'] = self.num_conditioning_candidates + unsupervised_audio_args['extra_sample_length'] = self.conditioning_length + if not self.collate: + unsupervised_audio_args['pad_to_samples'] = self.max_solo_audio_length + self.speech = UnsupervisedAudioDataset(unsupervised_audio_args) + self.text = HfDataset(**text_corpus_args) + + def fetch_text_at(self, i): + try: + txt = self.text[i % len(self.text)]['text'] + # This is a hack to get around the use of '*' to mask expletives in some text-only datasets. There really isn't a linguistic use for this character anyways. + assert '*' not in txt + tok = self.speech_and_text.get_text(txt) + padding_required = self.max_solo_text_length - tok.shape[0] + if padding_required < 0: + # Just truncate since there is no conditioning required. + tok = tok[:self.max_solo_text_length] + elif padding_required > 0: + tok = F.pad(tok, (0, padding_required)) + return txt, tok + except: + # This is fully expected: there are a lot of text strings we intentionally do not + # handle (e.g. ones with emojis, or other languages). Just return another one. + return self.fetch_text_at((i+1) % len(self.text)) + + def fetch_snt_at(self, i): + fetched = self.speech_and_text[i % len(self.speech_and_text)] + if self.collate: + tseq, wav, path, text, cond = fetched + res = { + 'real_text': text, + 'padded_text': tseq, + 'text_lengths': torch.tensor(tseq.shape[0], dtype=torch.long), + 'wav': wav, + 'wav_lengths': torch.tensor(wav.shape[-1], dtype=torch.long), + 'filenames': path + } + if self.num_conditioning_candidates > 0: + res['conditioning'] = cond + return res + else: + return fetched + + def optionally_add_conditioning_candidates(self, res, paired, solo=None): + if self.num_conditioning_candidates > 0: + if solo is None: + res['paired_audio_conditioning'] = paired['conditioning'] + res['speech_audio_conditioning'] = paired['conditioning'] + else: + res['paired_audio_conditioning'] = paired['conditioning'] + res['speech_audio_conditioning'] = solo['alt_clips'] + return res + + def __getitem__(self, i): + snt = self.fetch_snt_at(i) + if self.only_paired: + return self.optionally_add_conditioning_candidates({ + 'paired_audio': snt['wav'], + 'paired_audio_lengths': snt['wav_lengths'], + 'paired_text': snt['real_text'], + 'paired_text_tokens': snt['padded_text'], + 'paired_file': snt['filenames'], + 'speech_audio': snt['wav'], + 'speech_audio_lengths': snt['wav_lengths'], + 'speech_file': snt['filenames'], + 'text_text': snt['real_text'], + 'text_tokens': snt['padded_text'], + }, snt) + else: + txt, txt_tok = self.fetch_text_at(i % len(self.text)) + sp = self.speech[i % len(self.speech)] + # Set upper bound on solo speech lengths. This is handled automatically when collation is turned off, but needs to be done otherwise. + sp['clip'] = sp['clip'][:, :self.max_solo_audio_length] + sp['clip_lengths'] = sp['clip_lengths'].clamp( + 0, self.max_solo_audio_length) + return self.optionally_add_conditioning_candidates({ + 'paired_audio': snt['wav'], + 'paired_audio_lengths': snt['wav_lengths'], + 'paired_text': snt['real_text'], + 'paired_text_tokens': snt['padded_text'], + 'paired_file': snt['filenames'], + 'speech_audio': sp['clip'], + 'speech_audio_lengths': sp['clip_lengths'], + 'speech_file': sp['path'], + 'text_text': txt, + 'text_tokens': txt_tok, + }, snt, sp) + + def __len__(self): + if self.only_paired: + return len(self.speech_and_text) + else: + return max(len(self.speech), len(self.speech_and_text), len(self.text)) + + +if __name__ == '__main__': + batch_sz = 8 + train_params = { + 'mode': 'grand_conjoined_voice', + 'phase': 'train', + 'n_workers': 0, + 'batch_size': batch_sz, + + 'max_paired_audio_length': 255995, + 'max_paired_text_length': 100, + 'max_solo_text_length': 200, + 'max_solo_audio_length': 307195, + 'num_conditioning_candidates': 1, + 'conditioning_length': 44000, + 'needs_collate': True, + 'paired_dataset_args': { + 'path': ['Z:\\bigasr_dataset\\tedlium\\train-all.txt'], + 'fetcher_mode': ['libritts'], + 'use_bpe_tokenizer': False, + }, + 'unsupervised_audio_args': { + 'path': ['Y:\\clips\\podcasts-0\\6175_20170425-How the National Security Council Works'], + 'cache_path': 'test_cache_delete_me2.pth', + }, + 'text_corpus_args': { + 'corpi': [['bookcorpus', '']], + 'cache_path': 'Z:\\huggingface_datasets\\cache', + }, + } + val_params = { + 'mode': 'grand_conjoined_voice', + 'phase': 'val', + 'n_workers': 0, + 'batch_size': batch_sz, + + 'max_paired_audio_length': 255995, + 'max_paired_text_length': 200, + 'max_solo_text_length': 330, + 'max_solo_audio_length': 300000, + 'only_paired': True, + 'needs_collate': False, + 'paired_dataset_args': { + 'path': ['Z:\\bigasr_dataset\\libritts\\test-clean_list.txt'], + 'fetcher_mode': ['libritts'], + 'use_bpe_tokenizer': False, + }, + } + from data import create_dataloader, create_dataset + os.remove('test_cache_delete_me2.pth') + + ds, c = create_dataset(train_params, return_collate=True) + dl = create_dataloader(ds, train_params, collate_fn=c) + + def save(b, i, ib, key, c=None): + if c is not None: + torchaudio.save(f'{i}_clip_{ib}_{key}_{c}.wav', + b[key][ib][c], 22050) + else: + torchaudio.save(f'{i}_clip_{ib}_{key}.wav', b[key][ib], 22050) + + def decode(b, ib, key): + return ds.speech_and_text.tokenizer.decode(b[key][ib].cpu().numpy()) + + i = 0 + m = None + for i, b in tqdm(enumerate(dl)): + for ib in range(batch_sz): + # save(b, i, ib, 'paired_audio') + # save(b, i, ib, 'paired_audio_conditioning', 0) + # save(b, i, ib, 'paired_audio_conditioning', 1) + print( + f'Paired file: {b["paired_file"][ib]} text: {b["paired_text"][ib]}') + print( + f'Paired text decoded: {decode(b, ib, "paired_text_tokens")}') + # save(b, i, ib, 'speech_audio') + # save(b, i, ib, 'speech_audio_conditioning', 0) + # save(b, i, ib, 'speech_audio_conditioning', 1) + # print(f'Text: {b["text_text"][ib]}') + # print(f'Text decoded: {decode(b, ib, "text_tokens")}') + if i > 5: + break diff --git a/dlas/data/audio/nv_tacotron_dataset.py b/dlas/data/audio/nv_tacotron_dataset.py new file mode 100644 index 0000000..e09158f --- /dev/null +++ b/dlas/data/audio/nv_tacotron_dataset.py @@ -0,0 +1,260 @@ +import os +import random + +import torch +import torch.nn.functional as F +import torch.utils.data +import torchaudio +from tqdm import tqdm + +from dlas.data.audio.unsupervised_audio_dataset import load_audio +from dlas.data.util import find_files_of_type, is_audio_file +from dlas.models.audio.tts.tacotron2 import (load_filepaths_and_text, + text_to_sequence) +from dlas.utils.util import opt_get + + +def load_tsv(filename): + with open(filename, encoding='utf-8') as f: + components = [line.strip().split('\t') for line in f] + base = os.path.dirname(filename) + filepaths_and_text = [ + [os.path.join(base, f'{component[1]}'), component[0]] for component in components] + return filepaths_and_text + + +def load_mozilla_cv(filename): + with open(filename, encoding='utf-8') as f: + components = [line.strip().split('\t') + for line in f][1:] # First line is the header + base = os.path.dirname(filename) + filepaths_and_text = [[os.path.join( + base, f'clips/{component[1]}'), component[2]] for component in components] + return filepaths_and_text + + +def load_voxpopuli(filename): + with open(filename, encoding='utf-8') as f: + lines = [line.strip().split('\t') + for line in f][1:] # First line is the header + base = os.path.dirname(filename) + filepaths_and_text = [] + for line in lines: + if len(line) == 0: + continue + file, raw_text, norm_text, speaker_id, split, gender = line + year = file[:4] + filepaths_and_text.append( + [os.path.join(base, year, f'{file}.ogg.wav'), raw_text]) + return filepaths_and_text + + +class TextWavLoader(torch.utils.data.Dataset): + def __init__(self, hparams): + self.path = hparams['path'] + if not isinstance(self.path, list): + self.path = [self.path] + + fetcher_mode = opt_get(hparams, ['fetcher_mode'], 'lj') + if not isinstance(fetcher_mode, list): + fetcher_mode = [fetcher_mode] + assert len(self.path) == len(fetcher_mode) + + self.load_conditioning = opt_get(hparams, ['load_conditioning'], False) + self.conditioning_candidates = opt_get( + hparams, ['num_conditioning_candidates'], 3) + self.conditioning_length = opt_get( + hparams, ['conditioning_length'], 44100) + self.audiopaths_and_text = [] + for p, fm in zip(self.path, fetcher_mode): + if fm == 'lj' or fm == 'libritts': + fetcher_fn = load_filepaths_and_text + elif fm == 'tsv': + fetcher_fn = load_tsv + elif fm == 'mozilla_cv': + # Conditioning inputs are incompatible with mozilla_cv + assert not self.load_conditioning + fetcher_fn = load_mozilla_cv + elif fm == 'voxpopuli': + # Conditioning inputs are incompatible with voxpopuli + assert not self.load_conditioning + fetcher_fn = load_voxpopuli + else: + raise NotImplementedError() + self.audiopaths_and_text.extend(fetcher_fn(p)) + self.text_cleaners = hparams.text_cleaners + self.sample_rate = hparams.sample_rate + random.seed(hparams.seed) + random.shuffle(self.audiopaths_and_text) + self.max_wav_len = opt_get(hparams, ['max_wav_length'], None) + self.max_text_len = opt_get(hparams, ['max_text_length'], None) + # If needs_collate=False, all outputs will be aligned and padded at maximum length. + self.needs_collate = opt_get(hparams, ['needs_collate'], True) + if not self.needs_collate: + assert self.max_wav_len is not None and self.max_text_len is not None + + def get_wav_text_pair(self, audiopath_and_text): + # separate filename and text + audiopath, text = audiopath_and_text[0], audiopath_and_text[1] + text_seq = self.get_text(text) + wav = load_audio(audiopath, self.sample_rate) + return (text_seq, wav, text, audiopath_and_text[0]) + + def get_text(self, text): + text_norm = torch.IntTensor(text_to_sequence(text, self.text_cleaners)) + return text_norm + + def load_conditioning_candidates(self, path): + candidates = find_files_of_type( + 'img', os.path.dirname(path), qualifier=is_audio_file)[0] + # Sanity check to ensure we aren't loading "related files" that aren't actually related. + assert len(candidates) < 50000 + if len(candidates) == 0: + print( + f"No conditioning candidates found for {path} (not even the clip itself??)") + raise NotImplementedError() + # Sample with replacement. This can get repeats, but more conveniently handles situations where there are not enough candidates. + related_clips = [] + for k in range(self.conditioning_candidates): + rel_clip = load_audio(random.choice(candidates), self.sample_rate) + gap = rel_clip.shape[-1] - self.conditioning_length + if gap < 0: + rel_clip = F.pad(rel_clip, pad=(0, abs(gap))) + elif gap > 0: + rand_start = random.randint(0, gap) + rel_clip = rel_clip[:, rand_start:rand_start + + self.conditioning_length] + related_clips.append(rel_clip) + return torch.stack(related_clips, dim=0) + + def __getitem__(self, index): + try: + tseq, wav, text, path = self.get_wav_text_pair( + self.audiopaths_and_text[index]) + cond = self.load_conditioning_candidates( + self.audiopaths_and_text[index][0]) if self.load_conditioning else None + except: + print(f"error loading {self.audiopaths_and_text[index][0]}") + return self[index+1] + if wav is None or \ + (self.max_wav_len is not None and wav.shape[-1] > self.max_wav_len) or \ + (self.max_text_len is not None and tseq.shape[0] > self.max_text_len): + # Basically, this audio file is nonexistent or too long to be supported by the dataset. + # It's hard to handle this situation properly. Best bet is to return the a random valid token and skew the dataset somewhat as a result. + # if wav is not None: + # print(f"Exception {index} wav_len:{wav.shape[-1]} text_len:{tseq.shape[0]} fname: {path}") + rv = random.randint(0, len(self)-1) + return self[rv] + orig_output = wav.shape[-1] + orig_text_len = tseq.shape[0] + if not self.needs_collate: + if wav.shape[-1] != self.max_wav_len: + wav = F.pad(wav, (0, self.max_wav_len - wav.shape[-1])) + if tseq.shape[0] != self.max_text_len: + tseq = F.pad(tseq, (0, self.max_text_len - tseq.shape[0])) + res = { + 'real_text': text, + 'padded_text': tseq, + 'text_lengths': torch.tensor(orig_text_len, dtype=torch.long), + 'wav': wav, + 'wav_lengths': torch.tensor(orig_output, dtype=torch.long), + 'filenames': path + } + if self.load_conditioning: + res['conditioning'] = cond + return res + return tseq, wav, path, text, cond + + def __len__(self): + return len(self.audiopaths_and_text) + + +class TextMelCollate(): + """ Zero-pads model inputs and targets based on number of frames per step + """ + + def __call__(self, batch): + """Collate's training batch from normalized text and wav + PARAMS + ------ + batch: [text_normalized, wav, filename, text] + """ + # Right zero-pad all one-hot text sequences to max input length + input_lengths, ids_sorted_decreasing = torch.sort( + torch.LongTensor([len(x[0]) for x in batch]), + dim=0, descending=True) + max_input_len = input_lengths[0] + + text_padded = torch.LongTensor(len(batch), max_input_len) + text_padded.zero_() + filenames = [] + real_text = [] + conds = [] + for i in range(len(ids_sorted_decreasing)): + text = batch[ids_sorted_decreasing[i]][0] + text_padded[i, :text.size(0)] = text + filenames.append(batch[ids_sorted_decreasing[i]][2]) + real_text.append(batch[ids_sorted_decreasing[i]][3]) + c = batch[ids_sorted_decreasing[i]][4] + if c is not None: + conds.append(c) + + # Right zero-pad wav + num_wavs = batch[0][1].size(0) + max_target_len = max([x[1].size(1) for x in batch]) + + # include mel padded and gate padded + wav_padded = torch.FloatTensor(len(batch), num_wavs, max_target_len) + wav_padded.zero_() + output_lengths = torch.LongTensor(len(batch)) + for i in range(len(ids_sorted_decreasing)): + wav = batch[ids_sorted_decreasing[i]][1] + wav_padded[i, :, :wav.size(1)] = wav + output_lengths[i] = wav.size(1) + + res = { + 'padded_text': text_padded, + 'text_lengths': input_lengths, + 'wav': wav_padded, + 'wav_lengths': output_lengths, + 'filenames': filenames, + 'real_text': real_text, + } + if len(conds) > 0: + res['conditioning'] = torch.stack(conds) + return res + + +if __name__ == '__main__': + batch_sz = 8 + params = { + 'mode': 'nv_tacotron', + 'path': ['Z:\\bigasr_dataset\\libritts\\test-clean_list.txt'], + 'fetcher_mode': ['libritts'], + 'phase': 'train', + 'n_workers': 0, + 'batch_size': batch_sz, + 'needs_collate': False, + 'max_wav_length': 255995, + 'max_text_length': 200, + 'sample_rate': 22050, + 'load_conditioning': True, + 'num_conditioning_candidates': 3, + 'conditioning_length': 44100, + } + from data import create_dataloader, create_dataset + + ds, c = create_dataset(params, return_collate=True) + dl = create_dataloader(ds, params, collate_fn=c) + i = 0 + m = None + for i, b in tqdm(enumerate(dl)): + if i > 5: + break + w = b['wav'] + for ib in range(batch_sz): + print(f'{i} {ib} {b["real_text"][ib]}') + torchaudio.save(f'{i}_clip_{ib}.wav', b['wav'][ib], ds.sample_rate) + for c in range(3): + torchaudio.save(f'{i}_clip_{ib}_cond{c}.wav', + b['conditioning'][ib, c], ds.sample_rate) diff --git a/dlas/data/audio/paired_voice_audio_dataset.py b/dlas/data/audio/paired_voice_audio_dataset.py new file mode 100644 index 0000000..e3c8719 --- /dev/null +++ b/dlas/data/audio/paired_voice_audio_dataset.py @@ -0,0 +1,350 @@ +import os +import random +import sys + +import torch +import torch.nn.functional as F +import torch.utils.data +import torchaudio +from tqdm import tqdm + +from dlas.data.audio.unsupervised_audio_dataset import (load_audio, + load_similar_clips) +from dlas.models.audio.tts.tacotron2 import (load_filepaths_and_text, + load_filepaths_and_text_type, + sequence_to_text, + text_to_sequence) +from dlas.utils.util import opt_get + + +def load_tsv_type(filename, type): + with open(filename, encoding='utf-8') as f: + filepaths_and_text = [] + base = os.path.dirname(filename) + bad_lines = 0 + for line in f: + components = line.strip().split('\t') + if len(components) < 2: + bad_lines += 1 + if bad_lines > 1000: + print( + f'{filename} contains 1000+ bad entries. Failing. Sample last entry: {line}') + raise ValueError + continue + filepaths_and_text.append( + [os.path.join(base, f'{components[1]}'), components[0]] + [type]) + return filepaths_and_text + + +def load_tsv(filename): + with open(filename, encoding='utf-8') as f: + filepaths_and_text = [] + base = os.path.dirname(filename) + bad_lines = 0 + for line in f: + components = line.strip().split('\t') + if len(components) < 2: + bad_lines += 1 + if bad_lines > 1000: + print( + f'{filename} contains 1000+ bad entries. Failing. Sample last entry: {line}') + raise ValueError + continue + filepaths_and_text.append( + [os.path.join(base, f'{components[1]}'), components[0]]) + return filepaths_and_text + + +def convert_string_list_to_tensor(strlist): + if strlist.startswith('['): + strlist = strlist[1:] + if strlist.endswith(']'): + strlist = strlist[:-1] + as_ints = [int(s) for s in strlist.split(', ')] + return torch.tensor(as_ints) + + +def load_tsv_aligned_codes_type(filename, type): + with open(filename, encoding='utf-8') as f: + filepaths_and_text = [] + base = os.path.dirname(filename) + bad_lines = 0 + for line in f: + components = line.strip().split('\t') + if len(components) < 3: + bad_lines += 1 + if bad_lines > 1000: + print( + f'{filename} contains 1000+ bad entries. Failing. Sample last entry: {line}') + raise ValueError + continue + filepaths_and_text.append([os.path.join( + base, f'{components[1]}'), components[0], convert_string_list_to_tensor(components[2])] + [type]) + return filepaths_and_text + + +def load_tsv_aligned_codes(filename): + with open(filename, encoding='utf-8') as f: + filepaths_and_text = [] + base = os.path.dirname(filename) + bad_lines = 0 + for line in f: + components = line.strip().split('\t') + if len(components) < 3: + bad_lines += 1 + if bad_lines > 1000: + print( + f'{filename} contains 1000+ bad entries. Failing. Sample last entry: {line}') + raise ValueError + continue + filepaths_and_text.append([os.path.join( + base, f'{components[1]}'), components[0], convert_string_list_to_tensor(components[2])]) + return filepaths_and_text + + +def load_mozilla_cv(filename, type): + with open(filename, encoding='utf-8') as f: + components = [line.strip().split('\t') + for line in f][1:] # First line is the header + base = os.path.dirname(filename) + filepaths_and_text = [[os.path.join( + base, f'clips/{component[1]}'), component[2], type] for component in components] + return filepaths_and_text + + +def load_voxpopuli(filename, type): + with open(filename, encoding='utf-8') as f: + lines = [line.strip().split('\t') + for line in f][1:] # First line is the header + base = os.path.dirname(filename) + filepaths_and_text = [] + for line in lines: + if len(line) == 0: + continue + file, raw_text, norm_text, speaker_id, split, gender = line + year = file[:4] + filepaths_and_text.append( + [os.path.join(base, year, f'{file}.ogg.wav'), raw_text, type]) + return filepaths_and_text + + +class CharacterTokenizer: + def encode(self, txt): + return text_to_sequence(txt, ['english_cleaners']) + + def decode(self, seq): + return sequence_to_text(seq) + + +class TextWavLoader(torch.utils.data.Dataset): + def __init__(self, hparams): + self.path = hparams['path'] + if not isinstance(self.path, list): + self.path = [self.path] + self.types = opt_get(hparams, ['types'], [0 for _ in self.path]) + + fetcher_mode = opt_get(hparams, ['fetcher_mode'], 'lj') + if not isinstance(fetcher_mode, list): + fetcher_mode = [fetcher_mode] + assert len(self.path) == len(fetcher_mode) + + self.load_conditioning = opt_get(hparams, ['load_conditioning'], False) + self.conditioning_candidates = opt_get( + hparams, ['num_conditioning_candidates'], 1) + self.conditioning_length = opt_get( + hparams, ['conditioning_length'], 44100) + self.debug_failures = opt_get( + hparams, ['debug_loading_failures'], False) + self.load_aligned_codes = opt_get( + hparams, ['load_aligned_codes'], False) + self.aligned_codes_to_audio_ratio = opt_get( + hparams, ['aligned_codes_ratio'], 443) + self.audiopaths_and_text = [] + for p, fm, type in zip(self.path, fetcher_mode, self.types): + if fm == 'lj' or fm == 'libritts': + fetcher_fn = load_filepaths_and_text_type + elif fm == 'tsv': + fetcher_fn = load_tsv_aligned_codes_type if self.load_aligned_codes else load_tsv_type + elif fm == 'mozilla_cv': + # Conditioning inputs are incompatible with mozilla_cv + assert not self.load_conditioning + fetcher_fn = load_mozilla_cv + elif fm == 'voxpopuli': + # Conditioning inputs are incompatible with voxpopuli + assert not self.load_conditioning + fetcher_fn = load_voxpopuli + else: + raise NotImplementedError() + self.audiopaths_and_text.extend(fetcher_fn(p, type)) + self.text_cleaners = hparams.text_cleaners + self.sample_rate = hparams.sample_rate + random.seed(hparams.seed) + random.shuffle(self.audiopaths_and_text) + self.max_wav_len = opt_get(hparams, ['max_wav_length'], None) + if self.max_wav_len is not None: + self.max_aligned_codes = self.max_wav_len // self.aligned_codes_to_audio_ratio + self.max_text_len = opt_get(hparams, ['max_text_length'], None) + assert self.max_wav_len is not None and self.max_text_len is not None + self.use_bpe_tokenizer = opt_get(hparams, ['use_bpe_tokenizer'], True) + if self.use_bpe_tokenizer: + from dlas.data.audio.voice_tokenizer import VoiceBpeTokenizer + self.tokenizer = VoiceBpeTokenizer(opt_get( + hparams, ['tokenizer_vocab'], '../experiments/bpe_lowercase_asr_256.json')) + else: + self.tokenizer = CharacterTokenizer() + # records how many items are skipped when accessing an index. + self.skipped_items = 0 + + def get_wav_text_pair(self, audiopath_and_text): + # separate filename and text + audiopath, text, type = audiopath_and_text[0], audiopath_and_text[1], audiopath_and_text[2] + text_seq = self.get_text(text) + wav = load_audio(audiopath, self.sample_rate) + return (text_seq, wav, text, audiopath_and_text[0], type) + + def get_text(self, text): + tokens = self.tokenizer.encode(text) + tokens = torch.IntTensor(tokens) + if self.use_bpe_tokenizer: + # Assert if any UNK,start tokens encountered. + assert not torch.any(tokens == 1) + # The stop token should always be sacred. + assert not torch.any(tokens == 0) + return tokens + + def __getitem__(self, index): + self.skipped_items += 1 + try: + tseq, wav, text, path, type = self.get_wav_text_pair( + self.audiopaths_and_text[index]) + if text is None or len(text.strip()) == 0: + raise ValueError + if wav is None or wav.shape[-1] < (.6 * self.sample_rate): + # Ultra short clips are also useless (and can cause problems within some models). + raise ValueError + cond, cond_is_self = load_similar_clips(self.audiopaths_and_text[index][0], self.conditioning_length, self.sample_rate, + n=self.conditioning_candidates) if self.load_conditioning else (None, False) + except: + if self.skipped_items > 100: + raise # Rethrow if we have nested too far. + if self.debug_failures: + print( + f"error loading {self.audiopaths_and_text[index][0]} {sys.exc_info()}") + return self[(index+1) % len(self)] + + if self.load_aligned_codes: + aligned_codes = self.audiopaths_and_text[index][3] + + actually_skipped_items = self.skipped_items + self.skipped_items = 0 + if wav is None or \ + (self.max_wav_len is not None and wav.shape[-1] > self.max_wav_len) or \ + (self.max_text_len is not None and tseq.shape[0] > self.max_text_len): + # Basically, this audio file is nonexistent or too long to be supported by the dataset. + # It's hard to handle this situation properly. Best bet is to return the a random valid token and skew the dataset somewhat as a result. + if self.debug_failures: + print( + f"error loading {path}: ranges are out of bounds; {wav.shape[-1]}, {tseq.shape[0]}") + rv = random.randint(0, len(self)-1) + return self[rv] + orig_output = wav.shape[-1] + orig_text_len = tseq.shape[0] + if wav.shape[-1] != self.max_wav_len: + wav = F.pad(wav, (0, self.max_wav_len - wav.shape[-1])) + if self.load_aligned_codes: + # These codes are aligned to audio inputs, so make sure to pad them as well. + aligned_codes = F.pad( + aligned_codes, (0, self.max_aligned_codes-aligned_codes.shape[0])) + if tseq.shape[0] != self.max_text_len: + tseq = F.pad(tseq, (0, self.max_text_len - tseq.shape[0])) + res = { + 'real_text': text, + 'padded_text': tseq, + 'text_lengths': torch.tensor(orig_text_len, dtype=torch.long), + 'wav': wav, + 'wav_lengths': torch.tensor(orig_output, dtype=torch.long), + 'filenames': path, + 'skipped_items': actually_skipped_items, + 'type': type, + } + if self.load_conditioning: + res['conditioning'] = cond + res['conditioning_contains_self'] = cond_is_self + if self.load_aligned_codes: + res['aligned_codes'] = aligned_codes + return res + + def __len__(self): + return len(self.audiopaths_and_text) + + +class PairedVoiceDebugger: + def __init__(self): + self.total_items = 0 + self.loaded_items = 0 + self.self_conditioning_items = 0 + + def get_state(self): + return {'total_items': self.total_items, + 'loaded_items': self.loaded_items, + 'self_conditioning_items': self.self_conditioning_items} + + def load_state(self, state): + if isinstance(state, dict): + self.total_items = opt_get(state, ['total_items'], 0) + self.loaded_items = opt_get(state, ['loaded_items'], 0) + self.self_conditioning_items = opt_get( + state, ['self_conditioning_items'], 0) + + def update(self, batch): + self.total_items += batch['wav'].shape[0] + self.loaded_items += batch['skipped_items'].sum().item() + if 'conditioning' in batch.keys(): + self.self_conditioning_items += batch['conditioning_contains_self'].sum( + ).item() + + def get_debugging_map(self): + return { + 'total_samples_loaded': self.total_items, + 'percent_skipped_samples': (self.loaded_items - self.total_items) / self.loaded_items, + 'percent_conditioning_is_self': self.self_conditioning_items / self.loaded_items, + } + + +if __name__ == '__main__': + batch_sz = 8 + params = { + 'mode': 'paired_voice_audio', + 'path': ['Y:\\libritts/test-clean_list.txt'], + 'fetcher_mode': ['libritts'], + 'phase': 'train', + 'n_workers': 0, + 'batch_size': batch_sz, + 'max_wav_length': 255995, + 'max_text_length': 200, + 'sample_rate': 22050, + 'load_conditioning': True, + 'num_conditioning_candidates': 2, + 'conditioning_length': 44000, + 'use_bpe_tokenizer': True, + 'load_aligned_codes': False, + } + from data import create_dataloader, create_dataset + + def save(b, i, ib, key, c=None): + if c is not None: + torchaudio.save(f'{i}_clip_{ib}_{key}_{c}.wav', + b[key][ib][c], 22050) + else: + torchaudio.save(f'{i}_clip_{ib}_{key}.wav', b[key][ib], 22050) + + ds, c = create_dataset(params, return_collate=True) + dl = create_dataloader(ds, params, collate_fn=c) + i = 0 + m = None + for i, b in tqdm(enumerate(dl)): + for ib in range(batch_sz): + print(f'{i} {ib} {b["real_text"][ib]}') + save(b, i, ib, 'wav') + if i > 5: + break diff --git a/dlas/data/audio/preprocessed_mel_dataset.py b/dlas/data/audio/preprocessed_mel_dataset.py new file mode 100644 index 0000000..9a29bec --- /dev/null +++ b/dlas/data/audio/preprocessed_mel_dataset.py @@ -0,0 +1,76 @@ +import os +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F +import torch.utils.data +import torchaudio +import torchvision +from tqdm import tqdm + +from dlas.utils.util import opt_get + + +class PreprocessedMelDataset(torch.utils.data.Dataset): + + def __init__(self, opt): + path = opt['path'] + # Will fail when multiple paths specified, must be specified in this case. + cache_path = opt['cache_path'] + if os.path.exists(cache_path): + self.paths = torch.load(cache_path) + else: + print("Building cache..") + path = Path(path) + self.paths = [str(p) for p in path.rglob("*.npz")] + torch.save(self.paths, cache_path) + self.pad_to = opt_get(opt, ['pad_to_samples'], 10336) + self.squeeze = opt_get(opt, ['should_squeeze'], False) + + def __getitem__(self, index): + with np.load(self.paths[index]) as npz_file: + mel = torch.tensor(npz_file['arr_0']) + assert mel.shape[-1] <= self.pad_to + if self.squeeze: + mel = mel.squeeze() + padding_needed = self.pad_to - mel.shape[-1] + mask = torch.zeros_like(mel) + if padding_needed > 0: + mel = F.pad(mel, (0, padding_needed)) + mask = F.pad(mask, (0, padding_needed), value=1) + + output = { + 'mel': mel, + 'mel_lengths': torch.tensor(mel.shape[-1]), + 'mask': mask, + 'mask_lengths': torch.tensor(mask.shape[-1]), + 'path': self.paths[index], + } + return output + + def __len__(self): + return len(self.paths) + + +if __name__ == '__main__': + params = { + 'mode': 'preprocessed_mel', + 'path': 'Y:\\separated\\large_mel_cheaters', + 'cache_path': 'Y:\\separated\\large_mel_cheaters_win.pth', + 'pad_to_samples': 646, + 'phase': 'train', + 'n_workers': 0, + 'batch_size': 16, + } + from data import create_dataloader, create_dataset + + ds = create_dataset(params) + dl = create_dataloader(ds, params) + i = 0 + for b in tqdm(dl): + # pass + torchvision.utils.save_image((b['mel'].unsqueeze(1)+1)/2, f'{i}.png') + i += 1 + if i > 20: + break diff --git a/dlas/data/audio/unsupervised_audio_dataset.py b/dlas/data/audio/unsupervised_audio_dataset.py new file mode 100644 index 0000000..79451b6 --- /dev/null +++ b/dlas/data/audio/unsupervised_audio_dataset.py @@ -0,0 +1,228 @@ +import os +import random +import sys + +import torch +import torch.nn.functional as F +import torch.utils.data +import torchaudio +from audio2numpy import open_audio +from tqdm import tqdm + +from dlas.data.util import (find_files_of_type, is_audio_file, + load_paths_from_cache) +from dlas.models.audio.tts.tacotron2.taco_utils import load_wav_to_torch +from dlas.utils.util import opt_get + + +def load_audio(audiopath, sampling_rate): + if audiopath[-4:] == '.wav': + audio, lsr = load_wav_to_torch(audiopath) + elif audiopath[-4:] == '.mp3': + # https://github.com/neonbjb/pyfastmp3decoder - Definitely worth it. + from pyfastmp3decoder.mp3decoder import load_mp3 + audio, lsr = load_mp3(audiopath, sampling_rate) + audio = torch.FloatTensor(audio) + else: + audio, lsr = open_audio(audiopath) + audio = torch.FloatTensor(audio) + + # Remove any channel data. + if len(audio.shape) > 1: + if audio.shape[0] < 5: + audio = audio[0] + else: + assert audio.shape[1] < 5 + audio = audio[:, 0] + + if lsr != sampling_rate: + audio = torchaudio.functional.resample(audio, lsr, sampling_rate) + + # Check some assumptions about audio range. This should be automatically fixed in load_wav_to_torch, but might not be in some edge cases, where we should squawk. + # '10' is arbitrarily chosen since it seems like audio will often "overdrive" the [-1,1] bounds. + if torch.any(audio > 10) or not torch.any(audio < 0): + print(f"Error with {audiopath}. Max={audio.max()} min={audio.min()}") + audio.clip_(-1, 1) + + return audio.unsqueeze(0) + + +def load_similar_clips(path, sample_length, sample_rate, n=3, fallback_to_self=True): + sim_path = os.path.join(os.path.dirname(path), 'similarities.pth') + candidates = [] + if os.path.exists(sim_path): + similarities = torch.load(sim_path) + fname = os.path.basename(path) + if fname in similarities.keys(): + candidates = [os.path.join(os.path.dirname(path), s) + for s in similarities[fname]] + else: + print( + f'Similarities list found for {path} but {fname} was not in that list.') + # candidates.append(path) # Always include self as a possible similar clip. + if len(candidates) == 0: + if fallback_to_self: + candidates = [path] + else: + candidates = find_files_of_type( + 'img', os.path.dirname(path), qualifier=is_audio_file)[0] + + # Sanity check to ensure we aren't loading "related files" that aren't actually related. + assert len(candidates) < 50000 + if len(candidates) == 0: + print(f"No conditioning candidates found for {path}") + raise NotImplementedError() + + # Sample with replacement. This can get repeats, but more conveniently handles situations where there are not enough candidates. + related_clips = [] + contains_self = False + for k in range(n): + rel_path = random.choice(candidates) + contains_self = contains_self or (rel_path == path) + rel_clip = load_audio(rel_path, sample_rate) + gap = rel_clip.shape[-1] - sample_length + if gap < 0: + rel_clip = F.pad(rel_clip, pad=(0, abs(gap))) + elif gap > 0: + rand_start = random.randint(0, gap) + rel_clip = rel_clip[:, rand_start:rand_start+sample_length] + related_clips.append(rel_clip) + if n > 1: + return torch.stack(related_clips, dim=0), contains_self + else: + return related_clips[0], contains_self + + +class UnsupervisedAudioDataset(torch.utils.data.Dataset): + + def __init__(self, opt): + path = opt['path'] + # Will fail when multiple paths specified, must be specified in this case. + cache_path = opt['cache_path'] + exclusions = [] + if 'exclusions' in opt.keys(): + for exc in opt['exclusions']: + with open(exc, 'r') as f: + exclusions.extend(f.read().splitlines()) + ew = opt_get(opt, ['endswith'], []) + assert isinstance(ew, list) + not_ew = opt_get(opt, ['not_endswith'], []) + assert isinstance(not_ew, list) + self.audiopaths = load_paths_from_cache( + path, cache_path, exclusions, endswith=ew, not_endswith=not_ew) + + # Parse options + self.sampling_rate = opt_get(opt, ['sampling_rate'], 22050) + self.pad_to = opt_get(opt, ['pad_to_seconds'], None) + if self.pad_to is not None: + self.pad_to *= self.sampling_rate + self.pad_to = opt_get(opt, ['pad_to_samples'], self.pad_to) + self.min_length = opt_get(opt, ['min_length'], 0) + self.dont_clip = opt_get(opt, ['dont_clip'], False) + + # "Resampled clip" is audio data pulled from the basis of "clip" but with randomly different bounds. There are no + # guarantees that "clip_resampled" is different from "clip": in fact, if "clip" is less than pad_to_seconds/samples, + self.should_resample_clip = opt_get(opt, ['resample_clip'], False) + + # "Extra samples" are other audio clips pulled from wav files in the same directory as the 'clip' wav file. + self.extra_samples = opt_get(opt, ['extra_samples'], 0) + self.extra_sample_len = opt_get(opt, ['extra_sample_length'], 44000) + + self.debug_loading_failures = opt_get( + opt, ['debug_loading_failures'], True) + + def get_audio_for_index(self, index): + audiopath = self.audiopaths[index] + audio = load_audio(audiopath, self.sampling_rate) + assert audio.shape[1] > self.min_length + if self.dont_clip: + assert audio.shape[1] <= self.pad_to + return audio, audiopath + + def get_related_audio_for_index(self, index): + if self.extra_samples <= 0: + return None, 0 + audiopath = self.audiopaths[index] + return load_similar_clips(audiopath, self.extra_sample_len, self.sampling_rate, n=self.extra_samples) + + def __getitem__(self, index): + try: + # Split audio_norm into two tensors of equal size. + audio_norm, filename = self.get_audio_for_index(index) + alt_files, alt_is_self = self.get_related_audio_for_index(index) + except: + if self.debug_loading_failures: + print( + f"Error loading audio for file {self.audiopaths[index]} {sys.exc_info()}") + return self[random.randint(0, len(self))] + + # When generating resampled clips, skew is a bias that tries to spread them out from each other, reducing their + # influence on one another. + skew = [-1, 1] if self.should_resample_clip else [0] + # To increase variability, which skew is applied to the clip and resampled_clip is randomized. + random.shuffle(skew) + clips = [] + prepad_length = min(audio_norm.shape[-1], self.pad_to) + for sk in skew: + if self.pad_to is not None: + if audio_norm.shape[-1] <= self.pad_to: + clips.append(torch.nn.functional.pad( + audio_norm, (0, self.pad_to - audio_norm.shape[-1]))) + else: + gap = audio_norm.shape[-1] - self.pad_to + start = min( + max(random.randint(0, gap-1) + sk * gap // 2, 0), gap-1) + clips.append(audio_norm[:, start:start+self.pad_to]) + else: + clips.append(audio_norm) + + output = { + 'prepad_length': prepad_length, + 'clip': clips[0], + 'clip_lengths': torch.tensor(clips[0].shape[-1]), + 'path': filename, + } + if self.should_resample_clip: + output['resampled_clip'] = clips[1] + if self.extra_samples > 0: + output['alt_clips'] = alt_files + output['alt_contains_self'] = alt_is_self + return output + + def __len__(self): + return len(self.audiopaths) + + +if __name__ == '__main__': + params = { + 'mode': 'unsupervised_audio', + 'path': ['Y:\\separated\\yt-music-0', 'Y:\\separated\\yt-music-1', + 'Y:\\separated\\bt-music-1', 'Y:\\separated\\bt-music-2', + 'Y:\\separated\\bt-music-3', 'Y:\\separated\\bt-music-4', + 'Y:\\separated\\bt-music-5'], + 'cache_path': 'Y:\\separated\\no-vocals-cache-win.pth', + 'endswith': ['no_vocals.wav'], + 'sampling_rate': 22050, + 'pad_to_samples': 200000, + 'resample_clip': False, + 'extra_samples': 1, + 'extra_sample_length': 100000, + 'phase': 'train', + 'n_workers': 1, + 'batch_size': 16, + } + from data import create_dataloader, create_dataset + + ds = create_dataset(params) + dl = create_dataloader(ds, params) + i = 0 + for b in tqdm(dl): + for b_ in range(b['clip'].shape[0]): + # pass + torchaudio.save(f'{i}_clip_{b_}.wav', + b['clip'][b_], ds.sampling_rate) + torchaudio.save(f'{i}_alt_clip_{b_}.wav', + b['alt_clips'][b_], ds.sampling_rate) + i += 1 + if i > 200: + break diff --git a/dlas/data/audio/voice_tokenizer.py b/dlas/data/audio/voice_tokenizer.py new file mode 100644 index 0000000..790a4e0 --- /dev/null +++ b/dlas/data/audio/voice_tokenizer.py @@ -0,0 +1,187 @@ +import json +import re + +import torch +from tokenizers import Tokenizer +from tokenizers.models import BPE +from tokenizers.pre_tokenizers import Whitespace +from tokenizers.trainers import BpeTrainer + +from dlas.data.audio.paired_voice_audio_dataset import (load_mozilla_cv, + load_tsv, + load_voxpopuli) +from dlas.models.audio.tts.tacotron2 import load_filepaths_and_text +from dlas.models.audio.tts.tacotron2.text.cleaners import english_cleaners + + +def remove_extraneous_punctuation(word): + replacement_punctuation = { + '{': '(', '}': ')', + '[': '(', ']': ')', + '`': '\'', '—': '-', + '—': '-', '`': '\'', + 'ʼ': '\'' + } + replace = re.compile("|".join([re.escape(k) for k in sorted( + replacement_punctuation, key=len, reverse=True)]), flags=re.DOTALL) + word = replace.sub(lambda x: replacement_punctuation[x.group(0)], word) + + # TODO: some of these are spoken ('@', '%', '+', etc). Integrate them into the cleaners. + extraneous = re.compile(r'^[@#%_=\$\^&\*\+\\]$') + word = extraneous.sub('', word) + return word + + +def expand_numbers(text): + return normalize_numbers(text) + + +def lowercase(text): + return text.lower() + + +_whitespace_re = re.compile(r'\s+') + + +def collapse_whitespace(text): + return re.sub(_whitespace_re, ' ', text) + + +def convert_to_ascii(text): + return unidecode(text) + + +def basic_cleaners(text): + '''Basic pipeline that lowercases and collapses whitespace without transliteration.''' + text = lowercase(text) + text = collapse_whitespace(text) + return text + + +class VoiceBpeTokenizer: + def __init__(self, vocab_file, preprocess=None): + with open(vocab_file, 'r', encoding='utf-8') as f: + vocab = json.load(f) + + self.language = vocab['model']['language'] if 'language' in vocab['model'] else None + + if preprocess is None: + self.preprocess = 'pre_tokenizer' in vocab and vocab['pre_tokenizer'] + else: + self.preprocess = preprocess + + if vocab_file is not None: + self.tokenizer = Tokenizer.from_file(vocab_file) + + def preprocess_text(self, txt): + if self.language == 'ja': + import pykakasi + + kks = pykakasi.kakasi() + results = kks.convert(txt) + txt = " ".join([result['kana'] for result in results]) + txt = basic_cleaners(txt) + else: + txt = english_cleaners(txt) + + return txt + + def encode(self, txt): + if self.preprocess: + txt = self.preprocess_text(txt) + txt = txt.replace(' ', '[SPACE]') + return self.tokenizer.encode(txt).ids + + def decode(self, seq): + if isinstance(seq, torch.Tensor): + seq = seq.cpu().numpy() + txt = self.tokenizer.decode( + seq, skip_special_tokens=False).replace(' ', '') + txt = txt.replace('[SPACE]', ' ') + txt = txt.replace('[STOP]', '') + txt = txt.replace('[UNK]', '') + return txt + + +def build_text_file_from_priors(priors, output): + with open(output, 'w', encoding='utf-8') as out: + for p, fm in priors: + if fm == 'lj' or fm == 'libritts': + fetcher_fn = load_filepaths_and_text + elif fm == 'tsv': + fetcher_fn = load_tsv + elif fm == 'mozilla_cv': + fetcher_fn = load_mozilla_cv + elif fm == 'voxpopuli': + fetcher_fn = load_voxpopuli + else: + raise NotImplementedError() + apt = fetcher_fn(p) + for path, text in apt: + out.write(text + "\n") + out.flush() + + +def train(): + with open('all_texts.txt', 'r', encoding='utf-8') as at: + ttsd = at.readlines() + # bcd = datasets.load_dataset('bookcorpus', cache_dir='Z:\\huggingface_datasets\\cache')['train'] + + # allowed_characters_re = re.compile(r'^[0-9a-z!@#%_=:;"/, \-\$\^&\*\(\)\+\{\[\]\}\\\.\'\?—–ʼ]+$') + allowed_characters_re = re.compile(r'^[a-z!:;"/, \-\(\)\.\'\?ʼ]+$') + + def preprocess_word(word, report=False): + word = english_cleaners(word) + word = remove_extraneous_punctuation(word) + if not bool(allowed_characters_re.match(word)): + if report and word: + print(f"REPORTING: '{word}'") + return '' + return word + + def batch_iterator(batch_size=1000): + print("Processing ASR texts.") + for i in range(0, len(ttsd), batch_size): + yield [preprocess_word(t, True) for t in ttsd[i:i+batch_size]] + + # print("Processing bookcorpus.") + # for i in range(0, len(bcd), batch_size): + # yield [preprocess_word(t) for t in bcd[i:i+batch_size]['text']] + + trainer = BpeTrainer( + special_tokens=['[STOP]', '[UNK]', '[SPACE]'], vocab_size=255) + tokenizer = Tokenizer(BPE(unk_token="[UNK]")) + tokenizer.pre_tokenizer = Whitespace() + tokenizer.train_from_iterator( + batch_iterator(), trainer, length=len(ttsd)) # +len(bcd)) + + print(tokenizer.decode(tokenizer.encode( + "i was traveling throughhadslfghds the woods in 1235375t137{{}}").ids)) + + tokenizer.save('gpt_tts_tokenizer.json') + + +def test(): + tok = VoiceBpeTokenizer('gpt_tts_tokenizer.json') + with open('all_texts.txt', 'r', encoding='utf-8') as at: + ttsd = at.readlines() + for line in ttsd: + line = line.strip() + seq = tok.encode(line) + out = tok.decode(seq) + print(f">>>{line}") + print(f"<<<{out}") + + +if __name__ == '__main__': + ''' + build_text_file_from_priors([('Y:\\bigasr_dataset\\libritts\\train-all.txt', 'libritts'), + ('Y:\\bigasr_dataset\\libritts\\test-clean_list.txt', 'libritts'), + #('Y:\\bigasr_dataset\\voxpopuli\\audio\\transcribed_data\\en\\asr_en.tsv', 'voxpopuli'), + ('Y:\\bigasr_dataset\\voxpopuli\\audio\\transcribed_data\\en\\asr_train.tsv', 'voxpopuli'), + ('Y:\\clips\\books1-transcribed.tsv', 'tsv'), + ('Y:\\clips\\books2-transcribed.tsv', 'tsv'), + ('Y:\\clips\\podcasts-0-transcribed.tsv', 'tsv')], 'all_texts.txt') + ''' + # train() + test() diff --git a/dlas/data/audio/wav_aug.py b/dlas/data/audio/wav_aug.py new file mode 100644 index 0000000..5955b4f --- /dev/null +++ b/dlas/data/audio/wav_aug.py @@ -0,0 +1,64 @@ +import random + +import torch +import torchaudio.sox_effects + +from dlas.models.audio.tts.tacotron2.taco_utils import load_wav_to_torch + + +# Returns random double on [l,h] as a string +def rdstr(l=0, h=1): + assert h > l + i = h-l + return str(random.random() * i + l) + + +# Returns a randint on [s,e] as a string +def rdi(e, s=0): + return str(random.randint(s, e)) + + +class WavAugmentor: + def __init__(self): + pass + + def augment(self, wav, sample_rate): + speed_effect = ['speed', rdstr(.8, 1)] + ''' + Band effects are disabled until I can audit them better. + band_effects = [ + ['reverb', '-w'], + ['reverb'], + ['band', rdi(8000, 3000), rdi(1000, 100)], + ['bandpass', rdi(8000, 3000), rdi(1000, 100)], + ['bass', rdi(20,-20)], + ['treble', rdi(20,-20)], + ['dither'], + ['equalizer', rdi(3000, 100), rdi(1000, 100), rdi(10, -10)], + ['hilbert'], + ['sinc', '3k'], + ['sinc', '-4k'], + ['sinc', '3k-4k'] + ] + band_effect = random.choice(band_effects) + ''' + volume_effects = [ + ['loudness', rdi(10, -2)], + ['overdrive', rdi(20, 0), rdi(20, 0)], + ] + vol_effect = random.choice(volume_effects) + effects = [speed_effect, vol_effect] + out, sr = torchaudio.sox_effects.apply_effects_tensor( + wav, sample_rate, effects) + # Add a variable amount of noise + out = out + torch.rand_like(out) * random.random() * .03 + return out + + +if __name__ == '__main__': + sample, _ = load_wav_to_torch('obama1.wav') + sample = sample / 32768.0 + aug = WavAugmentor() + for j in range(10): + out = aug.augment(sample, 24000) + torchaudio.save(f'out{j}.wav', out, 24000) diff --git a/dlas/data/combined_dataset.py b/dlas/data/combined_dataset.py new file mode 100644 index 0000000..598db78 --- /dev/null +++ b/dlas/data/combined_dataset.py @@ -0,0 +1,35 @@ +import torch + +from dlas.data import create_dataset + + +# Simple composite dataset that combines multiple other datasets. +# Assumes that the datasets output dicts. +class CombinedDataset(torch.utils.data.Dataset): + def __init__(self, opt): + self.datasets = {} + for k, v in opt.items(): + if not isinstance(v, dict): + continue + # Scale&phase gets injected by options.py.. + v['scale'] = opt['scale'] + v['phase'] = opt['phase'] + self.datasets[k] = create_dataset(v) + self.items_fetched = 0 + + def __getitem__(self, i): + self.items_fetched += 1 + output = {} + for name, dataset in self.datasets.items(): + prefix = "" + # 'default' dataset gets no prefix, other ones get `key_` + if name != 'default': + prefix = name + "_" + + data = dataset[i % len(dataset)] + for k, v in data.items(): + output[prefix + k] = v + return output + + def __len__(self): + return max(len(d) for d in self.datasets.values()) diff --git a/dlas/data/data_sampler.py b/dlas/data/data_sampler.py new file mode 100644 index 0000000..43464d4 --- /dev/null +++ b/dlas/data/data_sampler.py @@ -0,0 +1,69 @@ +""" +Modified from torch.utils.data.distributed.DistributedSampler +Support enlarging the dataset for *iteration-oriented* training, for saving time when restart the +dataloader after each epoch +""" +import math + +import torch +import torch.distributed as dist +from torch.utils.data.sampler import Sampler + + +class DistIterSampler(Sampler): + """Sampler that restricts data loading to a subset of the dataset. + + It is especially useful in conjunction with + :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each + process can pass a DistributedSampler instance as a DataLoader sampler, + and load a subset of the original dataset that is exclusive to it. + + .. note:: + Dataset is assumed to be of constant size. + + Arguments: + dataset: Dataset used for sampling. + num_replicas (optional): Number of processes participating in + distributed training. + rank (optional): Rank of the current process within num_replicas. + """ + + def __init__(self, dataset, num_replicas=None, rank=None, ratio=100): + if num_replicas is None: + if not dist.is_available(): + raise RuntimeError( + "Requires distributed package to be available") + num_replicas = dist.get_world_size() + if rank is None: + if not dist.is_available(): + raise RuntimeError( + "Requires distributed package to be available") + rank = dist.get_rank() + self.dataset = dataset + self.num_replicas = num_replicas + self.rank = rank + self.epoch = 0 + self.num_samples = int( + math.ceil(len(self.dataset) * ratio / self.num_replicas)) + self.total_size = self.num_samples * self.num_replicas + + def __iter__(self): + # deterministically shuffle based on epoch + g = torch.Generator() + g.manual_seed(self.epoch) + indices = torch.randperm(self.total_size, generator=g).tolist() + + dsize = len(self.dataset) + indices = [v % dsize for v in indices] + + # subsample + indices = indices[self.rank:self.total_size:self.num_replicas] + assert len(indices) == self.num_samples + + return iter(indices) + + def __len__(self): + return self.num_samples + + def set_epoch(self, epoch): + self.epoch = epoch diff --git a/dlas/data/images/__init__.py b/dlas/data/images/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/dlas/data/images/base_unsupervised_image_dataset.py b/dlas/data/images/base_unsupervised_image_dataset.py new file mode 100644 index 0000000..0eb8525 --- /dev/null +++ b/dlas/data/images/base_unsupervised_image_dataset.py @@ -0,0 +1,143 @@ +import os + +import cv2 +import numpy as np +import torch +from torch.utils import data + +from dlas.data.images.chunk_with_reference import ChunkWithReference +from dlas.data.images.image_corruptor import ImageCorruptor + + +# Class whose purpose is to hold as much logic as can possibly be shared between datasets that operate on raw image +# data and nothing else (which also have a very specific directory structure being used, as dictated by +# ChunkWithReference). +class BaseUnsupervisedImageDataset(data.Dataset): + def __init__(self, opt): + self.opt = opt + self.corruptor = ImageCorruptor(opt) + self.target_hq_size = opt['target_size'] if 'target_size' in opt.keys( + ) else None + self.multiple = opt['force_multiple'] if 'force_multiple' in opt.keys( + ) else 1 + self.for_eval = opt['eval'] if 'eval' in opt.keys() else False + self.scale = opt['scale'] if not self.for_eval else 1 + self.paths = opt['paths'] + self.corrupt_before_downsize = opt['corrupt_before_downsize'] if 'corrupt_before_downsize' in opt.keys( + ) else False + # If we dont throw here, we get some really obscure errors. + assert (self.target_hq_size // self.scale) % self.multiple == 0 + if not isinstance(self.paths, list): + self.paths = [self.paths] + self.weights = [1] + else: + self.weights = opt['weights'] + + # See if there is a cached directory listing and use that rather than re-scanning everything. This will greatly + # reduce startup costs. + self.chunks = [] + for path, weight in zip(self.paths, self.weights): + cache_path = os.path.join(path, 'cache.pth') + if os.path.exists(cache_path): + chunks = torch.load(cache_path) + else: + print( + "Building chunk cache, this can take some time for large datasets..") + chunks = [ChunkWithReference(opt, d) for d in sorted( + os.scandir(path), key=lambda e: e.name) if d.is_dir()] + # Prune out chunks that have no images + res = [] + for c in chunks: + if len(c) != 0: + res.append(c) + chunks = res + # Save to a cache. + torch.save(chunks, cache_path) + for w in range(weight): + self.chunks.extend(chunks) + + # Indexing this dataset is tricky. Aid it by having a list of starting indices for each chunk. + start = 0 + self.starting_indices = [] + for c in self.chunks: + self.starting_indices.append(start) + start += len(c) + self.len = start + + def get_paths(self): + paths = [] + for c in self.chunks: + paths.extend(c.tiles) + return paths + + # Utility method for translating a point when the dimensions of an image change. + def resize_point(self, point, orig_dim, new_dim): + oh, ow = orig_dim + nh, nw = new_dim + dh, dw = float(nh) / float(oh), float(nw) / float(ow) + point = int(dh * float(point[0])), int(dw * float(point[1])) + return point + + # Given an HQ square of arbitrary size, resizes it to specifications from opt. + def resize_hq(self, imgs_hq, refs_hq, masks_hq, centers_hq): + # Enforce size constraints + h, w, _ = imgs_hq[0].shape + if self.target_hq_size is not None and self.target_hq_size != h: + hqs_adjusted, hq_refs_adjusted, hq_masks_adjusted, hq_centers_adjusted = [], [], [], [] + for hq, hq_ref, hq_mask, hq_center in zip(imgs_hq, refs_hq, masks_hq, centers_hq): + # It is assumed that the target size is a square. + target_size = (self.target_hq_size, self.target_hq_size) + hqs_adjusted.append(cv2.resize( + hq, target_size, interpolation=cv2.INTER_AREA)) + hq_refs_adjusted.append(cv2.resize( + hq_ref, target_size, interpolation=cv2.INTER_AREA)) + hq_masks_adjusted.append(cv2.resize( + hq_mask, target_size, interpolation=cv2.INTER_AREA)) + hq_centers_adjusted.append( + self.resize_point(hq_center, (h, w), target_size)) + h, w = self.target_hq_size, self.target_hq_size + else: + hqs_adjusted, hq_refs_adjusted, hq_masks_adjusted, hq_centers_adjusted = imgs_hq, refs_hq, masks_hq, centers_hq + # This is done implicitly above.. + hq_masks_adjusted = [m.squeeze(-1) for m in hq_masks_adjusted] + # Multiple must apply to LQ image. + hq_multiple = self.multiple * self.scale + if h % hq_multiple != 0 or w % hq_multiple != 0: + hqs_conformed, hq_refs_conformed, hq_masks_conformed, hq_centers_conformed = [], [], [], [] + for hq, hq_ref, hq_mask, hq_center in zip(hqs_adjusted, hq_refs_adjusted, hq_masks_adjusted, hq_centers_adjusted): + h, w = (h - h % hq_multiple), (w - w % hq_multiple) + hq_centers_conformed.append( + self.resize_point(hq_center, hq.shape[:2], (h, w))) + hqs_conformed.append(hq[:h, :w, :]) + hq_refs_conformed.append(hq_ref[:h, :w, :]) + hq_masks_conformed.append(hq_mask[:h, :w, :]) + return hqs_conformed, hq_refs_conformed, hq_masks_conformed, hq_centers_conformed + return hqs_adjusted, hq_refs_adjusted, hq_masks_adjusted, hq_centers_adjusted + + def synthesize_lq(self, hs, hrefs, hmasks, hcenters): + h, w, _ = hs[0].shape + ls, lrs, lms, lcs = [], [], [], [] + if self.corrupt_before_downsize and not self.for_eval: + hs = self.corruptor.corrupt_images(np.copy(hs)) + for hq, hq_ref, hq_mask, hq_center in zip(hs, hrefs, hmasks, hcenters): + if self.for_eval: + ls.append(hq) + lrs.append(hq_ref) + lms.append(hq_mask) + lcs.append(hq_center) + else: + ls.append(cv2.resize(hq, (h // self.scale, w // + self.scale), interpolation=cv2.INTER_AREA)) + lrs.append(cv2.resize(hq_ref, (h // self.scale, w // + self.scale), interpolation=cv2.INTER_AREA)) + lms.append(cv2.resize(hq_mask, (h // self.scale, w // + self.scale), interpolation=cv2.INTER_AREA)) + lcs.append(self.resize_point( + hq_center, (h, w), ls[0].shape[:2])) + # Corrupt the LQ image (only in eval mode) + if not self.corrupt_before_downsize and not self.for_eval: + ls = self.corruptor.corrupt_images(ls) + return ls, lrs, lms, lcs + + def __len__(self): + return self.len diff --git a/dlas/data/images/byol_attachment.py b/dlas/data/images/byol_attachment.py new file mode 100644 index 0000000..6434637 --- /dev/null +++ b/dlas/data/images/byol_attachment.py @@ -0,0 +1,392 @@ +import random +from time import time + +import kornia +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchvision +from kornia import augmentation as augs +from kornia import filters, geometry +from torch.utils.data import Dataset +# Wrapper for a DLAS Dataset class that applies random augmentations from the BYOL paper to BOTH the 'lq' and 'hq' +# inputs. These are then outputted as 'aug1' and 'aug2'. +from tqdm import tqdm + +from dlas.data import create_dataset +from dlas.models.arch_util import PixelUnshuffle +from dlas.utils.util import opt_get + + +class RandomApply(nn.Module): + def __init__(self, fn, p): + super().__init__() + self.fn = fn + self.p = p + + def forward(self, x): + if random.random() > self.p: + return x + return self.fn(x) + + +class ByolDatasetWrapper(Dataset): + def __init__(self, opt): + super().__init__() + self.wrapped_dataset = create_dataset(opt['dataset']) + self.cropped_img_size = opt['crop_size'] + self.key1 = opt_get(opt, ['key1'], 'hq') + self.key2 = opt_get(opt, ['key2'], 'lq') + # When set, color alterations and blurs are disabled. + for_sr = opt_get(opt, ['for_sr'], False) + + augmentations = [ + augs.RandomHorizontalFlip(), + augs.RandomResizedCrop((self.cropped_img_size, self.cropped_img_size))] + if not for_sr: + augmentations.extend([RandomApply(augs.ColorJitter(0.8, 0.8, 0.8, 0.2), p=0.8), + augs.RandomGrayscale(p=0.2), + RandomApply(filters.GaussianBlur2d((3, 3), (1.5, 1.5)), p=0.1)]) + if opt['normalize']: + # The paper calls for normalization. Most datasets/models in this repo don't use this. + # Recommend setting true if you want to train exactly like the paper. + augmentations.append(augs.Normalize(mean=torch.tensor( + [0.485, 0.456, 0.406]), std=torch.tensor([0.229, 0.224, 0.225]))) + self.aug = nn.Sequential(*augmentations) + + def __getitem__(self, item): + item = self.wrapped_dataset[item] + item.update({'aug1': self.aug(item[self.key1]).squeeze( + dim=0), 'aug2': self.aug(item[self.key2]).squeeze(dim=0)}) + return item + + def __len__(self): + return len(self.wrapped_dataset) + + +# Basically the same as ByolDatasetWrapper except only produces 1 augmentation and stores in the 'lr' key. Also applies +# crop&resize to 2D tensors in the state dict with the word "label" in them. +class DatasetRandomAugWrapper(Dataset): + def __init__(self, opt): + super().__init__() + self.wrapped_dataset = create_dataset(opt['dataset']) + self.cropped_img_size = opt['crop_size'] + self.includes_labels = opt['includes_labels'] + augmentations = [ + RandomApply(augs.ColorJitter(0.4, 0.4, 0.4, 0.2), p=0.8), + augs.RandomGrayscale(p=0.2), + RandomApply(filters.GaussianBlur2d((3, 3), (1.5, 1.5)), p=0.1)] + self.aug = nn.Sequential(*augmentations) + self.rrc = nn.Sequential(*[ + augs.RandomHorizontalFlip(), + augs.RandomResizedCrop((self.cropped_img_size, self.cropped_img_size))]) + + def __getitem__(self, item): + item = self.wrapped_dataset[item] + hq = self.aug(item['hq'].unsqueeze(0)).squeeze(0) + labels = [] + dtypes = [] + for k in item.keys(): + if 'label' in k and isinstance(item[k], torch.Tensor) and len(item[k].shape) == 3: + # Only supports a channel dim of 1. + assert item[k].shape[0] == 1 + labels.append(k) + dtypes.append(item[k].dtype) + hq = torch.cat([hq, item[k].type(torch.float)], dim=0) + hq = self.rrc(hq.unsqueeze(0)).squeeze(0) + for i, k in enumerate(labels): + # Strip out any label values that are not whole numbers. + item[k] = hq[3+i:3+i+1, :, :] + whole = (item[k].round() == item[k]) + item[k] = item[k] * whole + item[k] = item[k].type(dtypes[i]) + item['lq'] = hq[:3, :, :] + item['hq'] = item['lq'] + return item + + def __len__(self): + return len(self.wrapped_dataset) + + +def test_dataset_random_aug_wrapper(): + opt = { + 'dataset': { + 'mode': 'imagefolder', + 'name': 'amalgam', + 'paths': ['F:\\4k6k\\datasets\\ns_images\\512_unsupervised'], + 'weights': [1], + 'target_size': 512, + 'force_multiple': 1, + 'scale': 1, + 'fixed_corruptions': ['jpeg-broad'], + 'random_corruptions': ['noise-5', 'none'], + 'num_corrupts_per_image': 1, + 'corrupt_before_downsize': False, + 'labeler': { + 'type': 'patch_labels', + 'label_file': 'F:\\4k6k\\datasets\\ns_images\\512_unsupervised\\categories.json' + } + }, + 'crop_size': 512, + 'includes_labels': True, + } + + ds = DatasetRandomAugWrapper(opt) + import os + os.makedirs("debug", exist_ok=True) + for i in tqdm(range(0, len(ds))): + o = ds[random.randint(0, len(ds)-1)] + for k, v in o.items(): + # 'lq', 'hq', 'aug1', 'aug2', + if k in ['hq']: + torchvision.utils.save_image( + v.unsqueeze(0), "debug/%i_%s.png" % (i, k)) + masked = v * (o['labels_mask'] * .5 + .5) + # torchvision.utils.save_image(masked.unsqueeze(0), "debug/%i_%s_masked.png" % (i, k)) + # Pick a random (non-zero) label and spit it out with the textual label. + if len(o['labels'].unique()) > 1: + randlbl = np.random.choice(o['labels'].unique()[1:]) + moremask = v * ((1*(o['labels'] == randlbl))*.5+.5) + torchvision.utils.save_image(moremask.unsqueeze( + 0), "debug/%i_%s_%s.png" % (i, k, o['label_strings'][randlbl])) + + +def no_batch_interpolate(i, size, mode): + i = i.unsqueeze(0) + i = F.interpolate(i, size=size, mode=mode) + return i.squeeze(0) + + +# Performs a 1-d translation of "other": +# If other (self.multiple*3) + d = d // self.multiple + + # Step 2 + base_w = random.randint(d//2+1, d-1) + base_l = random.randint(0, d-base_w) + base_h = random.randint(base_w-1, base_w+1) + base_t = random.randint(0, d-base_h) + base_r, base_b = base_l+base_w, base_t+base_h + + # Step 3 + im2_w = random.randint(d//2+1, d-1) + im2_l = random.randint(0, d-im2_w) + im2_h = random.randint(im2_w-1, im2_w+1) + im2_t = random.randint(0, d-im2_h) + im2_r, im2_b = im2_l+im2_w, im2_t+im2_h + + # Step 4 + m = self.multiple + jl, jt = random.randint(-self.jitter_range, + self.jitter_range), random.randint(-self.jitter_range, self.jitter_range) + # If the top of a patch is zero, a negative jitter will cause it to go negative. + jt = jt if base_t != 0 else abs(jt) + # Likewise, jitter shouldn't allow the patch to go over-bounds. + jt = jt if (base_t+base_h)*m != i1.shape[1] else 0 + jl = jl if base_l != 0 else abs(jl) + jl = jl if (base_l+base_w)*m != i1.shape[1] else 0 + p1 = i1[:, base_t*m+jt:(base_t+base_h)*m+jt, + base_l*m+jl:(base_l+base_w)*m+jl] + p1_resized = no_batch_interpolate(p1, size=(d*m, d*m), mode="bilinear") + jl, jt = random.randint(-self.jitter_range, + self.jitter_range), random.randint(-self.jitter_range, self.jitter_range) + jt = jt if im2_t != 0 else abs(jt) + jt = jt if (im2_t+im2_h)*m != i2.shape[1] else 0 + jl = jl if im2_l != 0 else abs(jl) + jl = jl if (im2_l+im2_w)*m != i2.shape[1] else 0 + p2 = i2[:, im2_t*m+jt:(im2_t+im2_h)*m+jt, + im2_l*m+jl:(im2_l+im2_w)*m+jl] + p2_resized = no_batch_interpolate(p2, size=(d*m, d*m), mode="bilinear") + + # Step 5 + should_flip = random.random() < .5 + if should_flip: + should_flip = 1 + p2_resized = geometry.transform.hflip(p2_resized) + else: + should_flip = 0 + + # Step 6 + i1_shared_t, i1_shared_l = snap(base_t, im2_t), snap(base_l, im2_l) + i2_shared_t, i2_shared_l = snap(im2_t, base_t), snap(im2_l, base_l) + ix_h = min(base_b, im2_b) - max(base_t, im2_t) + ix_w = min(base_r, im2_r) - max(base_l, im2_l) + recompute_package = torch.tensor([d, base_h, base_w, i1_shared_t, i1_shared_l, im2_h, + im2_w, i2_shared_t, i2_shared_l, should_flip, ix_h, ix_w], dtype=torch.long) + + # Step 7 + mask1 = torch.full((1, base_h*m, base_w*m), fill_value=.5) + mask1[:, i1_shared_t*m:(i1_shared_t+ix_h)*m, + i1_shared_l*m:(i1_shared_l+ix_w)*m] = 1 + masked1 = pad_to(p1 * mask1, d*m) + mask2 = torch.full((1, im2_h*m, im2_w*m), fill_value=.5) + mask2[:, i2_shared_t*m:(i2_shared_t+ix_h)*m, + i2_shared_l*m:(i2_shared_l+ix_w)*m] = 1 + masked2 = pad_to(p2 * mask2, d*m) + mask = torch.full((1, d*m, d*m), fill_value=.33) + mask[:, base_t*m:(base_t+base_w)*m, base_l*m:(base_l+base_h)*m] += .33 + mask[:, im2_t*m:(im2_t+im2_w)*m, im2_l*m:(im2_l+im2_h)*m] += .33 + masked_dbg = i1 * mask + + # Step 8 - Rebuild shared regions for testing purposes. + p1_shuf, p2_shuf = PixelUnshuffle(self.multiple)(p1_resized.unsqueeze(0)), \ + PixelUnshuffle(self.multiple)(p2_resized.unsqueeze(0)) + i1_shared, i2_shared = reconstructed_shared_regions( + p1_shuf, p2_shuf, recompute_package.unsqueeze(0)) + i1_shared = pad_to(nn.PixelShuffle(self.multiple) + (i1_shared).squeeze(0), d * m) + i2_shared = pad_to(nn.PixelShuffle(self.multiple) + (i2_shared).squeeze(0), d*m) + + return p1_resized, p2_resized, recompute_package, masked1, masked2, masked_dbg, i1_shared, i2_shared + + +# Uses the recompute package returned from the above dataset to extract matched-size "similar regions" from two feature +# maps. +def reconstructed_shared_regions(fea1, fea2, recompute_package: torch.Tensor): + package = recompute_package.cpu() + res1 = [] + res2 = [] + pad_dim = torch.max(package[:, -2:]).item() + # It'd be real nice if we could do this at the batch level, but I don't see a really good way to do that outside + # of conforming the recompute_package across the entire batch. + for b in range(package.shape[0]): + expected_dim, f1_h, f1_w, f1s_t, f1s_l, f2_h, f2_w, f2s_t, f2s_l, should_flip, s_h, s_w = tuple( + package[b].tolist()) + # If you are hitting this assert, you specified `latent_multiple` in your dataset config wrong. + assert expected_dim == fea1.shape[2] and expected_dim == fea2.shape[2] + + # Unflip 2 if needed. + f2 = fea2[b] + if should_flip == 1: + f2 = kornia.geometry.transform.hflip(f2) + # Resize the input features to match + f1s = F.interpolate(fea1[b].unsqueeze(0), (f1_h, f1_w), mode="nearest") + f2s = F.interpolate(f2.unsqueeze(0), (f2_h, f2_w), mode="nearest") + # Outputs must be padded so they can "get along" with each other. + res1.append( + pad_to(f1s[:, :, f1s_t:f1s_t+s_h, f1s_l:f1s_l+s_w], pad_dim)) + res2.append( + pad_to(f2s[:, :, f2s_t:f2s_t+s_h, f2s_l:f2s_l+s_w], pad_dim)) + return torch.cat(res1, dim=0), torch.cat(res2, dim=0) + + +# Follows the general template of BYOL dataset, with the following changes: +# 1. Flip() is not applied. +# 2. Instead of RandomResizedCrop, a custom Transform, RandomSharedRegionCrop is used. +# 3. The dataset injects two integer tensors alongside the augmentations, which are used to index image regions shared +# by the joint augmentations. +# 4. The dataset injects an aug_shared_view for debugging purposes. +class StructuredCropDatasetWrapper(Dataset): + def __init__(self, opt): + super().__init__() + self.wrapped_dataset = create_dataset(opt['dataset']) + augmentations = [RandomApply(augs.ColorJitter(0.8, 0.8, 0.8, 0.2), p=0.8), + augs.RandomGrayscale(p=0.2), + RandomApply(filters.GaussianBlur2d((3, 3), (1.5, 1.5)), p=0.1)] + self.aug = nn.Sequential(*augmentations) + self.rrc = RandomSharedRegionCrop( + opt['latent_multiple'], opt_get(opt, ['jitter_range'], 0)) + + def __getitem__(self, item): + item = self.wrapped_dataset[item] + a1 = self.aug(item['hq']).squeeze(dim=0) + a2 = self.aug(item['lq']).squeeze(dim=0) + a1, a2, sr_dim, m1, m2, db, i1s, i2s = self.rrc(a1, a2) + item.update({'aug1': a1, 'aug2': a2, 'similar_region_dimensions': sr_dim, + 'masked1': m1, 'masked2': m2, 'aug_shared_view': db, + 'i1_shared': i1s, 'i2_shared': i2s}) + return item + + def __len__(self): + return len(self.wrapped_dataset) + + +# For testing this dataset. +def test_structured_crop_dataset_wrapper(): + opt = { + 'dataset': + { + 'mode': 'imagefolder', + 'name': 'amalgam', + 'paths': ['F:\\4k6k\\datasets\\ns_images\\512_unsupervised'], + 'weights': [1], + 'target_size': 256, + 'force_multiple': 32, + 'scale': 1, + 'fixed_corruptions': ['jpeg-broad', 'gaussian_blur'], + 'random_corruptions': ['noise-5', 'none'], + 'num_corrupts_per_image': 1, + 'corrupt_before_downsize': True, + }, + 'latent_multiple': 16, + 'jitter_range': 0, + } + + ds = StructuredCropDatasetWrapper(opt) + import os + os.makedirs("debug", exist_ok=True) + for i in tqdm(range(0, len(ds))): + o = ds[random.randint(0, len(ds)-1)] + # for k, v in o.items(): + # 'lq', 'hq', 'aug1', 'aug2', + # if k in [ 'aug_shared_view', 'masked1', 'masked2']: + # torchvision.utils.save_image(v.unsqueeze(0), "debug/%i_%s.png" % (i, k)) + rcpkg = o['similar_region_dimensions'] + pixun = PixelUnshuffle(16) + pixsh = nn.PixelShuffle(16) + rc1, rc2 = reconstructed_shared_regions(pixun(o['aug1'].unsqueeze( + 0)), pixun(o['aug2'].unsqueeze(0)), rcpkg.unsqueeze(0)) + # torchvision.utils.save_image(pixsh(rc1), "debug/%i_rc1.png" % (i,)) + # torchvision.utils.save_image(pixsh(rc2), "debug/%i_rc2.png" % (i,)) + + +if __name__ == '__main__': + test_dataset_random_aug_wrapper() diff --git a/dlas/data/images/chunk_with_reference.py b/dlas/data/images/chunk_with_reference.py new file mode 100644 index 0000000..f8341cb --- /dev/null +++ b/dlas/data/images/chunk_with_reference.py @@ -0,0 +1,58 @@ +import os.path as osp + +import numpy as np +import torch + +from dlas.data import util +# Iterable that reads all the images in a directory that contains a reference image, tile images and center coordinates. +from dlas.utils.util import opt_get + + +class ChunkWithReference: + def __init__(self, opt, path): + self.path = path.path + self.tiles, _ = util.find_files_of_type('img', self.path) + self.need_metadata = opt_get(opt, ['strict'], False) or opt_get( + opt, ['needs_metadata'], False) + self.need_ref = opt_get(opt, ['need_ref'], False) + if 'ignore_first' in opt.keys(): + self.tiles = self.tiles[opt['ignore_first']:] + + # Odd failures occur at times. Rather than crashing, report the error and just return zeros. + def read_image_or_get_zero(self, img_path): + img = util.read_img(None, img_path, rgb=True) + if img is None: + return np.zeros(128, 128, 3) + return img + + def __getitem__(self, item): + tile = self.read_image_or_get_zero(self.tiles[item]) + if self.need_ref and osp.exists(osp.join(self.path, "ref.jpg")): + tile_id = int(osp.splitext(osp.basename(self.tiles[item]))[0]) + ref = self.read_image_or_get_zero(osp.join(self.path, "ref.jpg")) + if self.need_metadata: + centers = torch.load(osp.join(self.path, "centers.pt")) + if tile_id in centers.keys(): + center, tile_width = centers[tile_id] + else: + print("Could not find the given tile id in the accompanying centers.pt. This generally means that " + "centers.pt was overwritten at some point e.g. by duplicate data. If you don't care about tile " + "centers, consider passing strict=false to the dataset options. (Note: you must re-build your" + "caches for this setting change to take effect.)") + raise FileNotFoundError(tile_id, self.tiles[item]) + else: + center = torch.tensor([128, 128], dtype=torch.long) + tile_width = 256 + mask = np.full(tile.shape[:2] + (1,), + fill_value=.1, dtype=tile.dtype) + mask[center[0] - tile_width // 2:center[0] + tile_width // 2, + center[1] - tile_width // 2:center[1] + tile_width // 2] = 1 + else: + ref = np.zeros_like(tile) + mask = np.zeros(tile.shape[:2] + (1,)) + center = (0, 0) + + return tile, ref, center, mask, self.tiles[item] + + def __len__(self): + return len(self.tiles) diff --git a/dlas/data/images/cifar.py b/dlas/data/images/cifar.py new file mode 100644 index 0000000..538bc07 --- /dev/null +++ b/dlas/data/images/cifar.py @@ -0,0 +1,179 @@ +# A copy of the cifar dataset from torch which also returns coarse labels. + +import os +import os.path +import pickle +from typing import Any, Callable, Optional, Tuple + +import numpy as np +from PIL import Image +from torchvision.datasets import VisionDataset +from torchvision.datasets.utils import (check_integrity, + download_and_extract_archive) + + +class CIFAR10(VisionDataset): + """`CIFAR10 `_ Dataset. + + Args: + root (string): Root directory of dataset where directory + ``cifar-10-batches-py`` exists or will be saved to if download is set to True. + train (bool, optional): If True, creates dataset from training set, otherwise + creates from test set. + transform (callable, optional): A function/transform that takes in an PIL image + and returns a transformed version. E.g, ``transforms.RandomCrop`` + target_transform (callable, optional): A function/transform that takes in the + target and transforms it. + download (bool, optional): If true, downloads the dataset from the internet and + puts it in root directory. If dataset is already downloaded, it is not + downloaded again. + + """ + base_folder = 'cifar-10-batches-py' + url = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz" + filename = "cifar-10-python.tar.gz" + tgz_md5 = 'c58f30108f718f92721af3b95e74349a' + train_list = [ + ['data_batch_1', 'c99cafc152244af753f735de768cd75f'], + ['data_batch_2', 'd4bba439e000b95fd0a9bffe97cbabec'], + ['data_batch_3', '54ebc095f3ab1f0389bbae665268c751'], + ['data_batch_4', '634d18415352ddfa80567beed471001a'], + ['data_batch_5', '482c414d41f54cd18b22e5b47cb7c3cb'], + ] + + test_list = [ + ['test_batch', '40351d587109b95175f43aff81a1287e'], + ] + meta = { + 'filename': 'batches.meta', + 'key': 'label_names', + 'md5': '5ff9c542aee3614f3951f8cda6e48888', + } + + def __init__( + self, + root: str, + train: bool = True, + transform: Optional[Callable] = None, + target_transform: Optional[Callable] = None, + download: bool = False, + ) -> None: + + super(CIFAR10, self).__init__(root, transform=transform, + target_transform=target_transform) + + self.train = train # training set or test set + + if download: + self.download() + + if not self._check_integrity(): + raise RuntimeError('Dataset not found or corrupted.' + + ' You can use download=True to download it') + + if self.train: + downloaded_list = self.train_list + else: + downloaded_list = self.test_list + + self.data: Any = [] + self.targets = [] + self.coarse_targets = [] + + # now load the picked numpy arrays + for file_name, checksum in downloaded_list: + file_path = os.path.join(self.root, self.base_folder, file_name) + with open(file_path, 'rb') as f: + entry = pickle.load(f, encoding='latin1') + self.data.append(entry['data']) + if 'labels' in entry: + self.targets.extend(entry['labels']) + else: + self.targets.extend(entry['fine_labels']) + self.coarse_targets.extend(entry['coarse_labels']) + + self.data = np.vstack(self.data).reshape(-1, 3, 32, 32) + self.data = self.data.transpose((0, 2, 3, 1)) # convert to HWC + + self._load_meta() + + def _load_meta(self) -> None: + path = os.path.join(self.root, self.base_folder, self.meta['filename']) + if not check_integrity(path, self.meta['md5']): + raise RuntimeError('Dataset metadata file not found or corrupted.' + + ' You can use download=True to download it') + with open(path, 'rb') as infile: + data = pickle.load(infile, encoding='latin1') + self.classes = data[self.meta['key']] + self.class_to_idx = {_class: i for i, + _class in enumerate(self.classes)} + + def __getitem__(self, index: int) -> Tuple[Any, Any]: + """ + Args: + index (int): Index + + Returns: + tuple: (image, target) where target is index of the target class. + """ + img, target = self.data[index], self.targets[index] + + # doing this so that it is consistent with all other datasets + # to return a PIL Image + img = Image.fromarray(img) + + if self.transform is not None: + img = self.transform(img) + + if self.target_transform is not None: + target = self.target_transform(target) + + if len(self.coarse_targets) > 0: + return img, target, self.coarse_targets[index] + + return img, target + + def __len__(self) -> int: + return len(self.data) + + def _check_integrity(self) -> bool: + root = self.root + for fentry in (self.train_list + self.test_list): + filename, md5 = fentry[0], fentry[1] + fpath = os.path.join(root, self.base_folder, filename) + if not check_integrity(fpath, md5): + return False + return True + + def download(self) -> None: + if self._check_integrity(): + print('Files already downloaded and verified') + return + download_and_extract_archive( + self.url, self.root, filename=self.filename, md5=self.tgz_md5) + + def extra_repr(self) -> str: + return "Split: {}".format("Train" if self.train is True else "Test") + + +class CIFAR100(CIFAR10): + """`CIFAR100 `_ Dataset. + + This is a subclass of the `CIFAR10` Dataset. + """ + base_folder = 'cifar-100-python' + url = "https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz" + filename = "cifar-100-python.tar.gz" + tgz_md5 = 'eb9058c3a382ffc7106e4002c42a8d85' + train_list = [ + ['train', '16019d7e3df5f24257cddd939b257f8d'], + ] + + test_list = [ + ['test', 'f0ef6b0ae62326f3e7ffdfab6717acfc'], + ] + meta = { + 'filename': 'meta', + 'key': 'fine_label_names', + 'md5': '7973b15100ade9c7d40fb424638fde48', + } diff --git a/dlas/data/images/full_image_dataset.py b/dlas/data/images/full_image_dataset.py new file mode 100644 index 0000000..8b37281 --- /dev/null +++ b/dlas/data/images/full_image_dataset.py @@ -0,0 +1,416 @@ +import random +from io import BytesIO + +import cv2 +import numpy as np +import torch +import torch.utils.data as data +import torchvision.transforms.functional as F +from PIL import Image, ImageOps + +import dlas.data.util as util + + +# Reads full-quality images and pulls tiles from them. Also extracts LR renderings of the full image with cues as to +# where those tiles came from. +class FullImageDataset(data.Dataset): + """ + Read LQ (Low Quality, e.g. LR (Low Resolution), blurry, etc) and GT image pairs. + If only GT images are provided, generate LQ images on-the-fly. + """ + + def get_lq_path(self, i): + which_lq = random.randint(0, len(self.paths_LQ)-1) + return self.paths_LQ[which_lq][i % len(self.paths_LQ[which_lq])] + + def __init__(self, opt): + super(FullImageDataset, self).__init__() + self.opt = opt + self.data_type = 'img' + self.paths_LQ, self.paths_GT = None, None + self.sizes_LQ, self.sizes_GT = None, None + self.LQ_env, self.GT_env = None, None + self.force_multiple = self.opt['force_multiple'] if 'force_multiple' in self.opt.keys( + ) else 1 + + self.paths_GT, self.sizes_GT = util.find_files_of_type( + self.data_type, opt['dataroot_GT'], opt['dataroot_GT_weights']) + if 'dataroot_LQ' in opt.keys(): + self.paths_LQ = [] + if isinstance(opt['dataroot_LQ'], list): + # Multiple LQ data sources can be given, in case there are multiple ways of corrupting a source image and + # we want the model to learn them all. + for dr_lq in opt['dataroot_LQ']: + lq_path, self.sizes_LQ = util.find_files_of_type( + self.data_type, dr_lq) + self.paths_LQ.append(lq_path) + else: + lq_path, self.sizes_LQ = util.find_files_of_type( + self.data_type, opt['dataroot_LQ']) + self.paths_LQ.append(lq_path) + + assert self.paths_GT, 'Error: GT path is empty.' + self.random_scale_list = [1] + + def motion_blur(self, image, size, angle): + k = np.zeros((size, size), dtype=np.float32) + k[(size - 1) // 2, :] = np.ones(size, dtype=np.float32) + k = cv2.warpAffine(k, cv2.getRotationMatrix2D( + (size / 2 - 0.5, size / 2 - 0.5), angle, 1.0), (size, size)) + k = k * (1.0 / np.sum(k)) + return cv2.filter2D(image, -1, k) + + # Selects the smallest dimension from the image and crops it randomly so the other dimension matches. The cropping + # offset from center is chosen on a normal probability curve. + def get_square_image(self, image): + h, w, _ = image.shape + if h == w: + return image + offset = max(min(np.random.normal(scale=.3), 1.0), -1.0) + if h > w: + diff = h - w + center = diff // 2 + top = int(center + offset * (center - 2)) + return image[top:top+w, :, :] + else: + diff = w - h + center = diff // 2 + left = int(center + offset * (center - 2)) + return image[:, left:left+h, :] + + def pick_along_range(self, sz, r, dev): + margin_sz = sz - r + margin_center = margin_sz // 2 + return min(max(int(min(np.random.normal(scale=dev), 1.0) * margin_sz + margin_center), 0), margin_sz) + + def resize_point(self, point, orig_dim, new_dim): + oh, ow = orig_dim + nh, nw = new_dim + dh, dw = float(nh) / float(oh), float(nw) / float(ow) + point[0] = int(dh * float(point[0])) + point[1] = int(dw * float(point[1])) + return point + + # - Randomly extracts a square from image and resizes it to opt['target_size']. + # - Fills a mask with zeros, then places 1's where the square was extracted from. Resizes this mask and the source + # image to the target_size and returns that too. + # Notes: + # - When extracting a square, the size of the square is randomly distributed [target_size, source_size] along a + # half-normal distribution, biasing towards the target_size. + # - A biased normal distribution is also used to bias the tile selection towards the center of the source image. + def pull_tile(self, image, lq=False): + if lq: + target_sz = self.opt['min_tile_size'] // self.opt['scale'] + else: + target_sz = self.opt['min_tile_size'] + h, w, _ = image.shape + possible_sizes_above_target = h - target_sz + if 'fixed_size' in self.opt.keys() and self.opt['fixed_size']: + square_size = target_sz + else: + tile_expansion_dev = self.opt['tile_scale_normal_stddev'] if 'tile_scale_normal_stddev' in self.opt.keys( + ) else .17 + square_size = int(target_sz + possible_sizes_above_target * + min(np.abs(np.random.normal(scale=tile_expansion_dev)), 1.0)) + + # Pick the left,top coords to draw the patch from + left = self.pick_along_range(w, square_size, .3) + top = self.pick_along_range(w, square_size, .3) + + mask = np.zeros((h, w, 1), dtype=image.dtype) + mask[top:top+square_size, left:left+square_size] = 1 + patch = image[top:top+square_size, left:left+square_size, :] + center = torch.tensor( + [top + square_size // 2, left + square_size // 2], dtype=torch.long) + + patch = cv2.resize(patch, (target_sz, target_sz), + interpolation=cv2.INTER_LINEAR) + image = cv2.resize(image, (target_sz, target_sz), + interpolation=cv2.INTER_LINEAR) + mask = cv2.resize(mask, (target_sz, target_sz), + interpolation=cv2.INTER_LINEAR) + center = self.resize_point(center, (h, w), image.shape[:2]) + + return patch, image, mask, center + + def augment_tile(self, img_GT, img_LQ, strength=1): + scale = self.opt['scale'] + GT_size = self.opt['target_size'] + + H, W, _ = img_GT.shape + assert H >= GT_size and W >= GT_size + + LQ_size = GT_size // scale + img_LQ = cv2.resize(img_LQ, (LQ_size, LQ_size), + interpolation=cv2.INTER_LINEAR) + img_GT = cv2.resize(img_GT, (GT_size, GT_size), + interpolation=cv2.INTER_LINEAR) + + if self.opt['use_blurring']: + # Pick randomly between gaussian, motion, or no blur. + blur_det = random.randint(0, 100) + blur_magnitude = 3 if 'blur_magnitude' not in self.opt.keys( + ) else self.opt['blur_magnitude'] + blur_magnitude = max(1, int(blur_magnitude*strength)) + if blur_det < 40: + blur_sig = int(random.randrange(0, int(blur_magnitude))) + img_LQ = cv2.GaussianBlur( + img_LQ, (blur_magnitude, blur_magnitude), blur_sig) + elif blur_det < 70: + img_LQ = self.motion_blur(img_LQ, random.randrange( + 1, int(blur_magnitude) * 3), random.randint(0, 360)) + + return img_GT, img_LQ + + # Converts img_LQ to PIL and performs JPG compression corruptions and grayscale on the image, then returns it. + def pil_augment(self, img_LQ, strength=1): + img_LQ = (img_LQ * 255).astype(np.uint8) + img_LQ = Image.fromarray(img_LQ) + if self.opt['use_compression_artifacts'] and random.random() > .25: + sub_lo = 90 * strength + sub_hi = 30 * strength + qf = random.randrange(100 - sub_lo, 100 - sub_hi) + corruption_buffer = BytesIO() + img_LQ.save(corruption_buffer, "JPEG", quality=qf, optimice=True) + corruption_buffer.seek(0) + img_LQ = Image.open(corruption_buffer) + + if 'grayscale' in self.opt.keys() and self.opt['grayscale']: + img_LQ = ImageOps.grayscale(img_LQ).convert('RGB') + + return img_LQ + + def perform_random_hr_augment(self, image, aug_code=None, augmentations=1): + if aug_code is None: + aug_code = [random.randint(0, 10) for _ in range(augmentations)] + else: + assert augmentations == 1 + aug_code = [aug_code] + if 0 in aug_code: + # Color quantization + pass + elif 1 in aug_code: + # Gaussian Blur (point or motion) + blur_magnitude = 3 + blur_sig = int(random.randrange(0, int(blur_magnitude))) + image = cv2.GaussianBlur( + image, (blur_magnitude, blur_magnitude), blur_sig) + elif 2 in aug_code: + # Median Blur + image = cv2.medianBlur(image, 3) + elif 3 in aug_code: + # Motion blur + image = self.motion_blur( + image, random.randrange(1, 9), random.randint(0, 360)) + elif 4 in aug_code: + # Smooth blur + image = cv2.blur(image, ksize=3) + elif 5 in aug_code: + # Block noise + pass + elif 6 in aug_code: + # Bicubic LR->HR + pass + elif 7 in aug_code: + # Linear compression distortion + pass + elif 8 in aug_code: + # Interlacing distortion + pass + elif 9 in aug_code: + # Chromatic aberration + pass + elif 10 in aug_code: + # Noise + pass + elif 11 in aug_code: + # JPEG compression + pass + elif 12 in aug_code: + # Lightening / saturation + pass + return image + + def __getitem__(self, index): + scale = self.opt['scale'] + + # get full size image + full_path = self.paths_GT[index % len(self.paths_GT)] + LQ_path = full_path + img_full = util.read_img(None, full_path, None) + img_full = util.channel_convert( + img_full.shape[2], 'RGB', [img_full])[0] + if self.opt['phase'] == 'train': + img_full = util.augment( + [img_full], self.opt['use_flip'], self.opt['use_rot'])[0] + img_full = self.get_square_image(img_full) + img_GT, gt_fullsize_ref, gt_mask, gt_center = self.pull_tile( + img_full) + else: + img_GT, gt_fullsize_ref = img_full, img_full + gt_mask = np.ones(img_full.shape[:2], dtype=gt_fullsize_ref.dtype) + gt_center = torch.tensor( + [img_full.shape[0] // 2, img_full.shape[1] // 2], dtype=torch.long) + orig_gt_dim = gt_fullsize_ref.shape[:2] + + # get LQ image + if self.paths_LQ: + LQ_path = self.get_lq_path(index) + img_lq_full = util.read_img(None, LQ_path, None) + if self.opt['phase'] == 'train': + img_lq_full = util.augment( + [img_lq_full], self.opt['use_flip'], self.opt['use_rot'])[0] + img_lq_full = self.get_square_image(img_lq_full) + img_LQ, lq_fullsize_ref, lq_mask, lq_center = self.pull_tile( + img_lq_full, lq=True) + else: + img_LQ, lq_fullsize_ref = img_lq_full, img_lq_full + lq_mask = np.ones( + img_lq_full.shape[:2], dtype=lq_fullsize_ref.dtype) + lq_center = torch.tensor( + [img_lq_full.shape[0] // 2, img_lq_full.shape[1] // 2], dtype=torch.long) + else: # down-sampling on-the-fly + # randomly scale during training + if self.opt['phase'] == 'train': + GT_size = self.opt['target_size'] + random_scale = random.choice(self.random_scale_list) + if len(img_GT.shape) == 2: + print("ERRAR:") + print(img_GT.shape) + print(full_path) + H_s, W_s, _ = img_GT.shape + + def _mod(n, random_scale, scale, thres): + rlt = int(n * random_scale) + rlt = (rlt // scale) * scale + return thres if rlt < thres else rlt + + H_s = _mod(H_s, random_scale, scale, GT_size) + W_s = _mod(W_s, random_scale, scale, GT_size) + img_GT = cv2.resize(img_GT, (W_s, H_s), + interpolation=cv2.INTER_LINEAR) + if img_GT.ndim == 2: + img_GT = cv2.cvtColor(img_GT, cv2.COLOR_GRAY2BGR) + + H, W, _ = img_GT.shape + + # using matlab imresize + img_LQ = util.imresize_np(img_GT, 1 / scale, True) + lq_fullsize_ref = util.imresize_np( + gt_fullsize_ref, 1 / scale, True) + if img_LQ.ndim == 2: + img_LQ = np.expand_dims(img_LQ, axis=2) + lq_mask, lq_center = gt_mask, self.resize_point( + gt_center.clone(), orig_gt_dim, lq_fullsize_ref.shape[:2]) + orig_lq_dim = lq_fullsize_ref.shape[:2] + + # Enforce force_resize constraints via clipping. + h, w, _ = img_LQ.shape + if h % self.force_multiple != 0 or w % self.force_multiple != 0: + h, w = (h - h % self.force_multiple), (w - w % self.force_multiple) + img_LQ = img_LQ[:h, :w, :] + lq_fullsize_ref = lq_fullsize_ref[:h, :w, :] + h *= scale + w *= scale + img_GT = img_GT[:h, :w] + gt_fullsize_ref = gt_fullsize_ref[:h, :w, :] + + if self.opt['phase'] == 'train': + img_GT, img_LQ = self.augment_tile(img_GT, img_LQ) + gt_fullsize_ref, lq_fullsize_ref = self.augment_tile( + gt_fullsize_ref, lq_fullsize_ref, strength=.2) + + # Scale masks. + lq_mask = cv2.resize( + lq_mask, (lq_fullsize_ref.shape[1], lq_fullsize_ref.shape[0]), interpolation=cv2.INTER_LINEAR) + gt_mask = cv2.resize( + gt_mask, (gt_fullsize_ref.shape[1], gt_fullsize_ref.shape[0]), interpolation=cv2.INTER_LINEAR) + + # Scale center coords + lq_center = self.resize_point( + lq_center, orig_lq_dim, lq_fullsize_ref.shape[:2]) + gt_center = self.resize_point( + gt_center, orig_gt_dim, gt_fullsize_ref.shape[:2]) + + # BGR to RGB, HWC to CHW, numpy to tensor + if img_GT.shape[2] == 3: + img_GT = cv2.cvtColor(img_GT, cv2.COLOR_BGR2RGB) + img_LQ = cv2.cvtColor(img_LQ, cv2.COLOR_BGR2RGB) + lq_fullsize_ref = cv2.cvtColor(lq_fullsize_ref, cv2.COLOR_BGR2RGB) + gt_fullsize_ref = cv2.cvtColor(gt_fullsize_ref, cv2.COLOR_BGR2RGB) + + # LQ needs to go to a PIL image to perform the compression-artifact transformation. + # if self.opt['phase'] == 'train': + # img_LQ = self.pil_augment(img_LQ) + # lq_fullsize_ref = self.pil_augment(lq_fullsize_ref, strength=.2) + + img_GT = torch.from_numpy(np.ascontiguousarray( + np.transpose(img_GT, (2, 0, 1)))).float() + gt_fullsize_ref = torch.from_numpy(np.ascontiguousarray( + np.transpose(gt_fullsize_ref, (2, 0, 1)))).float() + img_LQ = F.to_tensor(img_LQ) + lq_fullsize_ref = F.to_tensor(lq_fullsize_ref) + lq_mask = torch.from_numpy( + np.ascontiguousarray(lq_mask)).unsqueeze(dim=0) + gt_mask = torch.from_numpy( + np.ascontiguousarray(gt_mask)).unsqueeze(dim=0) + + if 'lq_noise' in self.opt.keys(): + lq_noise = torch.randn_like(img_LQ) * self.opt['lq_noise'] / 255 + img_LQ += lq_noise + lq_fullsize_ref += lq_noise + + # Apply the masks to the full images. + gt_fullsize_ref = torch.cat([gt_fullsize_ref, gt_mask], dim=0) + lq_fullsize_ref = torch.cat([lq_fullsize_ref, lq_mask], dim=0) + + d = {'lq': img_LQ, 'hq': img_GT, 'gt_fullsize_ref': gt_fullsize_ref, 'lq_fullsize_ref': lq_fullsize_ref, + 'lq_center': lq_center, 'gt_center': gt_center, + 'LQ_path': LQ_path, 'HQ_path': full_path} + return d + + def __len__(self): + return len(self.paths_GT) + + +if __name__ == '__main__': + ''' + opt = { + 'name': 'amalgam', + 'dataroot_GT': ['F:\\4k6k\\datasets\\ns_images\\imagesets\\images'], + 'dataroot_GT_weights': [1], + 'use_flip': True, + 'use_compression_artifacts': True, + 'use_blurring': True, + 'use_rot': True, + 'lq_noise': 5, + 'target_size': 128, + 'min_tile_size': 256, + 'scale': 2, + 'phase': 'train' + } + ''' + opt = { + 'name': 'amalgam', + 'dataroot_GT': ['F:\\4k6k\\datasets\\ns_images\\imagesets\\images'], + 'dataroot_GT_weights': [1], + 'force_multiple': 32, + 'scale': 2, + 'phase': 'test' + } + + ds = FullImageDataset(opt) + import os + os.makedirs("debug", exist_ok=True) + for i in range(300, len(ds)): + print(i) + o = ds[i] + for k, v in o.items(): + if 'path' not in k: + # if 'full' in k: + # masked = v[:3, :, :] * v[3] + # torchvision.utils.save_image(masked.unsqueeze(0), "debug/%i_%s_masked.png" % (i, k)) + # v = v[:3, :, :] + # import torchvision + # torchvision.utils.save_image(v.unsqueeze(0), "debug/%i_%s.png" % (i, k)) + pass diff --git a/dlas/data/images/image_corruptor.py b/dlas/data/images/image_corruptor.py new file mode 100644 index 0000000..726fabd --- /dev/null +++ b/dlas/data/images/image_corruptor.py @@ -0,0 +1,216 @@ +import functools +import random +from io import BytesIO +from math import cos, pi + +import cv2 +import kornia +import numpy as np +import torch +from data.util import read_img +from kornia.augmentation import ColorJitter +from PIL import Image + +# Get a rough visualization of the above distribution. (Y-axis is meaningless, just spreads data) +from dlas.utils.util import opt_get + +''' +if __name__ == '__main__': + import numpy as np + import matplotlib.pyplot as plt + data = np.asarray([get_rand() for _ in range(5000)]) + plt.plot(data, np.random.uniform(size=(5000,)), 'x') + plt.show() +''' + + +def kornia_color_jitter_numpy(img, setting): + if setting * 255 > 1: + # I'm using Kornia's ColorJitter, which requires pytorch arrays in b,c,h,w format. + img = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0) + img = ColorJitter(setting, setting, setting, setting)(img) + img = img.squeeze(0).permute(1, 2, 0).numpy() + return img + + +# Performs image corruption on a list of images from a configurable set of corruption +# options. +class ImageCorruptor: + def __init__(self, opt): + self.opt = opt + self.reset_random() + self.blur_scale = opt['corruption_blur_scale'] if 'corruption_blur_scale' in opt.keys( + ) else 1 + self.fixed_corruptions = opt['fixed_corruptions'] if 'fixed_corruptions' in opt.keys() else [ + ] + self.num_corrupts = opt['num_corrupts_per_image'] if 'num_corrupts_per_image' in opt.keys( + ) else 0 + self.cosine_bias = opt_get(opt, ['cosine_bias'], True) + if self.num_corrupts == 0: + return + else: + self.random_corruptions = opt['random_corruptions'] if 'random_corruptions' in opt.keys() else [ + ] + + def reset_random(self): + if 'random_seed' in self.opt.keys(): + self.rand = random.Random(self.opt['random_seed']) + else: + self.rand = random.Random() + + # Feeds a random uniform through a cosine distribution to slightly bias corruptions towards "uncorrupted". + # Return is on [0,1] with a bias towards 0. + def get_rand(self): + r = self.rand.random() + if self.cosine_bias: + return 1 - cos(r * pi / 2) + else: + return r + + def corrupt_images(self, imgs, return_entropy=False): + if self.num_corrupts == 0 and not self.fixed_corruptions: + if return_entropy: + return imgs, [] + else: + return imgs + + if self.num_corrupts == 0: + augmentations = [] + else: + augmentations = random.choices( + self.random_corruptions, k=self.num_corrupts) + + # Sources of entropy + corrupted_imgs = [] + entropy = [] + undo_fns = [] + applied_augs = augmentations + self.fixed_corruptions + for img in imgs: + for aug in augmentations: + r = self.get_rand() + img, undo_fn = self.apply_corruption(img, aug, r, applied_augs) + if undo_fn is not None: + undo_fns.append(undo_fn) + for aug in self.fixed_corruptions: + r = self.get_rand() + img, undo_fn = self.apply_corruption(img, aug, r, applied_augs) + entropy.append(r) + if undo_fn is not None: + undo_fns.append(undo_fn) + # Apply undo_fns after all corruptions are finished, in same order. + for ufn in undo_fns: + img = ufn(img) + corrupted_imgs.append(img) + + if return_entropy: + return corrupted_imgs, entropy + else: + return corrupted_imgs + + def apply_corruption(self, img, aug, rand_val, applied_augmentations): + undo_fn = None + if 'color_quantization' in aug: + # Color quantization + quant_div = 2 ** (int(rand_val * 10 / 3) + 2) + img = img * 255 + img = (img // quant_div) * quant_div + img = img / 255 + elif 'color_jitter' in aug: + lo_end = 0 + hi_end = .2 + setting = rand_val * (hi_end - lo_end) + lo_end + img = kornia_color_jitter_numpy(img, setting) + elif 'gaussian_blur' in aug: + img = cv2.GaussianBlur(img, (0, 0), self.blur_scale*rand_val*1.5) + elif 'motion_blur' in aug: + # Motion blur + intensity = self.blur_scale*rand_val * 3 + 1 + angle = random.randint(0, 360) + k = np.zeros((intensity, intensity), dtype=np.float32) + k[(intensity - 1) // 2, :] = np.ones(intensity, dtype=np.float32) + k = cv2.warpAffine(k, cv2.getRotationMatrix2D((intensity / 2 - 0.5, intensity / 2 - 0.5), angle, 1.0), + (intensity, intensity)) + k = k * (1.0 / np.sum(k)) + img = cv2.filter2D(img, -1, k) + elif 'block_noise' in aug: + # Large distortion blocks in part of an img, such as is used to mask out a face. + pass + elif 'lq_resampling' in aug: + # Random mode interpolation HR->LR->HR + if 'lq_resampling4x' == aug: + scale = 4 + else: + if rand_val < .3: + scale = 1 + elif rand_val < .7: + scale = 2 + else: + scale = 4 + if scale > 1: + interpolation_modes = [ + cv2.INTER_NEAREST, cv2.INTER_CUBIC, cv2.INTER_LINEAR, cv2.INTER_LANCZOS4] + mode = random.randint(0, 4) % len(interpolation_modes) + # Downsample first, then upsample using the random mode. + img = cv2.resize(img, dsize=( + img.shape[1]//scale, img.shape[0]//scale), interpolation=mode) + + def lq_resampling_undo_fn(scale, img): + return cv2.resize(img, dsize=(img.shape[1]*scale, img.shape[0]*scale), interpolation=cv2.INTER_LINEAR) + undo_fn = functools.partial(lq_resampling_undo_fn, scale) + elif 'color_shift' in aug: + # Color shift + pass + elif 'interlacing' in aug: + # Interlacing distortion + pass + elif 'chromatic_aberration' in aug: + # Chromatic aberration + pass + elif 'noise' in aug: + # Random noise + if 'noise-5' == aug: + noise_intensity = 5 / 255.0 + else: + noise_intensity = (rand_val*6) / 255.0 + img += np.random.rand(*img.shape) * noise_intensity + elif 'jpeg' in aug: + if 'noise' not in applied_augmentations and 'noise-5' not in applied_augmentations: + if aug == 'jpeg': + lo = 10 + range = 20 + elif aug == 'jpeg-low': + lo = 15 + range = 10 + elif aug == 'jpeg-medium': + lo = 23 + range = 25 + elif aug == 'jpeg-broad': + lo = 15 + range = 60 + elif aug == 'jpeg-normal': + lo = 47 + range = 35 + else: + raise NotImplementedError( + "specified jpeg corruption doesn't exist") + # JPEG compression + qf = (int((1-rand_val)*range) + lo) + # Use PIL to perform a mock compression to a data buffer, then swap back to cv2. + img = (img * 255).astype(np.uint8) + img = Image.fromarray(img) + buffer = BytesIO() + img.save(buffer, "JPEG", quality=qf, optimize=True) + buffer.seek(0) + jpeg_img_bytes = np.asarray( + bytearray(buffer.read()), dtype="uint8") + img = read_img("buffer", jpeg_img_bytes, rgb=True) + elif 'saturation' in aug: + # Lightening / saturation + saturation = rand_val * .3 + img = np.clip(img + saturation, a_max=1, a_min=0) + elif 'greyscale' in aug: + img = np.tile(np.mean(img, axis=2, keepdims=True), [1, 1, 3]) + elif 'none' not in aug: + raise NotImplementedError("Augmentation doesn't exist") + + return img, undo_fn diff --git a/dlas/data/images/image_folder_dataset.py b/dlas/data/images/image_folder_dataset.py new file mode 100644 index 0000000..572d627 --- /dev/null +++ b/dlas/data/images/image_folder_dataset.py @@ -0,0 +1,306 @@ +import functools +import os +import random + +import cv2 +import numpy as np +import torch +import torchvision +from data import util +from torch.utils.data import DataLoader +from torchvision.transforms import Normalize +from tqdm import tqdm + +# Builds a dataset created from a simple folder containing a list of training/test/validation images. +from dlas.data.images.image_corruptor import (ImageCorruptor, + kornia_color_jitter_numpy) +from dlas.data.images.image_label_parser import VsNetImageLabeler +from dlas.utils.util import opt_get + + +def ndarray_center_crop(crop, img): + y, x, c = img.shape + startx = x // 2 - crop // 2 + starty = y // 2 - crop // 2 + return img[starty:starty + crop, startx:startx + crop, :] + + +class ImageFolderDataset: + def __init__(self, opt): + self.opt = opt + self.corruptor = ImageCorruptor(opt) + if 'center_crop_hq_sz' in opt.keys(): + self.center_crop = functools.partial(ndarray_center_crop, opt['center_crop_hq_sz']) + self.target_hq_size = opt['target_size'] if 'target_size' in opt.keys() else None + self.multiple = opt['force_multiple'] if 'force_multiple' in opt.keys() else 1 + self.scale = opt['scale'] + self.paths = opt['paths'] + self.corrupt_before_downsize = opt['corrupt_before_downsize'] if 'corrupt_before_downsize' in opt.keys() else False + self.fetch_alt_image = opt['fetch_alt_image'] # If specified, this dataset will attempt to find a second image + # from the same video source. Search for 'fetch_alt_image' for more info. + self.fetch_alt_tiled_image = opt['fetch_alt_tiled_image'] # If specified, this dataset will attempt to find anoter tile from the same source image + # Search for 'fetch_alt_tiled_image' for more info. + assert not (self.fetch_alt_image and self.fetch_alt_tiled_image) # These are mutually exclusive. + self.skip_lq = opt_get(opt, ['skip_lq'], False) + self.disable_flip = opt_get(opt, ['disable_flip'], False) + self.rgb_n1_to_1 = opt_get(opt, ['rgb_n1_to_1'], False) + self.force_square = opt_get(opt, ['force_square'], True) + self.fixed_parameters = {k: torch.tensor(v) for k, v in opt_get(opt, ['fixed_parameters'], {}).items()} + self.all_image_color_jitter = opt_get(opt, ['all_image_color_jitter'], 0) + if 'normalize' in opt.keys(): + if opt['normalize'] == 'stylegan2_norm': + self.normalize = Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True) + elif opt['normalize'] == 'imagenet': + self.normalize = Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225), inplace=True) + else: + raise Exception('Unsupported normalize') + else: + self.normalize = None + if self.target_hq_size is not None: + assert (self.target_hq_size // self.scale) % self.multiple == 0 # If we dont throw here, we get some really obscure errors. + if not isinstance(self.paths, list): + self.paths = [self.paths] + self.weights = [1] + else: + self.weights = opt['weights'] + + if 'labeler' in opt.keys(): + if opt['labeler']['type'] == 'patch_labels': + self.labeler = VsNetImageLabeler(opt['labeler']['label_file']) + assert len(self.paths) == 1 # Only a single base-path is supported for labeled images. + self.image_paths = self.labeler.get_labeled_paths(self.paths[0]) + else: + self.labeler = None + + # Just scan the given directory for images of standard types. + supported_types = ['jpg', 'jpeg', 'png', 'gif'] + self.image_paths = [] + for path, weight in zip(self.paths, self.weights): + cache_path = os.path.join(path, 'cache.pth') + if os.path.exists(cache_path): + imgs = torch.load(cache_path) + else: + print("Building image folder cache, this can take some time for large datasets..") + imgs = util.find_files_of_type('img', path)[0] + torch.save(imgs, cache_path) + for w in range(weight): + self.image_paths.extend(imgs) + self.len = len(self.image_paths) + + def get_paths(self): + return self.image_paths + + # Given an HQ square of arbitrary size, resizes it to specifications from opt. + def resize_hq(self, imgs_hq): + # Enforce size constraints + h, w, _ = imgs_hq[0].shape + + if self.target_hq_size is not None and self.target_hq_size != h: + hqs_adjusted = [] + for hq in imgs_hq: + # It is assumed that the target size is a square. + target_size = (self.target_hq_size, self.target_hq_size) + hqs_adjusted.append(cv2.resize(hq, target_size, interpolation=cv2.INTER_AREA)) + h, w = self.target_hq_size, self.target_hq_size + else: + hqs_adjusted = imgs_hq + hq_multiple = self.multiple * self.scale # Multiple must apply to LQ image. + if h % hq_multiple != 0 or w % hq_multiple != 0: + hqs_conformed = [] + for hq in hqs_adjusted: + h, w = (h - h % hq_multiple), (w - w % hq_multiple) + hqs_conformed.append(hq[:h, :w, :]) + return hqs_conformed + return hqs_adjusted + + def synthesize_lq(self, hs): + h, w, _ = hs[0].shape + ls = [] + local_scale = self.scale + if self.corrupt_before_downsize: + # You can downsize to a specified scale, then corrupt, then continue the downsize further using this option. + if 'corrupt_before_downsize_factor' in self.opt.keys(): + special_factor = self.opt['corrupt_before_downsize_factor'] + hs = [cv2.resize(h_, (h // special_factor, w // special_factor), interpolation=cv2.INTER_AREA) for h_ in hs] + local_scale = local_scale // special_factor + else: + hs = [h.copy() for h in hs] + hs, ent = self.corruptor.corrupt_images(hs, return_entropy=True) + for hq in hs: + h, w, _ = hq.shape + ls.append(cv2.resize(hq, (h // local_scale, w // local_scale), interpolation=cv2.INTER_AREA)) + # Corrupt the LQ image (only in eval mode) + if not self.corrupt_before_downsize: + ls, ent = self.corruptor.corrupt_images(ls, return_entropy=True) + return ls, ent + + def reset_random(self): + self.corruptor.reset_random() + + def __len__(self): + return self.len + + def __getitem__(self, item): + hq = util.read_img(None, self.image_paths[item], rgb=True) + if hasattr(self, 'center_crop'): + hq = self.center_crop(hq) + if not self.disable_flip and random.random() < .5: + hq = hq[:, ::-1, :] + + if self.force_square: + h, w, _ = hq.shape + dim = min(h, w) + hq = hq[(h - dim) // 2:dim + (h - dim) // 2, (w - dim) // 2:dim + (w - dim) // 2, :] + + # Perform color jittering on the HQ image if specified. The given value should be between [0,1]. + if self.all_image_color_jitter > 0: + hq = kornia_color_jitter_numpy(hq, self.all_image_color_jitter) + + if self.labeler: + assert hq.shape[0] == hq.shape[1] # This just has not been accomodated yet. + dim = hq.shape[0] + + hs = self.resize_hq([hq]) + if not self.skip_lq: + for_lq = [hs[0]] + + # Convert to torch tensor + hq = torch.from_numpy(np.ascontiguousarray(np.transpose(hs[0], (2, 0, 1)))).float() + + out_dict = {'hq': hq, 'LQ_path': self.image_paths[item], 'HQ_path': self.image_paths[item], 'has_alt': False} + + if self.fetch_alt_image: + # This works by assuming a specific filename structure as would produced by ffmpeg. ex: + # 'Candied Walnutsxjktqhr_SYc.webm_00000478.jpg` and + # 'Candied Walnutsxjktqhr_SYc.webm_00000479.jpg` and + # 'Candied Walnutsxjktqhr_SYc.webm_00000480.jpg` + # The basic format is `%08d.`. This logic parses off that 8 digit number. If it is + # not found, the 'alt_image' returned is just the current image. If it is found, the algorithm searches for + # an image one number higher. If it is found - it is returned in the 'alt_hq' and 'alt_lq' keys, else the + # current image is put in those keys. + + imname_parts = self.image_paths[item] + while '.jpg.jpg' in imname_parts: + imname_parts = imname_parts.replace(".jpg.jpg", ".jpg") # Hack workaround to my own bug. + imname_parts = imname_parts.split('.') + if len(imname_parts) >= 2 and len(imname_parts[-2]) > 8: + try: + imnumber = int(imname_parts[-2][-8:]) + # When we're dealing with images in the 1M range, it's straight up faster to attempt to just open + # the file rather than searching the path list. Let the exception handler below do its work. + next_img = self.image_paths[item].replace(str(imnumber), str(imnumber+1)) + alt_hq = util.read_img(None, next_img, rgb=True) + alt_hs = self.resize_hq([alt_hq]) + alt_hq = torch.from_numpy(np.ascontiguousarray(np.transpose(alt_hs[0], (2, 0, 1)))).float() + out_dict['has_alt'] = True + if not self.skip_lq: + for_lq.append(alt_hs[0]) + except: + alt_hq = hq + if not self.skip_lq: + for_lq.append(hs[0]) + else: + alt_hq = hq + if not self.skip_lq: + for_lq.append(hs[0]) + out_dict['alt_hq'] = alt_hq + + if self.fetch_alt_tiled_image: + # This assumes the output format generated by the tiled image generation scripts included with DLAS. Specifically, + # all image read by this dataset are assumed to be in subfolders with other tiles from the same source image. When + # this option is set, another random image from the same folder is selected and returned as the alt image. + sel_path = os.path.dirname(self.image_paths[item]) + other_images = os.listdir(sel_path) + # Assume that the directory contains at least , , + try: + if len(other_images) <= 3: + alt_hq = hq # This is a fallback in case an alt image can't be found. + else: + random.shuffle(other_images) + for oi in other_images: + if oi == os.path.basename(self.image_paths[item]) or 'ref.' in oi or 'centers.pt' in oi: + continue + alt_hq = util.read_img(None, os.path.join(sel_path, oi), rgb=True) + alt_hs = self.resize_hq([alt_hq]) + alt_hq = torch.from_numpy(np.ascontiguousarray(np.transpose(alt_hs[0], (2, 0, 1)))).float() + except: + alt_hq = hq + print(f"Error with {self.image_paths[item]}") + out_dict['has_alt'] = True + out_dict['alt_hq'] = alt_hq + + + if not self.skip_lq: + lqs, ent = self.synthesize_lq(for_lq) + ls = lqs[0] + out_dict['lq'] = torch.from_numpy(np.ascontiguousarray(np.transpose(ls, (2, 0, 1)))).float() + out_dict['corruption_entropy'] = torch.tensor(ent) + if len(lqs) > 1: + alt_lq = lqs[1] + out_dict['alt_lq'] = torch.from_numpy(np.ascontiguousarray(np.transpose(alt_lq, (2, 0, 1)))).float() + + + if self.labeler: + base_file = self.image_paths[item].replace(self.paths[0], "") + while base_file.startswith("\\"): + base_file = base_file[1:] + assert dim % hq.shape[1] == 0 + lbls, lbl_masks, lblstrings = self.labeler.get_labels_as_tensor(hq, base_file, dim // hq.shape[1]) + out_dict['labels'] = lbls + out_dict['labels_mask'] = lbl_masks + out_dict['label_strings'] = lblstrings + + for k, v in out_dict.items(): + if isinstance(v, torch.Tensor) and len(v.shape) == 3: + if self.normalize: + v = self.normalize(v) + if self.rgb_n1_to_1: + v = v * 2 - 1 + out_dict[k] = v + + out_dict.update(self.fixed_parameters) + return out_dict + +if __name__ == '__main__': + opt = { + 'name': 'amalgam', + 'paths': ['E:\\4k6k\\datasets\\ns_images\\imagesets\\256_only_humans_masked'], + 'weights': [1], + 'target_size': 256, + 'scale': 1, + 'corrupt_before_downsize': True, + 'fetch_alt_image': False, + 'fetch_alt_tiled_image': True, + 'disable_flip': True, + 'fixed_corruptions': ['lq_resampling', 'jpeg-medium', 'gaussian_blur', 'noise', 'color_jitter'], + 'num_corrupts_per_image': 0, + 'corruption_blur_scale': 1, + 'all_image_color_jitter': .1, + } + + ds = DataLoader(ImageFolderDataset(opt), shuffle=True, num_workers=0, batch_size=64) + import os + output_path = 'F:\\tmp' + os.makedirs(output_path, exist_ok=True) + res = [] + for i, d in tqdm(enumerate(ds)): + ''' + x = d['hq'] + b,c,h,w = x.shape + x_c = x.view(c*b, h, w) + x_c = torch.view_as_real(torch.fft.rfft(x_c)) + # Log-normalize spectrogram + x_c = (x_c.abs() ** 2).clip(min=1e-8, max=1e16) + x_c = torch.log(x_c) + res.append(x_c) + if i % 100 == 99: + stacked = torch.cat(res, dim=0) + print(stacked.mean(dim=[0,1,2]), stacked.std(dim=[0,1,2])) + ''' + + for k, v in d.items(): + if isinstance(v, torch.Tensor) and len(v.shape) >= 3: + os.makedirs(f'{output_path}\\{k}', exist_ok=True) + torchvision.utils.save_image(v, f'{output_path}\\{k}\\{i}.png') + if i >= 200000: + break diff --git a/dlas/data/images/image_label_parser.py b/dlas/data/images/image_label_parser.py new file mode 100644 index 0000000..7a114fc --- /dev/null +++ b/dlas/data/images/image_label_parser.py @@ -0,0 +1,130 @@ +import os +from collections import OrderedDict + +import orjson as json +# Given a JSON file produced by the VS.net image labeler utility, produces a dict where the keys are image file names +# and the values are a list of object with the following properties: +# [patch_top, patch_left, patch_height, patch_width, label] +import torch + + +class VsNetImageLabeler: + def __init__(self, label_file): + if not isinstance(label_file, list): + label_file = [label_file] + self.labeled_images = {} + for lfil in label_file: + with open(lfil, "r") as read_file: + self.label_file = label_file + # Format of JSON file: + # "key_binding" { + # "label": "