拥有2000+农场结构训练图，如何用TensorFlow Object Detection API自定义模型？

Got it, let's walk through the concrete steps to build your custom object detection model with TensorFlow for those 2000+ farm structure images—way more hands-on than the official docs. Here's a step-by-step breakdown tailored to your use case:

First, get your environment ready with all required tools:

• Install TensorFlow 2.x is recommended for newer APIs and dependencies:

  pip install tensorflow opencv-python pillow lxml matplotlib
  

• Clone the TensorFlow Models repository and set up the Object Detection API:

  git clone https://github.com/tensorflow/models.git
  cd models/research
  # Compile Protobuf files
  protoc object_detection/protos/*.proto --python_out=.
  # Install the object detection package
  cp object_detection/packages/tf2/setup.py .
  python -m pip install .
  

• Verify the installation runs successfully with a test script:

  python object_detection/builders/model_builder_tf2_test.py
  

Your 2000 images need proper annotation and formatting for TensorFlow:

• Annotate images: Use a tool like LabelImg to draw bounding boxes around each farm structure (e.g., barns, fences, irrigation systems) and save annotations in Pascal VOC XML format. Ensure each XML file matches its corresponding image filename.
• Split the dataset: Divide your images and XML files into an 80/20 split (1600 for training, 400 for validation). Create two folders: train with subfolders images and xmls and val with the same structure.
• Generate TFRecord files: TensorFlow uses TFRecords for efficient data loading. Here's a simplified script to convert your XML annotations and images into TFRecords:

  import tensorflow as tf
  from object_detection.utils import dataset_util
  import xml.etree.ElementTree as ET
  import os
  
  # Define your farm structure categories map names to IDs
  category_map = {'barn': 1, 'fence': 2, 'irrigation_system': 3}
  
  def create_tf_example(xml_path, image_path):
      tree = ET.parse(xml_path)
      root = tree.getroot()
  
      filename = root.find('filename').text.encode('utf8')
      width = int(root.find('size/width').text)
      height = int(root.find('size/height').text)
  
      xmins, xmaxs, ymins, ymaxs = [], [], [], []
      classes_text, classes = [], []
  
      for obj in root.findall('object'):
          cls_name = obj.find('name').text
          classes_text.append(cls_name.encode('utf8'))
          classes.append(category_map[cls_name])
  
          bbox = obj.find('bndbox')
          xmins.append(float(bbox.find('xmin').text) / width)
          xmaxs.append(float(bbox.find('xmax').text) / width)
          ymins.append(float(bbox.find('ymin').text) / height)
          ymaxs.append(float(bbox.find('ymax').text) / height)
  
      with tf.io.gfile.GFile(image_path, 'rb') as fid:
          encoded_jpg = fid.read()
  
      tf_example = tf.train.Example(features=tf.train.Features(feature={
          'image/height': dataset_util.int64_feature(height),
          'image/width': dataset_util.int64_feature(width),
          'image/filename': dataset_util.bytes_feature(filename),
          'image/encoded': dataset_util.bytes_feature(encoded_jpg),
          'image/format': dataset_util.bytes_feature(b'jpg'),
          'image/object/bbox/xmin': dataset_util.float_list_feature(xmins),
          'image/object/bbox/xmax': dataset_util.float_list_feature(xmaxs),
          'image/object/bbox/ymin': dataset_util.float_list_feature(ymins),
          'image/object/bbox/ymax': dataset_util.float_list_feature(ymaxs),
          'image/object/class/text': dataset_util.bytes_list_feature(classes_text),
          'image/object/class/label': dataset_util.int64_list_feature(classes),
      }))
      return tf_example
  
  # Generate training TFRecord
  writer = tf.io.TFRecordWriter('train.record')
  for xml_file in os.listdir('train/xmls'):
      img_file = xml_file.replace('.xml', '.jpg')
      tf_example = create_tf_example(f'train/xmls/{xml_file}', f'train/images/{img_file}')
      writer.write(tf_example.SerializeToString())
  writer.close()
  
  # Generate validation TFRecord
  writer = tf.io.TFRecordWriter('val.record')
  for xml_file in os.listdir('val/xmls'):
      img_file = xml_file.replace('.xml', '.jpg')
      tf_example = create_tf_example(f'val/xmls/{xml_file}', f'val/images/{img_file}')
      writer.write(tf_example.SerializeToString())
  writer.close()
  

• Create a label map file: Save this as label_map.pbtxt matches your category IDs:

  item {
    id: 1
    name: 'barn'
  }
  item {
    id: 2
    name: 'fence'
  }
  item {
    id: 3
    name: 'irrigation_system'
  }
  

Choose a pre-trained model based on your speed/accuracy needs:

• Fast inference for edge devices: SSD MobileNet V2
• Higher accuracy: Faster R-CNN ResNet50
  Download the model's config file and checkpoint, then modify the config to match your dataset:
  1. Update num_classes to the number of farm structure categories you have.
  2. In train_input_reader, set input_path to train.record and label_map_path to your label_map.pbtxt.
  3. In eval_input_reader, do the same for val.record.
  4. Adjust batch_size based on your GPU memory start with 8 if you have 8GB VRAM and num_steps try 20000 as a starting point.

Start training with this command:

python model_main_tf2.py \
  --model_dir=./farm_model_weights \
  --pipeline_config_path=./pipeline.config

• Use TensorBoard to monitor training progress loss, mAP:

  tensorboard --logdir=./farm_model_weights
  

• Stop training once the validation mAP stops improving or reaches your target accuracy.

Convert the trained checkpoint to a saved model for detecting farm structures in new images:

python exporter_main_v2.py \
  --input_type image_tensor \
  --pipeline_config_path ./pipeline.config \
  --trained_checkpoint_dir ./farm_model_weights \
  --output_directory ./exported_farm_model

Write a simple inference script to test your model:

import tensorflow as tf
from object_detection.utils import label_map_util
from object_detection.utils import visualization_utils as vis_util
import cv2

# Load the trained model
detect_fn = tf.saved_model.load('./exported_farm_model/saved_model')

# Load label map
category_index = label_map_util.create_category_index_from_labelmap('./label_map.pbtxt', use_display_name=True)

# Load test image
image = cv2.imread('test_farm_image.jpg')
input_tensor = tf.convert_to_tensor(image)
input_tensor = input_tensor[tf.newaxis, ...]

# Run detection
detections = detect_fn(input_tensor)

# Visualize results
vis_util.visualize_boxes_and_labels_on_image_array(
    image,
    detections['detection_boxes'][0].numpy(),
    detections['detection_classes'][0].numpy().astype(int),
    detections['detection_scores'][0].numpy(),
    category_index,
    use_normalized_coordinates=True,
    min_score_thresh=0.5,  # Adjust based on desired precision
    agnostic_mode=False)

# Save or display the result
cv2.imwrite('detection_result.jpg', image)
cv2.imshow('Farm Structure Detection', image)
cv2.waitKey(0)

Pro tip: If you notice poor performance, try augmenting your dataset flip images, adjust brightness, add noise to improve model generalization—this helps a lot with farm scenes that have varying lighting and weather conditions.

内容的提问来源于stack exchange，提问作者Ajinkya