SYNERWEEE

Synthetic Data for Enhanced Refurbishment of Waste Electrical and Electronic Equipment

This project addresses the pervasive issue of data scarcity when training machine learning algorithms for the recycling of Waste Electrical and Electronic Equipment (WEEE), specifically laptops.

In industrial settings, obtaining sufficient real-world data for tasks like laptop refurbishment is challenging, time-consuming, and expensive due to the manual effort required for data collection and annotation.

We propose a method to generate synthetic images of laptops with surface defects to enhance the datasets used to train detection models for the identification of damaged laptops and improve their performance.

This project is based on the work presented in the paper Diffusion-based Image Generation for In-distribution Data Augmentation in Surface Defect Detection.

Instructions

1. Set up the repository:

   • $ git clone https://github.com/Clearbox-AI/SYNERWEEE.git

2. Set up the environment:

   • $ conda create -n in_and_out python=3.10
   • $ conda activate in_and_out
   • $ pip install torch==1.12.0+cu113 torchvision==0.13.0+cu113 torchaudio==0.12.0 --extra-index-url https://download.pytorch.org/whl/cu113
   • $ cd sd_utilities/
   • $ pip install --upgrade -r requirements.txt
   • $ pip install xformers==0.0.20
   • $ pip install bitsandbytes==0.38.1
   • $ accelerate config

3. Download the Stable Diffusion model:

   • sd-v1-5-pruned-noema-fp16.safetensors.

   And store it within /sd_utilities/models/.

Fine tuning and image generation with Stable Diffusion

Firstly, enter inside the sd_utilities/ folder with the command:

$ cd sd_utilities/

Make sure the .sh files are executable by running:

$ chmod +x sd_utilities/*.sh

Then:

1. To generate new images from the pretrain of SD, use the script 3_launch_generate_imgs.sh. In particular: Use the parameter --ckpt to specify the path of the SD model, set as default models/sd-v1-5-pruned-noema-fp16.safetensors.
2. To finetune the pretrain of SD:

   • For each image in the dataset, create the corresponding label using the script utils/generate_lbls.py.

     • The token that you specify must have the form sks type_of_item.
     • Specifically, your dataset folder should then have the following structure:

       dataset_folder
           |-- img001.png
           |-- img001.txt
           |-- img002.png
           |-- img002.txt
           |-- img00N.png
           |-- img00N.txt

   • Since we use the Deambooth technique, we need to generate at least 200 regularization images. In order to do this:

     • Generate 200 new images through the script 3_launch_generate_imgs.sh, using the token type_of_item.
     • Generate the corresponding labels using the script utils/generate_lbls.py and the token type_of_item.

   • Put in the same folder your own images (and labels) and the regularization images (and labels).

   • Use the file 1_launch_finetuning.sh to finetune the SD model.

     • Check that all the parameters in configs/dataset_config.toml and configs/training_config.toml are correct.

   • Since we are also using the LoRA technique, use the file 2_launch_merge_lora.sh to merge the weights.

   • Use the file 3_launch_generate_imgs.sh with the prompt sks type_of_item to generate the new images.

     • Use the parameter --ckpt to specify the path of the finetuned SDLoRA model.

The fine-tuning of the model was performed on a machine equipped with two Nvidia T4 GPUs of 16 GB each.

---

Contacts

  ​info@clearbox.ai

  www.clearbox.ai

---

This project is developed in the framework of the euROBIN consortium.

  Funded by the European Union