Resolving the "Can't pickle local object" Error in Custom Dataset Training for English Text Detection

[navaneeth-nagu]

Hi team,

CC'ed : @felixdittrich92

I encountered an issue when trying to train a custom dataset for English text detection using Doctr. The error occurs when attempting to start the training process. Below are the details of the command I ran and the error message I received.

Command

python doctr/references/detection/train_pytorch.py ./train ./validation db_resnet50

Error Message

python doctr/references/detection/train_pytorch.py ./train ./validation db_resnet50
Namespace(train_path='./train', val_path='./validation', arch='db_resnet50', name=None, epochs=10, batch_size=2, device=None, save_interval_epoch=False, input_size=1024, lr=0.001, weight_decay=0, workers=None, resume=None, test_only=False, freeze_backbone=False, show_samples=False, wb=False, push_to_hub=False, pretrained=False, rotation=False, eval_straight=False, sched='poly', amp=False, find_lr=False, early_stop=False, early_stop_epochs=5, early_stop_delta=0.01)
Validation set loaded in 0.008291s (5 samples in 3 batches)
WARNING:root:No accessible GPU, targe device set to CPU.
Train set loaded in 0.01626s (9 samples in 4 batches)
                                                                                               Traceback (most recent call last):                                         | 0/4 [00:00<?, ?it/s]
  File "/Users/navaneethn/code/doctr/doctr/references/detection/train_pytorch.py", line 481, in <module>
    main(args)
  File "/Users/navaneethn/code/doctr/doctr/references/detection/train_pytorch.py", line 388, in main
    fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, amp=args.amp)
  File "/Users/navaneethn/code/doctr/doctr/references/detection/train_pytorch.py", line 108, in fit_one_epoch
    pbar = tqdm(train_loader, position=1)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/doc_train/lib/python3.9/site-packages/tqdm/asyncio.py", line 33, in __init__
    self.iterable_iterator = iter(iterable)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/doc_train/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 439, in __iter__
    return self._get_iterator()
  File "/opt/homebrew/Caskroom/miniforge/base/envs/doc_train/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 387, in _get_iterator
    return _MultiProcessingDataLoaderIter(self)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/doc_train/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 1040, in __init__
    w.start()
  File "/opt/homebrew/Caskroom/miniforge/base/envs/doc_train/lib/python3.9/multiprocessing/process.py", line 121, in start
    self._popen = self._Popen(self)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/doc_train/lib/python3.9/multiprocessing/context.py", line 224, in _Popen
    return _default_context.get_context().Process._Popen(process_obj)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/doc_train/lib/python3.9/multiprocessing/context.py", line 284, in _Popen
    return Popen(process_obj)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/doc_train/lib/python3.9/multiprocessing/popen_spawn_posix.py", line 32, in __init__
    super().__init__(process_obj)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/doc_train/lib/python3.9/multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/doc_train/lib/python3.9/multiprocessing/popen_spawn_posix.py", line 47, in _launch
    reduction.dump(process_obj, fp)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/doc_train/lib/python3.9/multiprocessing/reduction.py", line 60, in dump
    ForkingPickler(file, protocol).dump(obj)
AttributeError: Can't pickle local object 'main.<locals>.<lambda>'
  0%|                                                                     | 0/4 [00:00<?, ?it/s]

Thank you for the incredible work you are doing with Doctr. It is a fantastic software, and your efforts are highly appreciated!

Thank you in advance for your support!

[felixdittrich92]

Hi @navaneeth-nagu 👋,
For debugging could you run

python doctr/references/detection/train_pytorch.py ./train ./validation db_resnet50 -j 1

please and check if that works ? :)
Thanks for the kind words 🤗

[navaneeth-nagu]

Context: 

annotations = document.get("label", [])
polygons = []
for result in annotations:
    # Each result is expected to have x, y, width, height, and original dimensions
    x, y, width, height = result["x"], result["y"], result["width"], result["height"]
    corners = [[x, y], [x + width, y], [x + width, y + height], [x, y + height]]
    polygons.append(corners)

img_dimensions = (result["original_width"], result["original_height"])
return {'img_dimensions': img_dimensions, 'img_hash': img_hash, 'polygons': polygons}

Logic to fetch img_hash:
```
def get_image_hash(image_path):
"""Calculate SHA256 hash of an image file."""
try:
    with open(image_path, "rb") as file:
        bytes = file.read()
        readable_hash = hashlib.sha256(bytes).hexdigest()
        return readable_hash
except FileNotFoundError:
    logging.error(f"File not found: {image_path}")
    return "File not found"
```

[navaneeth-nagu]

Yes, each sample has more than 1 annotation.

Attached labels.json file for your reference:
labels.json

btw, thank you for quick turnaround. Much appreciated for your time. :)

[navaneeth-nagu]

Here is the script:

import json
import hashlib
import os
import shutil
import random
from PIL import Image
import logging

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def load_config(path='config.json'):
    """Load configuration from a JSON file."""
    try:
        with open(path, 'r') as file:
            return json.load(file)
    except FileNotFoundError:
        logging.error(f"Configuration file not found: {path}")
        return None

