Skip to content

README.md

Latest commit

 

History

History
89 lines (57 loc) · 2.22 KB

README.md

File metadata and controls

89 lines (57 loc) · 2.22 KB
Raw

Trash Can Segmentation

Description

This project focuses on segmenting trash cans in images using deep learning techniques. Two custom models based on the ResUNet architecture were developed:

  • Small ResUNet: A lightweight model with fewer parameters, optimized for faster inference and lower resource consumption.
  • Large ResUNet: A deeper network with increased capacity, aiming for improved accuracy.

Additionally, a pre-trained ResNet model from PyTorch's torchvision.models was fine-tuned to adapt to the trash can segmentation task, serving as a baseline for comparison.

Approaches Used

  • Custom ResUNet Architectures:

    • Small Version: Designed for environments with limited computational resources.
    • Large Version: Enhanced depth and complexity for higher segmentation accuracy.

  • Transfer Learning:

    • Fine-tuning a pre-trained ResNet model to leverage existing feature extraction capabilities for the specific task of trash can segmentation.

Installation

  1. Clone the Repository:

    git clone https://github.com/Micz26/Trash-Can-Segmentation.git
cd Trash-Can-Segmentation

  2. Create a Virtual Environment:

    conda create --name trashcan-env python=3.9

  3. Activate the Virtual Environment:

    conda activate trashcan-env

  4. Install Dependencies:

    pip install .

  5. Download and Place the Dataset:

    • Download the TrashCan dataset.
    • Extract and place the dataset into the data/ directory.

Usage

Training

To train and save the model:

python scripts/train_resunet.py

Docker Deployment

  1. Build the Docker Image:

    docker build -t trashcan-app .

  2. Run the Docker Container:

    docker run -p 8501:8501 trashcan-app

Streamlit Interface (Development Mode)

To run the application using Streamlit:

streamlit run src/trashcan_frontend/frontend.py

Note: If you choose to run the application via Streamlit, ensure that you adjust the paths in the following file:

src/trashcan_core/components/constants/file_paths.py