Generating New Watch Designs With StyleGAN

Machine Learning

Category

In early 2019, Nvidia open sourced its “hyperrealistic face generator”, titled StyleGAN. While generating faces demonstrates just how impressive this GAN is, we can also use it to generate really any image we want. That’s considering we have an appropriate dataset of images of course.

I’m not really interested in faces at the moment, so let’s look at some generated watches instead! We will look at what StyleGAN created, then compare it to a GAN I made for fun a couple of months back.

First, the dataset. My goal was 50,000 images of watches, all at 128x128 resolution. I sourced the images from the /r/watches subreddit out of convenience. I previously wrote a script to download a predefined number of images from a given subreddit, so I just reused it here. After the images were downloaded, I manually shifted through the 50,000+ dataset, removing any images that did not fit. I then wrote another script to resize all of the images to 128x128, and voila, I was done.

Now that we have everything in place, It was simply a matter of downloading the source code from GitHub, making the configurations I needed for my situation, then start training. I didn’t need the best quality possible, so I stopped training after a little more than 48 hours. As a quick note, this was all done on a GTX 980 Ti. Unfortunately, this GPU is peanuts compared to what Nvidia recommends for quick training. As a point of reference, Nvidia recommends 8 Tesla V100s if you would like to generate high quality 1024x1024 images in about 6 days. Anyway, here are the results!

While not perfect, they look pretty close to real watches at a quick glance! I could have trained for longer and achieved better results, but I had other projects that required my attention, and I only have one machine. If you want to see what type of results you can get with 14 days of training on a GTX 2080 Ti, check out Christian Freischlag’s article on generating watches!

Also, here is a neat little video on the training process, if you are interested!

Okay, so now we should take a look at what my GAN was able to produce using the same dataset and 5 hours of training.

It’s nowhere near StyleGAN, but I’m still proud of it. It was a fun exercise, I learned a lot, and that’s really all I need!

I hope you enjoyed this quick overview of StyleGAN. If you are interested in doing something like this yourself, I will leave some helpful tips below to help you out!

---

If you would like to do this yourself, you can follow my tutorial on using custom datasets with StyleGaN here!