initial commit

dataset.py

@@ -0,0 +1,103 @@
+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))