generator.py

import os.path
import json
import scipy.misc
import numpy as np
import matplotlib.pyplot as plt
from skimage.transform import resize

# In this exercise task you will implement an image generator. Generator objects in python are defined as having a next function.
# This next function returns the next generated object. In our case it returns the input of a neural network each time it gets called.
# This input consists of a batch of images and its corresponding labels.
class ImageGenerator:
    def __init__(self, file_path, label_path, batch_size, image_size, rotation=False, mirroring=False, shuffle=False):
        # Define all members of your generator class object as global members here.
        self.file_path=file_path
        self.label_path=label_path
        self.batch_size=batch_size
        self.image_size = image_size
        self.rotation = rotation
        self.mirroring = mirroring
        self.shuffle = shuffle

        self.Curr_Index=0  # index count
        self.Curr_Epoch=0 # epoch count


        with open(self.label_path, 'r') as f: #To load JSON data
            self.labels = json.load(f)


        self.image_files = [f for f in os.listdir(self.file_path)] #loading data files

        self.no_images= len(self.image_files) #number of images



        if self.shuffle == True:
            np.random.shuffle(self.image_files) # shuffle image files if true


        self.class_dict = {0: 'airplane', 1: 'automobile', 2: 'bird', 3: 'cat', 4: 'deer', 5: 'dog', 6: 'frog',
                           7: 'horse', 8: 'ship', 9: 'truck'}



        # These need to include:
        # the batch size
        # the image size
        # flags for different augmentations and whether the data should be shuffled for each epoch
         # The labels are stored in json format and can be directly loaded as dictionary.
        # Note that the file names correspond to the dicts of the label dictionary.



        #TODO: implement constructor

    def next(self):
        labels=[]
        images=[]

        if self.Curr_Index >= self.no_images:
            self.Curr_Epoch=self.Curr_Epoch+1 # increasing epoch count after one cycle is done
            self.Curr_Index=0 # start by the first index again for the new epoch
            if self.shuffle == True:
                np.random.shuffle(self.image_files)
        Curr_Batch=self.image_files[self.Curr_Index:self.Curr_Index+self.batch_size]  #making batches

        if len(Curr_Batch)<self.batch_size:
            remaining_batch_size = self.batch_size - len(Curr_Batch)
            Curr_Batch = self.image_files[self.Curr_Index:] + self.image_files[:remaining_batch_size]  #reuse images from the start

        for file in Curr_Batch :
            Image_Path = os.path.join(self.file_path,file) #full file path /as/1.npy

            img = np.load(Image_Path)
            Resized_img= resize(img,(self.image_size[0],self.image_size[1],self.image_size[2]))


            root, ext = os.path.splitext(file) # root ='1' ,ext= 'npy'

            json_label = self.labels[root]


            labels.append(json_label)  # append number

            if self.rotation== True or self.mirroring == True:
                Resized_img = self.augment(Resized_img)
            images.append(Resized_img)

        self.Curr_Index = self.Curr_Index + self.batch_size # updating index
        images=np.array(images)
        labels=np.array(labels)

        # This function creates a batch of images and corresponding labels and returns them.
        # In this context a "batch" of images just means a bunch, say 10 images that are forwarded at once.
        # Note that your amount of total data might not be divisible without remainder with the batch_size.
        # Think about how to handle such cases
        #TODO: implement next method
        #print(labels)
        pass
        return images, labels

    def augment(self,img):
        if self.rotation==True:
            arr = np.array([1, 2, 3,])
            np.random.shuffle(arr)
            img= np.rot90(img,arr[0])

        if self.mirroring== True:
            condition=np.array([1,2])
            np.random.shuffle(condition)
            if condition[0]==1:
                img=np.fliplr(img)

        # this function takes a single image as an input and performs a random transformation
        # (mirroring and/or rotation) on it and outputs the transformed image
        #TODO: implement augmentation function

        return img

    def current_epoch(self):

        # return the current epoch number
        return self.Curr_Epoch

    def class_name(self, x):
        # This function returns the class name for a specific input
        #TODO: implement class name function
        return self.class_dict.get(x)

    def show(self):
        # In order to verify that the generator creates batches as required, this functions calls next to get a
        # batch of images and labels and visualizes it.
        images, labels = self.next()

        columns = 3
        for i, image in enumerate(images):

            rows = (len(images) + columns - 1) // columns # to calculate correct number of rows
            plt.subplot(rows, columns, i + 1)
            plt.imshow(image)
            plt.title(self.class_name(labels[i]))
            plt.axis('off')
        plt.subplots_adjust( hspace=0.6)  # making sure the titles doesnt overlap with the images
        plt.show()

        #TODO: implement show method