davidtvs / PyTorch-ENet

PyTorch implementation of ENet
MIT License
383 stars 129 forks source link

How to use the 960x720 CamVid dataset? #16

Closed MrLinNing closed 4 years ago

MrLinNing commented 5 years ago

Hello, I want to train the model by big resolution CamVid dataset, and I find the big image can be download here https://github.com/mostafaizz/camvid, but its label image is not consistent with small (480x360) https://github.com/alexgkendall/SegNet-Tutorial/tree/master/CamVid . I am a new segmentation player, I don't know how to convert the color label image to the annotation image. Can you give me some advice about it? Thank you!

davidtvs commented 5 years ago

I would start with:

  1. Make an ordered list of all the colors (by ordered I mean ordered by the annotation label) (colors)
  2. Load each image with Pillow, opencv or scipy and convert it to a numpy array (img)
  3. Create another numpy array with the same height and width as the image (label)
  4. For each color in the list, mask the image and the array
  5. Assign the masked elements to the list index
  6. Save the label to a file

Example code (not tested):

import numpy as np
from PIL import Image

colors = [(255, 0, 0), (0, 255, 0), ...]
img = np.array(Image.open(file))
label = np.zeros(img[:-1])
for anno, c in enumerate(colors):
    label[img == c] = anno

label = Image.fromarray(label).astype(np.uint8))
label.save(save_file)