{"id":17794146,"url":"https://github.com/khdlr/augmax","last_synced_at":"2025-04-13T15:35:34.136Z","repository":{"id":47341750,"uuid":"387887720","full_name":"khdlr/augmax","owner":"khdlr","description":"Efficiently Composable Data Augmentation on the GPU with Jax","archived":false,"fork":false,"pushed_at":"2024-07-18T12:49:34.000Z","size":2016,"stargazers_count":33,"open_issues_count":2,"forks_count":5,"subscribers_count":2,"default_branch":"master","last_synced_at":"2025-04-07T18:21:12.909Z","etag":null,"topics":["gpu","jax","pre-processing"],"latest_commit_sha":null,"homepage":"https://augmax.readthedocs.io/en/latest/","language":"Python","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/khdlr.png","metadata":{"files":{"readme":"README.md","changelog":"CHANGELOG.md","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":"2021-07-20T18:59:30.000Z","updated_at":"2025-02-25T20:20:02.000Z","dependencies_parsed_at":"2024-10-27T11:15:26.951Z","dependency_job_id":"a50c065d-0455-44e4-9a35-8f4172514c7e","html_url":"https://github.com/khdlr/augmax","commit_stats":{"total_commits":68,"total_committers":1,"mean_commits":68.0,"dds":0.0,"last_synced_commit":"318bd314e023e2b3f2c8c6d8d805262cad8b0c06"},"previous_names":[],"tags_count":9,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/khdlr%2Faugmax","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/khdlr%2Faugmax/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/khdlr%2Faugmax/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/khdlr%2Faugmax/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/khdlr","download_url":"https://codeload.github.com/khdlr/augmax/tar.gz/refs/heads/master","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":248737144,"owners_count":21153715,"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":["gpu","jax","pre-processing"],"created_at":"2024-10-27T11:15:16.521Z","updated_at":"2025-04-13T15:35:34.112Z","avatar_url":"https://github.com/khdlr.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Augmax\n[![PyPI version](https://badge.fury.io/py/augmax.svg)](https://pypi.org/project/augmax/) [![Documentation Status](https://readthedocs.org/projects/augmax/badge/?version=latest)](https://augmax.readthedocs.io/en/latest/?badge=latest)\n\nAugmax is an image data augmentation framework supporting efficiently-composable transformations\nwith support for JAX function transformations.\nIts strengths are efficient execution of complex augmentation pipelines and batched data augmentation on the GPU/TPU via the use of [`jax.jit`](jax:jax-jit) and `jax.vmap`.\n\nIn existing data augmentation frameworks,\neach transformation is executed separately,\nleading to many unnecessary memory reads and writes that could be avoided.\nIn contrast, Augmax tries its best to fuse transformations together,\nso that these data-intensive operations are be minimized.\n\n## Getting Started\n\nAugmax aims to implement an API similar to that of [Albumentations](https://albumentations.ai).\nAn augmentation pipeline is defined as a sequence of transformations,\nwhich are then randomly applied to the input images.\n\n```python\nimport jax\nimport augmax\n\ntransform = augmax.Chain(\n  augmax.RandomCrop(256, 256),\n  augmax.HorizontalFlip(),\n  augmax.Rotate(),\n)\n\nimage = ...\n\nrng = jax.random.PRNGKey(27)\n\ntransformed_image = transform(rng, image)\n```\n\n## Batch-wise Augmentation on the GPU\n\nLeveraging the JAX infrastructure,\nit is possible to greatly accelerate data augmentation by using Just-in-Time compilation (`jax.jit`),\nwhich can execute the code on the GPU, as well as batched augmentation (`jax.vmap`).\n\n### Augmenting a single image on the GPU\n```python\ntransformed_image = jax.jit(transform)(rng, image)\n```\n\n### Augmenting an entire batch of images on the GPU\n```python\nsub_rngs = jax.random.split(rng, images.shape[0])\ntransformed_images = jax.jit(jax.vmap(transform))(sub_rngs, images)\n```\n\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fkhdlr%2Faugmax","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fkhdlr%2Faugmax","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fkhdlr%2Faugmax/lists"}