{"id":15636792,"url":"https://github.com/sayakpaul/maxim-tf","last_synced_at":"2025-04-14T03:34:05.481Z","repository":{"id":62400295,"uuid":"543537127","full_name":"sayakpaul/maxim-tf","owner":"sayakpaul","description":"Implementation of MAXIM in TensorFlow.","archived":false,"fork":false,"pushed_at":"2025-04-08T08:18:32.000Z","size":11688,"stargazers_count":138,"open_issues_count":3,"forks_count":18,"subscribers_count":6,"default_branch":"main","last_synced_at":"2025-04-08T09:28:58.124Z","etag":null,"topics":["all-mlp","computer-vision","conv","deblurring","dehazing","denoising","deraining","enhancement","gmlp","keras","tensorflow"],"latest_commit_sha":null,"homepage":"https://huggingface.co/models?pipeline_tag=image-to-image\u0026sort=downloads\u0026search=maxim","language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/sayakpaul.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2022-09-30T10:01:19.000Z","updated_at":"2025-04-08T09:03:00.000Z","dependencies_parsed_at":"2024-10-23T03:52:58.158Z","dependency_job_id":null,"html_url":"https://github.com/sayakpaul/maxim-tf","commit_stats":{"total_commits":30,"total_committers":1,"mean_commits":30.0,"dds":0.0,"last_synced_commit":"7b17c3729ae35970f5693fd21b1f05a8db6e1435"},"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sayakpaul%2Fmaxim-tf","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sayakpaul%2Fmaxim-tf/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sayakpaul%2Fmaxim-tf/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sayakpaul%2Fmaxim-tf/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/sayakpaul","download_url":"https://codeload.github.com/sayakpaul/maxim-tf/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":248815751,"owners_count":21165977,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["all-mlp","computer-vision","conv","deblurring","dehazing","denoising","deraining","enhancement","gmlp","keras","tensorflow"],"created_at":"2024-10-03T11:07:42.073Z","updated_at":"2025-04-14T03:34:05.474Z","avatar_url":"https://github.com/sayakpaul.png","language":"Jupyter Notebook","funding_links":[],"categories":[],"sub_categories":[],"readme":"# MAXIM in TensorFlow\n\n[![HugginFace badge](https://img.shields.io/badge/🤗%20Hugging%20Face-Spaces-yellow.svg)](https://huggingface.co/spaces/sayakpaul/maxim-spaces)\n[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/sayakpaul/maxim-tf/blob/main/notebooks/inference-dynamic-resize.ipynb) [![TensorFlow 2.10](https://img.shields.io/badge/TensorFlow-2.10-FF6F00?logo=tensorflow)](https://github.com/tensorflow/tensorflow/releases/tag/v2.8.0)\n[![Models on TF-Hub](https://img.shields.io/badge/TF--Hub-Models%20on%20TF--Hub-orange)](https://tfhub.dev/sayakpaul/collections/maxim/1)\n[![HugginFace badge](https://img.shields.io/badge/🤗%20Hugging%20Face-Hub-yellow.svg)](https://huggingface.co/models?pipeline_tag=image-to-image\u0026sort=downloads\u0026search=maxim)\n\nImplementation of MAXIM [1] in TensorFlow. This project received the [#TFCommunitySpotlight Award](https://twitter.com/TensorFlow/status/1611469033714470919?s=20\u0026t=a5LMpYXrPg6E0WGudsYezw).\n\nMAXIM introduces a backbone that can tackle image denoising, dehazing, deblurring, deraining, and enhancement.\n\n\u003cdiv align=\"center\"\u003e\n    \u003cimg src=\"images/overview.png\"/\u003e\n    \u003csup\u003eTaken from the MAXIM paper\u003c/sup\u003e\n\u003c/div\u003e\n\nThe weights of different MAXIM variants are in JAX and they're available in [2]. \n\nYou can find all the TensorFlow MAXIM models [here](https://tfhub.dev/sayakpaul/collections/maxim/1) on TensorFlow Hub as\nwell as on [Hugging Face Hub](https://huggingface.co/models?pipeline_tag=image-to-image\u0026sort=downloads\u0026search=maxim).\n\nYou can try out the models on Hugging Face Spaces:\n\n* [Denoising](https://huggingface.co/spaces/sayakpaul/sidd-denoising-maxim)\n* [Low-light enhancement](https://huggingface.co/spaces/sayakpaul/lol-enhancement-maxim)\n* [Image retouching](https://huggingface.co/spaces/sayakpaul/fivek-retouching-maxim)\n* [Dehazing indoors](https://huggingface.co/spaces/sayakpaul/sots-indoor-dehazing-maxim)\n* [Dehazing outdoors](https://huggingface.co/spaces/sayakpaul/sots-outdoor-dehazing-maxim)\n* [Image deraining](https://huggingface.co/spaces/sayakpaul/rain13k-deraining-maxim)\n* [Image deblurring](https://huggingface.co/spaces/sayakpaul/gopro-deblurring-maxim)\n\nIf you prefer Colab Notebooks, then you can check them out [here](https://github.com/sayakpaul/maxim-tf/tree/main/notebooks). \n## Model conversion to TensorFlow from JAX\n\nBlocks and layers related to MAXIM are implemented in the `maxim` directory. \n\n`convert_to_tf.py` script is leveraged to initialize a particular MAXIM model variant and a pre-trained checkpoint and then run the conversion to TensorFlow. Refer to the usage section of the script to know more. \n\nThis script serializes the model weights in `.h5` as as well pushes the `SavedModel` to Hugging Face Hub. For the latter, you need to authenticate yourself if not already done (`huggingface-cli login`).\n\nThis TensorFlow implementation is in close alignment with [2]. The author of this repository has reused some code blocks from [2] (with credits) to do. \n\n## Results and model variants\n\nA comprehensive table is available [here](https://github.com/google-research/maxim#results-and-pre-trained-models). The author of this repository validated the results with the converted models qualitatively. \n\n\u003cdiv align=\"center\"\u003e\n\u003cimg src=\"images/maxim.gif\" width=575\u003e\n\u003c/div\u003e\n\n## Inference with the provided sample images\n\nYou can run the `run_eval.py` script for this purpose.\n\n\u003cdetails\u003e\n  \u003csummary\u003e\u003cstrong\u003eImage Denoising\u003c/strong\u003e (click to expand) \u003c/summary\u003e\n\n```\npython3 maxim/run_eval.py --task Denoising --ckpt_path gs://tfhub-modules/sayakpaul/maxim_s-3_denoising_sidd/1/uncompressed \\\n  --input_dir images/Denoising --output_dir images/Results --has_target=False --dynamic_resize=True\n```\n\u003c/details\u003e\n\n\u003cdetails\u003e\n  \u003csummary\u003e\u003cstrong\u003eImage Deblurring\u003c/strong\u003e (click to expand) \u003c/summary\u003e\n\n```\npython3 maxim/run_eval.py --task Deblurring --ckpt_path gs://tfhub-modules/sayakpaul/maxim_s-3_deblurring_gopro/1/uncompressed \\\n  --input_dir images/Deblurring --output_dir images/Results --has_target=False --dynamic_resize=True\n```\n\u003c/details\u003e\n\n\u003cdetails\u003e\n  \u003csummary\u003e\u003cstrong\u003eImage Deraining\u003c/strong\u003e (click to expand) \u003c/summary\u003e\n\nRain streak:\n```\npython3 maxim/run_eval.py --task Deraining --ckpt_path gs://tfhub-modules/sayakpaul/maxim_s-2_deraining_rain13k/1/uncompressed \\\n  --input_dir images/Deraining --output_dir images/Results --has_target=False --dynamic_resize=True\n```\n\nRain drop:\n```\npython3 maxim/run_eval.py --task Deraining --ckpt_path gs://tfhub-modules/sayakpaul/maxim_s-2_deraining_raindrop/1/uncompressed \\\n  --input_dir images/Deraining --output_dir images/Results --has_target=False --dynamic_resize=True\n```\n\u003c/details\u003e\n\n\u003cdetails\u003e\n  \u003csummary\u003e\u003cstrong\u003eImage Dehazing\u003c/strong\u003e (click to expand) \u003c/summary\u003e\n\nIndoor:\n```\npython3 maxim/run_eval.py --task Dehazing --ckpt_path gs://tfhub-modules/sayakpaul/maxim_s-2_dehazing_sots-indoor/1/uncompressed \\\n  --input_dir images/Dehazing --output_dir images/Results --has_target=False --dynamic_resize=True\n```\n\nOutdoor:\n```\npython3 maxim/run_eval.py --task Dehazing --ckpt_path gs://tfhub-modules/sayakpaul/maxim_s-2_dehazing_sots-outdoor/1/uncompressed \\\n  --input_dir images/Dehazing --output_dir images/Results --has_target=False --dynamic_resize=True\n```\n\u003c/details\u003e\n\n\u003cdetails\u003e\n  \u003csummary\u003e\u003cstrong\u003eImage Enhancement\u003c/strong\u003e (click to expand) \u003c/summary\u003e\n\nLow-light enhancement:\n```\npython3 maxim/run_eval.py --task Enhancement --ckpt_path gs://tfhub-modules/sayakpaul/maxim_s-2_enhancement_lol/1/uncompressed \\\n  --input_dir images/Enhancement --output_dir images/Results --has_target=False --dynamic_resize=True\n```\n\nRetouching:\n```\npython3 maxim/run_eval.py --task Enhancement --ckpt_path gs://tfhub-modules/sayakpaul/maxim_s-2_enhancement_fivek/1/uncompressed \\\n  --input_dir images/Enhancement --output_dir images/Results --has_target=False --dynamic_resize=True\n```\n\u003c/details\u003e\n\u003cbr\u003e\n\n**Notes**:\n\n* The `run_eval.py` script is heavily inspired by the [original one](https://github.com/google-research/maxim/blob/main/maxim/run_eval.py). \n* You can set `dynamic_resize` to False to obtain faster latency compromising the prediction quality.\n\n\n## XLA support\n\nThe models are XLA-supported. It can drammatically reduce the latency. Refer to the `benchmark_xla.py` script for more. \n\n## Known limitations\n\nThese are some of the known limitations of the current implementation. These are all\nopen for contributions.\n\n### Supporting arbitrary image resolutions\n\nMAXIM supports arbitrary image resolutions. However, the available TensorFlow models were exported with `(256, 256, 3)` resolution. So, a crude form of resizing is done on the input images to perform inference with the available models. This impacts the results quite a bit. This issue is discussed in more details [here](https://github.com/sayakpaul/maxim-tf/issues/11). [Some work](https://github.com/sayakpaul/maxim-tf/pull/20) has been started to fix this behaviour (without ETA). I am thankful to [Amy Roberts](https://uk.linkedin.com/in/amy-roberts-70903a6a) from Hugging Face for guiding me in the right direction.\n\nBut these models can be extended to support arbitrary resolution. Refer to [this notebook](https://colab.research.google.com/github/sayakpaul/maxim-tf/blob/main/notebooks/inference-dynamic-resize.ipynb) for more details. Specifically, for a given task and an image, a new version of the model is instantiated and the weights of the available model are copied into the new model instance. This is a time-consuming process and isn't very efficient. \n\n#### Changes to achieve arbitrary image resolution on TF\n\n- Substitute einops calls for pure TF operations: Einops operations are not intended operate on data-dependent (unknown) dimensionality [https://github.com/data-apis/array-api/issues/494](https://github.com/data-apis/array-api/issues/494). Thus, it was necessary to re-write BlockImages and UnblockImages as full TF ops. For convenience, we separate BlockImages into TFBlockImages and TFBlockImagesByGrid. We also rewrote UnblockImages as TFUnblockImages. \n- Make [dim_u](https://github.com/sayakpaul/maxim-tf/pull/24/files#diff-8b281bcfc137b53489e1b19b29735462d5deac19b8c2c2f82cf0383680908063R121) and [dim_v](https://github.com/sayakpaul/maxim-tf/pull/24/files#diff-8b281bcfc137b53489e1b19b29735462d5deac19b8c2c2f82cf0383680908063R130) parameters independent of the input image size. This can be done by computing dim_u and dim_v from the provided grid_size and/or block_size. \n- Change resizing layers so as to receive a ratio independent of the image size. It was important to use the float ratios to compute the final image size, just then converting back to int, to avoid loss of information.\n \n### Output mismatches\n\nThe outputs of the TF and JAX models vary slightly. This is because of the differences in the implementation of different layers (resizing layer mainly). Even though the differences in the outputs of individual blocks of TF and JAX models are small, they add up, in the end, to be larger than one might expect. \n\nWith all that said, the qualitative performance doesn't seem to be disturbed at all.\n\n## Call for contributions\n\n- [ ] Add a minimal training notebook.\n- [ ] Fix any of the known limitations stated above\n\n## Acknowledgements\n\n* ML Developer Programs' team at Google for providing Google Cloud credits.\n* [Gustavo Martins](https://twitter.com/gusthema?lang=en) from Google for initial discussions and reviews of the codebase.\n* [Amy Roberts](https://uk.linkedin.com/in/amy-roberts-70903a6a) from Hugging Face for guiding me in the right direction for handling arbitrary input shapes.\n\n## References\n\n[1] MAXIM paper: https://arxiv.org/abs/2201.02973\n\n[2] MAXIM official GitHub: https://github.com/google-research/maxim\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fsayakpaul%2Fmaxim-tf","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fsayakpaul%2Fmaxim-tf","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fsayakpaul%2Fmaxim-tf/lists"}