I am looking to train a semantic segmentation model with Fully-Convolutional Network model in PyTorch.
I found this semantic segmentation train.py from the dusty-nv repo
This file has been truncated.
# Note -- this training script is tweaked from the original version at:
# It's also meant to be used against this fork of torchvision, which includes
# some patches for exporting to ONNX and adds fcn_resnet18 and fcn_resnet34:
from torch import nn
I went through train.py and found that the train_one_epoch function accepts image and target from the data loader. I was wondering what the format and shape of the image and the target is accepted by the function.
I want to define a dataset class more like PennFudanDataset class from the following link
@harsha.tejas2002, this code is original from torchvision example: vision/references/segmentation at master · pytorch/vision · GitHub
The input is RGB image and the target is a single-channel class ID mask, where each pixel corresponds to the class ID.
Thanks for the reply.
def __init__(self, root, transforms = None):
self.root = root
self.transforms = transforms
self.imgs = list(sorted(os.listdir(os.path.join(root, "PNGImages"))))
self.masks = list(sorted(os.listdir(os.path.join(root, "PedMasks"))))
def __getitem__(self, idx):
img_path = os.path.join(self.root, "PNGImages", self.imgs[idx])
mask_path = os.path.join(self.root, "PedMasks", self.masks[idx])
img = Image.open(img_path).convert("RGB")
target = Image.open(mask_path)
target = np.array(target)
target = torch.as_tensor(target, dtype=torch.int64)
if self.transforms is not None:
img, target = self.transforms(img, target)
return img, target
Please let me know if I should make any changes in the code.
I am not completely sure about the shape of the target.
On the surface it looks okay, but I’m not familiar with this dataset so hard for me to say for sure at a glance. Recommend that you consult the other segmentation dataloaders to compare and aid in your debugging.
This topic was automatically closed 60 days after the last reply. New replies are no longer allowed.