tae898/room-classification-cuda repository overview

⁠room-classification

⁠Data

I got data from here⁠. You should apply for permission first. Download it and name it as data.

⁠Model

You can download a pretrained model from here⁠. It's a fine-tuned EfficientNet-b3⁠ using pytorch-lightning⁠

⁠Prerequisites

1. A unix or unix-like x86 machine
2. Python 3.8 or higher.
3. Running in a virtual environment (e.g., conda, virtualenv, etc.) is highly recommended so that you don't mess up with the system Python.
4. pip install -r requirements.txt

⁠Understand data, training, and evaluation

1. Check out explore-data.ipynb⁠ to explore dataset.
2. Run train.py⁠ to train.

   train.py [-h] [--seed SEED] [--batch-size BATCH_SIZE] [--image-size IMAGE_SIZE] [--limit-data LIMIT_DATA] [--num-classes NUM_CLASSES] [--data-dir DATA_DIR] [--efficientnet EFFICIENTNET] [--epochs EPOCHS]
                   [--use-gpu] [--precision PRECISION] [--patience PATIENCE]
   

   Copy
3. Check out evaluation.py⁠ to evaluate and see some samples.

⁠Deploying

⁠Server

This Flask server is implemented in app.py⁠. It receives an image and outputs a probability distribution over the seven room types (i.e., interior, bathroom, bedroom, exterior, living_room, kitchen, and dining_room).

There are two ways to run the server. You can either run it natively in Python or as a docker container. The docker way is recommended.

• The docker⁠ way (recommended)

  1. Install docker.
  2. For GPU,

     docker run -it --rm -p 10005:10005 --gpus all tae898/room-classification-cuda
     

     Copy
  3. For CPU,

     docker run -it --rm -p 10005:10005 tae898/room-classification
     

     Copy

• Running natively in Python (GPUs are supported)

  1. First do Prerequisites⁠
  2. Download model.ckpt⁠ and place it in the root repo directory.

  3. python app.py
     

     Copy

⁠Client

The Python client sends an image to the server. This image is either an image saved in your disk or your webcam stream.

1. First do Prerequisites⁠
2. Install requirements

   pip install -r requirements-client.txt
   

   Copy
3. Run the server⁠

• Run the room-classifier on the image saved in your disk:

  python client.py --mode image --image-path path/to/the/image.jpg
  

  Copy

  This will save the results at path/to/the/image.jpg.json. It'll look something like this:

  "interior": 0.0012067470233887434,
  "bathroom": 0.004016552586108446,
  "bedroom": 0.9945330619812012,
  "exterior": 9.643802059144946e-07,
  "living_room": 0.00022632408945355564,
  "kitchen": 9.742022484715562e-06,
  "dining_room": 6.571232916030567e-06
  

  Copy

• Run the room-classifier on webcam:

  python client.py --mode webcam
  

  Copy

  

  Also check out the video⁠!

⁠TODOs

1. Make it lighter.
2. Better docstring.
3. More love.

⁠Contributing

Contributions are what make the open source community such an amazing place to be learn, inspire, and create. Any contributions you make are greatly appreciated.

1. Fork the Project
2. Create your Feature Branch (git checkout -b feature/AmazingFeature)
3. Run make style && quality in the root repo directory, to ensure code quality.
4. Commit your Changes (git commit -m 'Add some AmazingFeature')
5. Push to the Branch (git push origin feature/AmazingFeature)
6. Open a Pull Request

⁠Cite our work

⁠Authors

• Taewoon Kim⁠