StyleGAN2 for medical datasets

StyleGAN2 for medical datasets

In this project, we would train a StyleGAN2 model for medical datasets. This could be beneficial for synthetic data augmentation, or potentially encoding into and studying the latent space could be useful for other medical applications.


Personally, I am more interested in histopathological datasets:

However, I don’t mind looking into other medical datasets and training a StyleGAN2 for them:

Training scripts

This repository has a JAX model implementation of StyleGAN2 but unfortunately not a training script.

Desired project outcome

At the end of the project we would hopefully have:

  • A JAX-based pipeline for training StyleGAN2 models on arbitrary datasets and for exploring the latent space of these models.
  • A trained StyleGAN2 model for the selected datasets (with an online demo)

Let me know if you find this project interesting and if you have any related ideas…


Hey @tmabraham, super interested in being part of this project. I’m working in the intersection on Deep Learning and Medical imaging and have worked on histopathological datasets.

Really looking forward to train StyleGAN2 using JAX and see how the latent space on medical imaging pans out for other downstream tasks.

1 Like

It should definitely be feasible to write a training script during the week! Let’s officially create this project :slight_smile: GANs are usually quite difficult / unstable to train, but let’s give it a shot!


I want to check it out too.

Can I get added to the sheet?

I’m a Master’s student from India.
This is a very interesting project. I have some experience in working with GANs and used them to generate synthetic data (agricultural sector).

I am unfamiliar with JAX. This will a great opportunity for me to learn and contribute. Looking forward to working on this project.

Glad to hear all of you are interested in this project!

@vasudev-sharma @Tahsin-Mayeesha @Ankit-Kumar-Saini I have added the #stylegan-medical channel to the Flax-HuggingFace-Community-Week discord server. Let’s discuss there our next steps over there!

Thank you @patrickvonplaten for finalizing the project!

Hey @ tmabraham ! This project sounds interesting! I would love to join it. Although I have not worked with JAX yet, I have experience with GANs and Computer Vision using PyTorch. It would be a great learning experience for me. Please let me know how can I join?

1 Like

Patrick will be adding people officially to the project, but we will be having discussions in the Discord server under the #stylegan-medical channel.


Adding you @sachdevkartik !


Hey, I have some experience with GANs and would like to join :slight_smile:

Hey @tmabraham can you share the link of the channel on the Discord server?

Can you try this link?:

@Ankit-Kumar-Saini I see you are not on the spreadsheet:

If you are still interested in this project, please try to add yourself to the spreadsheet for the “StyleGAN2 for medical datasets”.

@tmabraham I am unable to add myself to the spreadsheet

@sachdevkartik I don’t see you on the Discord yet…

I’m there with the name Kitadake. I’ll message in the channel.

hi @tmabraham I find this project interesting, can you please add me to the team.
Thank you

Hi @patrickvonplaten @tmabraham I have been working on digital pathology for my masters, and I feel I can contribute to this project. I have experience with GANs too. Please add me to the project sheet if there are still spots left

Are you still interested? Please contact Patrick Von Platen (over here or on Slack) regarding this. Also I’ll need your Discord username, that’s where we are discussing.