Проблемы с созданием пользовательского набора данных TensorFlow с использованием аннотаций и изображений COCO ⇐ Python

[Anonymous]

Я загрузил набор данных Coco 2017, и все файлы (аннотации, обучение, проверка и т. д.) доступны. Чтобы создать набор данных для обучения моей модели тензорного потока, я написал класс CocoDataset.

Код: Выделить всё

class CocoDataset:
def __init__(self, annotation_file, image_dir):
self.coco = COCO(annotation_file)
self.image_dir = image_dir
self.img_ids = list(self.coco.imgs.keys())

def load_image_and_annotations(self, image_id):
# Load image metadata
img_data = self.coco.loadImgs(image_id)[0]
img_path = os.path.join(self.image_dir, img_data["file_name"])
#tf.print("Loaded image data:", img_path)

# Read and decode image
img = tf.io.read_file(img_path)
img = tf.image.decode_jpeg(img, channels=3)

# Original image dimensions
original_height = img_data['height']
original_width = img_data['width']

# Resize image to 640x640
img = tf.image.resize(img, [640, 640])

# Calculate resize ratios
x_ratio = 640 / original_width
y_ratio = 640 / original_height

# Get the annotations (e.g., bounding boxes)
annotation_ids = self.coco.getAnnIds(imgIds=image_id)
annotations = self.coco.loadAnns(annotation_ids)

bboxes, masks, categories = [], [], []

for ann in annotations:
# Add bounding boxes
if "bbox" in ann:
xmin, ymin, width, height = ann["bbox"][0]*x_ratio, ann["bbox"][1]*y_ratio, ann["bbox"][2]*x_ratio, ann["bbox"][3]*y_ratio
if width == 0:
width = 1
if height == 0:
height = 1
bbox = [xmin, ymin, xmin + width, ymin + height]
bboxes.append(bbox)

# Add labels (category_id)
categories.append(ann["category_id"])

# Add masks (convert to binary masks using pycocotools)
if "segmentation"  in ann:
mask = self.coco.annToMask(ann)
mask = tf.image.resize(tf.convert_to_tensor(mask, dtype=tf.float32), [640, 640])
masks.append(mask)

# Convert lists to tensors
bboxes = tf.convert_to_tensor(bboxes, dtype=tf.float32)  # Shape: (num_boxes, 4)
masks = tf.convert_to_tensor(masks, dtype=tf.float32) if masks else tf.zeros((0, 640, 640))  # Shape: (num_boxes, 640, 640)
categories = tf.convert_to_tensor(categories, dtype=tf.int32)  # Shape: (num_boxes,)

return img, bboxes, masks, categories

def tf_load_image_and_annotations(self, image_id):
tf.print(image_id) # for debugging
return tf.py_function(self.load_image_and_annotations, [image_id], [tf.float32, tf.float32, tf.float32, tf.int32])

def get_dataset(self, batch_size=4, shuffle_buffer_size=1000):
# Create a tf.data.Dataset
dataset = tf.data.Dataset.from_tensor_slices(self.img_ids)

# Map image IDs to images and annotations
dataset = dataset.map(self.tf_load_image_and_annotations, num_parallel_calls=tf.data.AUTOTUNE)

# Shuffle, batch, and prefetch the data
dataset = dataset.shuffle(shuffle_buffer_size).batch(batch_size).prefetch(tf.data.AUTOTUNE)
return dataset

def get_size(self):
return len(self.img_ids)

Класс вызывается (и создается набор данных) с помощью кода

Код: Выделить всё

coco_dataset = CocoDataset(val_ann_file, val_img_dir)
dataset = coco_dataset.get_dataset()

который загружает аннотации в память и создает индекс.
Когда я хочу просмотреть содержимое первого пакета с помощью следующего кода

Код: Выделить всё

for batch in dataset.take(1):
images, bboxes, masks, categories = batch
print(images.shape)      # Expecting (batch_size, 640, 640, 3)
print(bboxes.shape)      # Expecting (batch_size, num_boxes, 4)
print(masks.shape)       # Expecting (batch_size, num_boxes, 640, 640)
print(categories.shape)   # Expecting (batch_size, num_boxes)

Я вижу результат (идентификатор изображения создан с помощью tf.print(image_id))

Код: Выделить всё

37777
6818
480985
.
.
.
522713
181666
2024-10-20 21:57:00.980492: W tensorflow/core/framework/op_kernel.cc:1827] INVALID_ARGUMENT: TypeError: Cannot iterate over a scalar tensor.
Traceback (most recent call last):

Чего я действительно не понимаю, так это сообщения об ошибке «Невозможно выполнить итерацию по скалярному тензору». У меня такое впечатление, что код в load_image_and_annotations(self, image_id) не выполняется, но, к сожалению, я не могу решить проблему самостоятельно. Буду очень благодарен за любую помощь, предложения и пояснения!

Подробнее здесь: https://stackoverflow.com/questions/791 ... and-images