Сегментация DataLoader не останавливается - PullRequest
Новая
       57

Сегментация DataLoader не останавливается

0 голосов
/ 07 мая 2020

Я написал загрузчик данных для конвейера сегментации. Загрузчик данных отлично загружает изображения и маски сегментации, но выполнение загрузчика данных никогда не останавливается, так как он продолжает работать вечно.

Мой код прилагается ниже. Пожалуйста, ознакомьтесь с ним и сообщите по этому поводу. 

import random
import cv2
import itertools
import numpy as np
import keras
from keras.preprocessing.image import ImageDataGenerator, array_to_img, img_to_array, load_img


def padImage(image, channels, input_height, input_width, output_height, output_width):
  pad = int((output_height-input_height)/2)
  canvas = np.zeros((output_height,output_width,channels))
  canvas[pad:pad+input_height, pad:pad+input_width] = image
  return canvas

def image_segmentation_generator(img_seg_pairs, batch_size,
                                 input_height, input_width,
                                 output_height, output_width,
                                 classes):

    img_seg_pairs = img_seg_pairs
    random.shuffle(img_seg_pairs)
    zipped = itertools.cycle(img_seg_pairs)
    print(img_seg_pairs.shape)

    while True:
        X = []
        Y = []
        for _ in range(batch_size):
            im, seg = next(zipped)

            #im = cv2.imread(im, cv2.IMREAD_GRAYSCALE)
            im = img_to_array(load_img(im, color_mode='grayscale'))
            im = im.reshape(input_height,input_width,4)
            im = padImage(im, 4, input_height, input_width, output_height, output_width)
            im = im/255.0

            #seg = cv2.imread(seg, cv2.IMREAD_GRAYSCALE)
            seg = img_to_array(load_img(seg, color_mode='grayscale'))
            seg = seg.reshape(input_height,input_width,1)
            seg = padImage(seg, 1, input_height, input_width, output_height, output_width)
            seg = keras.utils.to_categorical(seg,5)

            X.append(im)
            Y.append(seg)

        yield np.array(X), np.array(Y)


import pandas as pd
train = pd.read_csv("./train.csv")
train = train.values
#print(test)
train_gen = image_segmentation_generator(train,8,
                                            240,240,
                                            256,256,5)


count = 0 
for value in train_gen:  
    #print(value[1].shape[0]) 
    count = count + value[1].shape[0]
    #print(count)
print("Total Count is "+ count)
...