def setup_directories(config):
    """Create directories for storing images and labels."""
    try:
        os.makedirs(f"{config['detection_training_dir']}/images", exist_ok=True)
        os.makedirs(f"{config['detection_validation_dir']}/images", exist_ok=True)
    except Exception as e:
        logging.error(f"Error setting up directories: {e}")

def get_image_hash(image_path):
    """Calculate SHA256 hash of an image file."""
    try:
        with open(image_path, "rb") as file:
            bytes = file.read()
            readable_hash = hashlib.sha256(bytes).hexdigest()
            return readable_hash
    except FileNotFoundError:
        logging.error(f"File not found: {image_path}")
        return "File not found"

def convert_jpg_to_png(source_path, target_path):
    """Convert a JPEG image to a PNG image and save it."""
    try:
        img = Image.open(source_path)
        img.save(target_path, 'PNG')
    except IOError as e:
        logging.error(f"Error converting image from JPG to PNG: {e}")

def process_document(document, config, training_set):
    """Process a single document for image conversion, hash, and annotation creation."""
    image_path = document["image"]
    image_name = image_path.split("/")[-1]
    full_image_path = os.path.join(config['annotation_image_directory'], image_name)
    png_image_name = os.path.splitext(image_name)[0] + '.png'
    target_folder = f"{config['detection_training_dir']}/images" if training_set else f"{config['detection_validation_dir']}/images"
    target_image_path = os.path.join(target_folder, png_image_name)

    # Convert image to PNG and get hash
    convert_jpg_to_png(full_image_path, target_image_path)
    img_hash = get_image_hash(target_image_path)

    # Create polygon annotations from bounding box data
    annotations = document.get("label", [])
    polygons = []
    for result in annotations:
        # Each result is expected to have x, y, width, height, and original dimensions
        x, y, width, height = result["x"], result["y"], result["width"], result["height"]
        corners = [[x, y], [x + width, y], [x + width, y + height], [x, y + height]]
        polygons.append(corners)

    img_dimensions = (result["original_width"], result["original_height"])
    return png_image_name, {'img_dimensions': img_dimensions, 'img_hash': img_hash, 'polygons': polygons}

def load_data_from_file(filename):
    """Load data from a JSON file."""
    try:
        with open(filename, 'r') as file:
            return json.load(file)
    except FileNotFoundError:
        logging.error(f"Data file not found: {filename}")
        return []

def convert_documents(data, config):
    """Convert documents and split into training and validation sets."""
    training_label_output = {}
    validation_label_output = {}
    random.shuffle(data)
    split_point = int(0.7 * len(data))
    training_data = data[:split_point]
    validation_data = data[split_point:]

    for document in training_data:
        image_name, details = process_document(document, config, True)
        training_label_output[image_name] = details

    for document in validation_data:
        image_name, details = process_document(document, config, False)
        validation_label_output[image_name] = details

    return training_label_output, validation_label_output

def save_data_to_file(data, filename):
    """Save processed data to a JSON file."""
    try:
        with open(filename, 'w') as file:
            json.dump(data, file, indent=4)
    except IOError as e:
        logging.error(f"Error saving data to file: {filename} - {e}")

def main():
    config = load_config()
    if config is None:
        return

    setup_directories(config)
    data = load_data_from_file(config['label_studio_json'])
    if not data:
        return

    training_label_output, validation_label_output = convert_documents(data, config)
    save_data_to_file(training_label_output, f"{config['detection_training_dir']}/labels.json")
    save_data_to_file(validation_label_output, f"{config['detection_validation_dir']}/labels.json")

    logging.info("Data processing complete and saved successfully.")

if __name__ == "__main__":
    main()

[felixdittrich92]

labels.json format is correct 👍 but take a look at the polygon values that would be super super small polygons ^^
You should check your x, y, width, height values again (to verify i suggest to plot some samples)
And for doctr format please save the polygon values as integers not floats :)

[navaneeth-nagu]

Ah, got it. Let me check and get back to you. Thanks

[navaneeth-nagu]

Hi @felixdittrich92,

I have a question regarding fine-tuning a custom model after completing its initial training. How would I go about improving this new model as its base?

Thanks in advance. :)

[felixdittrich92]

Hi @navaneeth-nagu 👋,

You mean how you can fine tune a model on your own dataset ?

Take a look here: https://mindee.github.io/doctr/using_doctr/custom_models_training.html

Best regards,
Felix

[navaneeth-nagu]

Yes, I've seen this. If I want to fine-tune a model that has already been fine-tuned before, how can I do it?

[navaneeth-nagu]

Hi @navaneeth-nagu 👋,

You mean how you can fine tune a model on your own dataset ?

Take a look here: https://mindee.github.io/doctr/using_doctr/custom_models_training.html

Best regards,

Felix

Yes, I've seen this. If I want to fine-tune a model that has already been fine-tuned before, how can I do it?

[felixdittrich92]

You can use --pretrained :)

Or use --resume=<PATH TO YOUR MODEL> to resume the training from a saved checkpoint

[navaneeth-nagu]

You can use --pretrained :)

Or use --resume=<PATH TO YOUR MODEL> to resume the training from a saved checkpoint

Ah,okay. Thanks! :)