Spaces:
Running
Running
File size: 2,156 Bytes
4875d48 |
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 59 60 61 62 63 64 65 66 67 68 69 |
import os
from PIL import Image
import torch
from torch.utils.data import Dataset
import numpy as np
class SingleClassSegmentationDataset(Dataset):
def __init__(self, dataset, class_labels, image_size=352, transform=None):
self.items = dataset
self.class_labels = class_labels
self.image_size = image_size
self.transform = transform
def __len__(self):
return len(self.items)
def __getitem__(self, idx):
item = self.items[idx]
image = Image.open(item["img_path"]).convert("RGB")
mask = Image.open(item["mask_path"]).convert("L")
class_name = item["label"]
class_index = self.class_labels.index(class_name)
background_index = 0
mask_np = np.array(mask) > 0
final_mask = np.full(mask_np.shape, background_index, dtype=np.uint8)
final_mask[mask_np] = class_index
image = image.resize((self.image_size, self.image_size), Image.BILINEAR)
final_mask = Image.fromarray(final_mask).resize((self.image_size, self.image_size), Image.NEAREST)
if self.transform:
image, final_mask = self.transform(image, final_mask)
return {
"image": image,
"labels": torch.from_numpy(np.array(final_mask)).long()
}
class SegmentationCollator:
def __init__(self, processor, class_labels):
self.processor = processor
self.class_labels = class_labels
def __call__(self, batch):
images = [item["image"] for item in batch]
labels = [item["labels"] for item in batch]
prompts = self.class_labels * len(images)
expanded_images = [img for img in images for _ in self.class_labels]
inputs = self.processor(
images=expanded_images,
text=prompts,
return_tensors="pt",
padding=True,
truncation=True
)
return {
"pixel_values": inputs["pixel_values"],
"input_ids": inputs["input_ids"],
"labels": torch.stack(labels)
}
|