I’ve recently made a tool to turn my smart lights on and off with an automated people counter, based on the AMG8833, and a raspberry pi zero. For details on how I made the pi talk to the AMG8833, see the previous blog post.
In this post, I’ll outline:
- the general approach I’ve used to build the people counter, using pytorch and python
- how I used onnx to embed my pytorch model into a rust executable.
By the end of the process, I have a single executable that I can run on my pi zero. It gets good FPS - each frame takes about 30-40ms to process, plenty fast enough for me.
And this is the result!
How to make a simple people counter with unsupervised learning
My goal is to make a simple classifier that will classify frames into one of three categories:
- Nothing happening
- Someone entering a zone
- Someone leaving a zone
I opted to use a neural networks approach, because it has a few valuable properties.
- It’s easy to express interesting unsupervised learning ideas as neural networks. I knew I didn’t want to spend time labelling data, so unsupervised learning would definitely be part of my approach.
- Pytorch supports Onnx export, so I knew I’d be able to run my code on the pi through a pi Onnx library.
- If I needed to do any data cleanup, I knew that as long as it was expressed in Pytorch, it would be exported by Onnx.
In the end I used the following architecture, and trained in on a small dataset of me walking back and forth.
- A small autoencoder is trained on single frames from the AM8833
- The encoded frames have dimensionality 3-5
- The encoded frames were run through a K-means classifier, with three clusters. This corresponds to the three states above. Note that this was the sklearn kmeans implementation.
- I extracted the centroids from KMeans, and created a Pytorch version of a KMeans classifier (just the classifier, not the training)
Then, I put all this work end to end to make a single Pytorch module - that I could export with Onnx.
# The autoencoder decoder = nn.Sequential( nn.Linear(64, 32), nn.PReLU(), nn.Linear(32, 4), ) encoder = nn.Sequential( nn.Linear(4, 16), nn.PReLU(), nn.Linear(16, 64), ) net = nn.Sequential( decoder, encoder )
# The KMeans implementation class KMeansLookup(nn.Module): def __init__(self, centroids, net): super().__init__() self.centroids = nn.Parameter(centroids) self.net = net def forward(self, x): cls_x = self.net(x) dists =  for ix, c in enumerate(self.centroids): dists.append(torch.norm(cls_x - c)) dists = torch.stack(dists) return dists.argmin()
# The export using Onnx mdl = KMeansLookup(torch.tensor(cls.cluster_centers_).float(), decoder) torch.onnx.export(mdl, x, 'test.onnx', verbose=True, training=False, example_outputs=r)
For all the code, check out the notebook with all the code
To get a sense for how this “looks”, I’ve extracted a few plots. The first is of the encoded frames from the autoencoder.
The spikes are me walking in and out of frame - you can see why it might be easy for KMeans to sort these into three clusters! In fact, this is what KMeans produces:
So all I need to do now is run this model in my rust executable, and then transform the state transitions into lights changing!
How to compile a model into a rust executable
This was a little bit fiddly - mostly because it takes a bit of work to get data from the internal format my program uses, into the format that Onnx expects. However, with some experimentation I was able to do so, and you can take a look at the code to see how I did it.
The other thing I really liked was that I could compile my model into the executable itself, rather than keeping it as a separate file. The rust
include_bytes macro helped here. This keeps the deployment simple and easy!
The only other piece of significant work was making a small state machine to manage the lights as people move around.
I was quite happy with most of this process. The ML side needs some refinement, I think it still isn’t robust enough. I’m very happy with how Onnx worked out, and I’ll continue to push models into edge devices in this way in future projects.