Scrape2Segment

Table of Contents Introduction Features Installation Usage Project Structure Examples Contributing License Contact

Introduction

Scrape2Segment automates the creation of custom image datasets and trains YOLO models based on user-defined search queries. From web scraping to segmentation, it’s your end-to-end solution for object detection and recognition.

Features

• Image Scraping (scraper.py): Automatically scrape images from the internet based on a keyword using Selenium. The script supports multiple search queries and allows for organizing images into a common directory structure.

• Image Masking (masking.py): Apply masking techniques using YOLOv8 segmentation model to the scraped images. This script handles the preprocessing of images, preparing them for further analysis.

• Latent Representation (training.py): Generate latent representations of the images, which are essential for clustering. This process involves feature extraction using the top layers of VGG16 model to prepare the data for clustering.

• Clustering (clustering.py): Use K-means clustering to group images based on their latent representations. The script identifies the main cluster with the least variability, which is then used to create the final dataset.

• Dataset Splitting and Configuration (yolo_config.py): Split the selected cluster of images into training and validation sets, and prepare the data for YOLO training. This script also configures the dataset directory structure and manages the data flow.

• YOLO Training (final_model.py): Retrain a YOLO model on the generated dataset for segmentation of the given query. The script handles the training loop, model checkpointing, and evaluation metrics.

Installation

1. Clone this repository:

   git clone https://github.com/EleftheriaTtl/Scrape2Segment.git
   cd Scrape2Segment

2. Install the required dependencies:

   pip install -r requirements.txt

3. Install a Selenium WebDriver for Google Chrome:

   • Download ChromeDriver from here.
   • Make sure it is added to your system's PATH.

Usage

Running the Script

To run the script, use the following command:

python main.py --search_queries "plastic cup, red plastic cup, white plastic cup" --search_engine google

• Replace "plastic cup, red plastic cup, white plastic cup" with your desired search queries.

• You can specify multiple search queries separated by commas.

Interactive Prompts

After running the script and if you have inserted multiple queries the script will ask:

"Do you want to use a common directory for all search queries? (yes/no):"

If you answer "yes", you will then be asked to provide a name for the directory:

"Enter the name of the common directory:"

Example: "PLASTIC CUP"

Project Structure

├── downloaded_images/ # The scraped images

├── final_dataset/ # The selected images after the k-means filtering

├── history/ # A pickle file on the scraping image history, to avoid downloading duplicates

├── latent_space_representations/ # The latent space representation

├── masked_images/ # The masks produced by the masking.py

├── scripts/ # The python scripts to perform the various feature of the project

├── webdriver/ # A folder where the latest webdriver is located

├── YAML files/ # The final_dataset images in a YOLO ready format

├── requirements.txt # Python dependencies

└── README.md # Project documentation

Example

Contributing

Contributions are welcome! Please fork this repository, create a new branch for your feature or bugfix, and submit a pull request.

Contact

For any questions or suggestions, feel free to reach out:

Eleftheria Tetoula Tsonga Email: tetoula.tsonga@hotmail.com
LinkedIn: Eleftheria Tetoula Tsonga