104 lines
3.9 KiB
Python
104 lines
3.9 KiB
Python
import os
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
from torch.utils.data import Dataset
|
|
|
|
class FruitVegetableDataset(Dataset):
|
|
def __init__(self, images_dir, labels_dir, transform=None):
|
|
"""
|
|
Args:
|
|
images_dir (str): Path to images directory.
|
|
labels_dir (str): Path to labels directory.
|
|
transform (albumentations.Compose, optional): Albumentations transform pipeline.
|
|
"""
|
|
self.images_dir = images_dir
|
|
self.labels_dir = labels_dir
|
|
self.transform = transform
|
|
|
|
# List all image files
|
|
self.image_files = sorted([
|
|
f for f in os.listdir(images_dir)
|
|
if f.lower().endswith(('.jpg', '.jpeg', '.png'))
|
|
])
|
|
|
|
def __len__(self):
|
|
return len(self.image_files)
|
|
|
|
def __getitem__(self, idx):
|
|
img_name = self.image_files[idx]
|
|
img_path = os.path.join(self.images_dir, img_name)
|
|
|
|
# Load image via OpenCV (BGR) and convert to RGB
|
|
image = cv2.imread(img_path)
|
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
height, width, _ = image.shape
|
|
|
|
# Load corresponding label file
|
|
label_name = os.path.splitext(img_name)[0] + ".txt"
|
|
label_path = os.path.join(self.labels_dir, label_name)
|
|
|
|
boxes = []
|
|
class_ids = []
|
|
|
|
if os.path.exists(label_path):
|
|
with open(label_path, "r") as f:
|
|
for line in f:
|
|
parts = line.strip().split()
|
|
if len(parts) == 5:
|
|
class_id = int(parts[0])
|
|
# YOLO format: x_center, y_center, width, height (normalized)
|
|
x_c, y_c, w, h = map(float, parts[1:])
|
|
|
|
# Convert to Pascal VOC absolute coordinates [x_min, y_min, x_max, y_max]
|
|
xmin = (x_c - w / 2.0) * width
|
|
ymin = (y_c - h / 2.0) * height
|
|
xmax = (x_c + w / 2.0) * width
|
|
ymax = (y_c + h / 2.0) * height
|
|
|
|
# Ensure coordinates are within image boundaries
|
|
xmin = max(0.0, min(xmin, float(width - 1)))
|
|
ymin = max(0.0, min(ymin, float(height - 1)))
|
|
xmax = max(xmin + 1.0, min(xmax, float(width)))
|
|
ymax = max(ymin + 1.0, min(ymax, float(height)))
|
|
|
|
boxes.append([xmin, ymin, xmax, ymax])
|
|
class_ids.append(class_id)
|
|
|
|
boxes = np.array(boxes, dtype=np.float32).reshape(-1, 4)
|
|
class_ids = np.array(class_ids, dtype=np.int64)
|
|
|
|
# Apply Albumentations transformations
|
|
if self.transform:
|
|
augmented = self.transform(
|
|
image=image,
|
|
bboxes=boxes,
|
|
category_ids=class_ids
|
|
)
|
|
image = augmented["image"]
|
|
boxes = np.array(augmented["bboxes"], dtype=np.float32).reshape(-1, 4)
|
|
class_ids = np.array(augmented["category_ids"], dtype=np.int64)
|
|
|
|
# Convert image to CHW tensor and scale to [0, 1]
|
|
image_tensor = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0
|
|
|
|
# Convert targets to PyTorch Tensors
|
|
# NOTE: torchvision models reserve class 0 for background.
|
|
# We must shift our class IDs by +1 (0 -> 1, 1 -> 2, etc.)
|
|
torch_boxes = torch.as_tensor(boxes, dtype=torch.float32)
|
|
torch_labels = torch.as_tensor(class_ids + 1, dtype=torch.int64)
|
|
|
|
target = {
|
|
"boxes": torch_boxes,
|
|
"labels": torch_labels,
|
|
"image_id": torch.tensor([idx])
|
|
}
|
|
|
|
return image_tensor, target
|
|
|
|
def collate_fn(batch):
|
|
"""
|
|
Collate function for DataLoader. Returns a tuple of (images, targets).
|
|
"""
|
|
return tuple(zip(*batch))
|