Spaces:
Sleeping
Sleeping
File size: 2,456 Bytes
86e22bf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
import os
import pandas as pd
import torch
import torchvision.transforms as T
from PIL import Image
from torch.utils.data import Dataset
class PlaneDataset(Dataset):
def __init__(self, images_folder, annotations_folder, transform=None):
self.images_folder = images_folder
self.annotations_folder = annotations_folder
self.transform = transform or T.ToTensor()
self.image_filenames = [f for f in os.listdir(images_folder) if f.endswith(".jpg")]
def __len__(self):
return len(self.image_filenames)
def __getitem__(self, idx):
img_filename = self.image_filenames[idx]
img_path = os.path.join(self.images_folder, img_filename)
# Load and convert image
image = Image.open(img_path).convert("RGB")
image = self.transform(image)
# Read bounding boxes from CSV
annotation_file = os.path.join(self.annotations_folder, img_filename.replace(".jpg", ".csv"))
if not os.path.exists(annotation_file) or os.path.getsize(annotation_file) == 0:
print(f"⚠️ Warning: Annotation file {annotation_file} is missing or empty!")
return image, {"boxes": torch.empty((0, 4), dtype=torch.float32), "labels": torch.empty((0,), dtype=torch.int64)}
try:
bboxes_df = pd.read_csv(annotation_file, header=None, skiprows=1, sep=r"\s+")
# Check if valid bounding boxes exist (at least 4 values per row)
if bboxes_df.shape[1] != 4:
print(f"⚠️ Warning: Invalid bounding boxes in {annotation_file}, skipping...")
return image, {"boxes": torch.empty((0, 4), dtype=torch.float32), "labels": torch.empty((0,), dtype=torch.int64)}
bboxes_df.columns = ["xmin", "ymin", "xmax", "ymax"]
boxes = torch.tensor(bboxes_df[["xmin", "ymin", "xmax", "ymax"]].values, dtype=torch.float32)
labels = torch.ones((boxes.shape[0],), dtype=torch.int64)
except Exception as e:
print(f"❌ Error reading CSV {annotation_file}: {e}")
return image, {"boxes": torch.empty((0, 4), dtype=torch.float32), "labels": torch.empty((0,), dtype=torch.int64)}
target = {"boxes": boxes, "labels": labels}
return image, target
transform = T.Compose([
T.Resize((512, 512)),
T.ToTensor()
])
dataset = PlaneDataset(images_folder="Images", annotations_folder="Annotations", transform=transform)
|