Upload 74 files

.gitattributes

@@ -32,3 +32,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+pretrained_models/angle_model.hdf5 filter=lfs diff=lfs merge=lfs -text
+pretrained_models/length_model.hdf5 filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,34 @@
----
-title: Deep Blind Motion Deblurring
-emoji: 🐠
-colorFrom: blue
-colorTo: green
-sdk: gradio
-sdk_version: 3.16.2
-app_file: app.py
-pinned: false
-license: apache-2.0
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Blind Motion Deblurring for Legible License Plates using Deep Learning
+
+This project uses deep learning techniques to estimate a length and angle parameter for the point-spread function responsible for motion-deblurring of an image. This estimation is achieved by training a deep CNN model on the fast-fourier transformation of the blurred images. By using enough random examples of motion blurred images, the model learns how to estimate any kind of motion blur (upto a certain blur degree), making this approach a truly blind motion deblurring example. Once a length and angle of motion blur is estimated by the model, one can easily deblur the image using Wiener Deconvolution. This technique can have many applications, but we used it specifically for deblurring and making license plates legible. As seen below, the images demonstrate our model in action. With the introduction of some artifacts, the model manages to deblur the images to a point where the license plates are legible.  
+
+<img src="readme_imgs/img1.jpg" width="360px"> <img src="readme_imgs/img1_result.jpg" width="360px">
+
+<img src="readme_imgs/img2.jpg" width="360px"> <img src="readme_imgs/img2_result.jpg" width="360px">
+
+<img src="readme_imgs/img3.jpg" width="360px"> <img src="readme_imgs/img3_result.jpg" width="360px">
+
+## Package Requirements:-
+1. Python3
+2. Numpy
+3. OpenCV 4
+4. Tensorflow 2
+5. H5py
+6. Imutils
+7. Progressbar
+8. Scikit-Learn
+## How to Run Code:-
+
+### Training the length and angle models:-
+
+1. Download the dataset of images from [here](https://cocodataset.org/#download). Download atleast 20000 images to train models optimally. (We used the COCO dataset to train our model. But any other dataset of general images will also suffice)
+2. Use the create_blurred.py to generate the motion blurred dataset as ```python create_blurred.py -i <path_to_input_dir> -o <path_to_output_dir> [-m <optional_number_of_images_to_generate>```. The output directory to store images must exist. The script randomly blurs the images using a random blur length and angle. The range of blur length and angle can be changes on lines 38-39. The script also generates a json file to store the labels for blur length and angle. Note that for blur angle we consider all angle over 180 degrees to be cyclic and wrap around (example 240 is 240-180=60) as it doesn't affect the PSF and significantly reduces the number of classes.
+3. Use the create_fft.py to generate the fast-fourier transform images of the blurred images to use for training. Run the script as ```python create_fft.py -i <path_to_input_dir> -o <path_to_output_dir>```. The input directory is the folder where the blurred images are stored. The output directory must be created manually.
+4. Use the build_dataset.py to generate the hdf5 dataset to train. We use this to overcome the bottleneck of working with a large number of images in memory. Run the script as ```python build_dataset.py -m <flag to determine which model is being trained: use either "angle" or "length"> -i <path to input fft images> -to <output hdf5 train file name/path. Must end with .hd5f extension> -vo <path/filename to output hd5f val data> -l <path to input labels json file. properly input either angle or length labels>```. We have resized our images to (224x224) to facilitate training. If you plan to use a different size change the lines 51 and 64. Before this script is run, make sure to delete any previously present .hdf5 files.
+5. Use the angle_model_train.py script to train the model to estimate the angle parameter of the blur. Change the path to the train and val hdf5 files on lines 17 and 18 and run the script as ```python angle_model_train -o <path to store output metrics. Must be created and empyt at the start of training> [-m <model checkpoint path to resume training> [-e <current epoch to restart training from>```. 
+6. Similarly, length_model_train.py can be used to train the length model.
+7. Remember to properly modify all variables in the train files
+
+### Testing the models to deblur images
+
+Run the deblur_img.py script as ```python deblur_img.py -i <path to input blur image> -a <path to trained angle model> -l <path to trained length model>```. The final deblurred image is saved as result.jpg on the same directory as the script.

angle_model_train.py

@@ -0,0 +1,55 @@
+from sidekick.nn.conv.angle_model import MiniVgg
+from sidekick.io.hdf5datagen import Hdf5DataGen
+from sidekick.callbs.manualcheckpoint import ManualCheckpoint
+from sidekick.callbs.trainmonitor import TrainMonitor
+from sidekick.prepro.process import Process
+from sidekick.prepro.imgtoarrayprepro import ImgtoArrPrePro
+from tensorflow.keras.optimizers import SGD
+from tensorflow.keras.models import load_model
+import argparse
+
+ap= argparse.ArgumentParser()
+ap.add_argument('-o','--output', type=str, required=True ,help="Path to output directory to store metrics")
+ap.add_argument('-m', '--model', help='Path to checkpointed model')
+ap.add_argument('-e','--epoch', type=int, default=0, help="Starting epoch of training")
+args= vars(ap.parse_args())
+
+hdf5_train_path= "train.hdf5"
+hdf5_val_path= "val.hdf5"
+epochs= 50
+lr= 1e-2
+batch_size= 32
+num_classes= 180
+fig_path= args['output']+"train_plot.jpg"
+json_path= args['output']+"train_values.json"
+
+print('[NOTE]:- Building Dataset...\n')
+pro= Process(224, 224)
+i2a= ImgtoArrPrePro()
+
+train_gen= Hdf5DataGen(hdf5_train_path, batch_size, num_classes, preprocessors=[pro, i2a])
+val_gen= Hdf5DataGen(hdf5_val_path, batch_size, num_classes, preprocessors=[pro, i2a])
+
+
+if args['model'] is None:
+    print("[NOTE]:- Building model from scratch...")
+    model= MiniVgg.build(224, 224, 1, num_classes)
+    opt= SGD(learning_rate=lr, momentum=0.9, nesterov=True)
+    model.compile(loss="categorical_crossentropy", metrics=['accuracy'], optimizer=opt)
+else:
+    print("[NOTE]:- Building model {}\n".format(args['model']))
+    model= load_model(args['model'])
+
+callbacks= [ManualCheckpoint(args['output'], save_at=1, start_from=args['epoch']),
+            TrainMonitor(figPath= fig_path, jsonPath= json_path, startAt=args['epoch'])]
+
+print("[NOTE]:- Training model...\n")
+
+model.fit_generator(train_gen.generator(),
+                    steps_per_epoch=train_gen.data_length//batch_size,
+                    validation_data= val_gen.generator(),
+                    validation_steps= val_gen.data_length//batch_size,
+                    epochs=epochs,
+                    max_queue_size=10,
+                    callbacks= callbacks,
+                    initial_epoch=args['epoch'])

build_dataset.py

@@ -0,0 +1,73 @@
+import numpy as np
+from sklearn.preprocessing import LabelEncoder, LabelBinarizer
+from sklearn.model_selection import train_test_split
+from sidekick.io.hdf5_writer import Hdf5Writer
+from imutils import paths
+import cv2
+import os
+import progressbar
+import json
+import argparse
+
+ap= argparse.ArgumentParser()
+ap.add_argument('--model_training', '-m', required=True, help='Flag to determine which model is trained. Choose from "angle" and "length".')
+ap.add_argument('--input_dir', '-i', required=True, help='Path to input dir for images')
+ap.add_argument('--train_output_file', '-to', required=True, help='Path to train output file. Must not exist by default.')
+ap.add_argument('--val_output_file', '-vo', required=True, help='Path to val output file. Must not exist by default.')
+ap.add_argument('--label_file', '-l', required=True, help='Path to input training labels.')
+
+args= vars(ap.parse_args())
+
+model_flag= args['model_training']
+data_path= args['input_dir']
+hdf5_train= args['train_output_file']
+hdf5_test= args['val_output_file']
+label_file= args['label_file']
+
+class_to_use= []
+f= open(label_file, 'r')
+label_dict= json.loads(f.read())
+
+
+train_paths= list(paths.list_images(data_path))
+train_labels= [label_dict[t.split(os.path.sep)[-1]] for t in train_paths]
+
+if model_flag=='angle':
+    le= LabelEncoder()
+    train_labels= le.fit_transform(train_labels)
+    print(le.classes_)
+    print("Number of classes are: {}".format(len(le.classes_)))
+
+train_paths, test_paths, train_labels, test_labels= train_test_split(train_paths,train_labels,
+                                               test_size=0.2)
+
+print(train_paths[10], train_labels[10], test_paths[10], test_labels[10])
+
+files= [('train', train_paths, train_labels, hdf5_train),
+        ('val', test_paths, test_labels, hdf5_test)]
+
+for optype, paths, labels, output_path in files:
+
+    dat_writer= Hdf5Writer((len(paths), 224, 224), output_path)
+
+    # Initializing the progress bar display
+    display=["Building Dataset: ", progressbar.Percentage(), " ",
+             progressbar.Bar(), " ", progressbar.ETA()]
+
+    # Start the progress bar
+    progress= progressbar.ProgressBar(maxval=len(paths), widgets=display).start()
+
+    # Iterate through each img path
+    for (i, (p, l)) in enumerate(zip(paths,labels)):
+        img= cv2.imread(p)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        img = cv2.resize(img, (224, 224))
+        img= img.astype('float') / 255.0
+
+
+        dat_writer.add([img], [l])
+        progress.update(i)
+
+    # Finish the progress for one type
+    progress.finish()
+    dat_writer.close()

create_blurred.py

@@ -0,0 +1,59 @@
+import numpy as np
+import os
+import cv2
+import random
+import json
+import argparse
+
+ap= argparse.ArgumentParser()
+ap.add_argument('--input_dir', '-i', required=True, help='Path to input dir for images')
+ap.add_argument('--output_dir', '-o', required=True, help='Path to output dir to store files. Must be created')
+ap.add_argument('--max_imgs', '-m', default=20000, type=int, help='Max number of images to generate')
+
+args= vars(ap.parse_args())
+
+def apply_motion_blur(image, size, angle):
+	k = np.zeros((size, size), dtype=np.float32)
+	k[ (size-1)// 2 , :] = np.ones(size, dtype=np.float32)
+	k = cv2.warpAffine(k, cv2.getRotationMatrix2D( (size / 2 -0.5 , size / 2 -0.5 ) , angle, 1.0), (size, size) )  
+	k = k * ( 1.0 / np.sum(k) )        
+	return cv2.filter2D(image, -1, k) 
+
+
+folder = args['input_dir']
+folder_save = args['output_dir']
+max_images = args['max_imgs']
+
+print(max_images)
+
+
+labels_angle = {}
+labels_length= {}
+images_done = 0
+for filename in os.listdir(folder):
+	img = cv2.imread(os.path.join(folder,filename))
+	if img is not None and img.shape[1] > img.shape[0]:
+		img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+		img_resized = cv2.resize(img_gray, (640,480), interpolation = cv2.INTER_AREA)
+		length = random.randint(20,40)
+		angle = random.randint(0,359)
+		blurred = apply_motion_blur(img_resized, length, angle)
+		cv2.imwrite(os.path.join(folder_save,filename), blurred)
+		if angle>=180:
+			angle_a= angle - 180
+		else:
+			angle_a= angle
+		labels_angle[filename] = angle_a
+		labels_length[filename]= length
+		images_done += 1
+		print("%s done"%images_done)
+		if(images_done == max_images):
+			print('Done!!!')
+			break
+
+with open('angle_labels.json', 'w') as file:
+	json.dump(labels_angle, file)
+with open('length_labels.json', 'w') as file:
+	json.dump(labels_length, file)
+    
+

create_fft.py

@@ -0,0 +1,31 @@
+import numpy as np
+import os
+import cv2
+import argparse
+
+ap = argparse.ArgumentParser()
+ap.add_argument('--input_dir', '-i', required=True, help='Path to input dir for images')
+ap.add_argument('--output_dir', '-o', required=True, help='Path to output dir to store files. Must be created')
+
+args= vars(ap.parse_args())
+
+
+folder = args['input_dir']
+folder_save = args['output_dir']
+
+
+labels = {}
+images_done = 0
+for filename in os.listdir(folder):
+	img = cv2.imread(os.path.join(folder,filename))
+	if img is not None:
+		img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+		img_gray = np.float32(img_gray) / 255.0
+		f = np.fft.fft2(img_gray)
+		fshift = np.fft.fftshift(f)
+		mag_spec = 20 * np.log(np.abs(fshift))
+		mag_spec = np.asarray(mag_spec, dtype=np.uint8)
+		cv2.imwrite(os.path.join(folder_save,filename), mag_spec)
+		images_done += 1
+		print("%s done"%images_done)
+

deblur_img.py

@@ -0,0 +1,85 @@
+import cv2
+import numpy as np
+from tensorflow.keras.models import load_model
+from tensorflow.keras.preprocessing.image import img_to_array
+import argparse
+
+ap= argparse.ArgumentParser()
+ap.add_argument('--image', '-i', required=True, help='Path to input blurred image')
+ap.add_argument('--angle_model', '-a', required=True, help='Path to trained angle model')
+ap.add_argument('--length_model', '-l', required=True, help='Path to trained length model')
+args= vars(ap.parse_args())
+
+def process(ip_image, length, deblur_angle):
+    noise = 0.01
+    size = 200
+    length= int(length)
+    angle = (deblur_angle*np.pi) /180
+
+    psf = np.ones((1, length), np.float32) #base image for psf
+    costerm, sinterm = np.cos(angle), np.sin(angle)
+    Ang = np.float32([[-costerm, sinterm, 0], [sinterm, costerm, 0]])
+    size2 = size // 2
+    Ang[:,2] = (size2, size2) - np.dot(Ang[:,:2], ((length-1)*0.5, 0))
+    psf = cv2.warpAffine(psf, Ang, (size, size), flags=cv2.INTER_CUBIC) #Warp affine to get the desired psf
+#     cv2.imshow("PSF",psf)
+#     cv2.waitKey(0)
+#     cv2.destroyAllWindows()
+
+    gray = ip_image
+    gray = np.float32(gray) / 255.0
+    gray_dft = cv2.dft(gray, flags=cv2.DFT_COMPLEX_OUTPUT) #DFT of the image
+    psf /= psf.sum() #Dividing by the sum
+    psf_mat = np.zeros_like(gray)
+    psf_mat[:size, :size] = psf
+    psf_dft = cv2.dft(psf_mat, flags=cv2.DFT_COMPLEX_OUTPUT) #DFT of the psf
+    PSFsq = (psf_dft**2).sum(-1)
+    imgPSF = psf_dft / (PSFsq + noise)[...,np.newaxis] #H in the equation for wiener deconvolution
+    gray_op = cv2.mulSpectrums(gray_dft, imgPSF, 0)
+    gray_res = cv2.idft(gray_op,flags = cv2.DFT_SCALE | cv2.DFT_REAL_OUTPUT) #Inverse DFT
+    gray_res = np.roll(gray_res, -size//2,0)
+    gray_res = np.roll(gray_res, -size//2,1)
+
+    return gray_res
+
+
+# Function to visualize the Fast Fourier Transform of the blurred images.
+def create_fft(img):
+    img = np.float32(img) / 255.0
+    f = np.fft.fft2(img)
+    fshift = np.fft.fftshift(f)
+    mag_spec = 20 * np.log(np.abs(fshift))
+    mag_spec = np.asarray(mag_spec, dtype=np.uint8)
+
+    return mag_spec
+
+# Change this variable with the name of the trained models.
+angle_model_name= args['angle_model']
+length_model_name= args['length_model']
+model1= load_model(angle_model_name)
+model2= load_model(length_model_name)
+
+# read blurred image
+ip_image = cv2.imread(args['image'])
+ip_image=  cv2.cvtColor(ip_image, cv2.COLOR_BGR2GRAY)
+ip_image= cv2.resize(ip_image, (640, 480))
+# FFT visualization of the blurred image
+fft_img= create_fft(ip_image)
+
+# Predicting the psf parameters of length and angle.
+img= cv2.resize(create_fft(ip_image), (224,224))
+img= np.expand_dims(img_to_array(img), axis=0)/ 255.0
+preds= model1.predict(img)
+# angle_value= np.sum(np.multiply(np.arange(0, 180), preds[0]))
+angle_value = np.mean(np.argsort(preds[0])[-3:])
+
+print("Predicted Blur Angle: ", angle_value)
+length_value= model2.predict(img)[0][0]
+print("Predicted Blur Length: ",length_value)
+
+op_image = process(ip_image, length_value, angle_value)
+op_image = (op_image*255).astype(np.uint8)
+op_image = (255/(np.max(op_image)-np.min(op_image))) * (op_image-np.min(op_image))
+
+cv2.imwrite("result.jpg", op_image)
+

length_model_train.py

@@ -0,0 +1,52 @@
+from sidekick.nn.conv.length_model import MiniVgg
+from sidekick.io.hdf5datagen import Hdf5DataGen
+from sidekick.callbs.manualcheckpoint import ManualCheckpoint
+from tensorflow.keras.models import load_model
+from sidekick.prepro.process import Process
+from sidekick.prepro.imgtoarrayprepro import ImgtoArrPrePro
+from tensorflow.keras.optimizers import SGD
+import argparse
+
+ap= argparse.ArgumentParser()
+ap.add_argument('-o','--output', type=str, required=True ,help="Path to output directory")
+ap.add_argument('-m', '--model', help='Path to checkpointed model')
+ap.add_argument('-e','--epoch', type=int, default=0, help="Starting epoch of training")
+args= vars(ap.parse_args())
+
+hdf5_train_path= "train.hdf5"
+hdf5_val_path= "val.hdf5"
+epochs= 50
+lr= 1e-2
+batch_size= 32
+num_classes= 1
+fig_path= args['output']+"train_plot.jpg"
+json_path= args['output']+"train_values.json"
+
+print('[NOTE]:- Building Dataset...\n')
+pro= Process(224, 224)
+i2a= ImgtoArrPrePro()
+
+train_gen= Hdf5DataGen(hdf5_train_path, batch_size, num_classes, encode=False, preprocessors=[pro, i2a])
+val_gen= Hdf5DataGen(hdf5_val_path, batch_size, num_classes, encode=False, preprocessors=[pro, i2a])
+
+if args['model'] is None:
+    print("[NOTE]:- Building model from scratch...")
+    model= MiniVgg.build(224, 224, 1, num_classes)
+    opt= SGD(learning_rate=lr, momentum=0.9, nesterov=True)
+    model.compile(loss="mean_absolute_percentage_error", optimizer=opt)
+else:
+    print("[NOTE]:- Building model {}\n".format(args['model']))
+    model= load_model(args['model'])
+
+callbacks= [ManualCheckpoint(args['output'], save_at=1, start_from=args['epoch'])]
+
+print("[NOTE]:- Training model...\n")
+
+model.fit_generator(train_gen.generator(),
+                    steps_per_epoch=train_gen.data_length//batch_size,
+                    validation_data= val_gen.generator(),
+                    validation_steps= val_gen.data_length//batch_size,
+                    epochs=epochs,
+                    max_queue_size=10,
+                    callbacks=callbacks,
+                    initial_epoch=args['epoch'])

pretrained_models/angle_model.hdf5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3fbe443169ad14a6799704313d16f9066ad8643be0b90806df89311d916c0746
+size 49559272

pretrained_models/length_model.hdf5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d6fd68204c720d0ed2643d8ef8f6dc68c677307c9c1a877edef427c5f466ac7
+size 48092168

sidekick/callbs/__init__.py

@@ -0,0 +1,2 @@
+from .trainmonitor import TrainMonitor
+from .manualcheckpoint import ManualCheckpoint

sidekick/callbs/manualcheckpoint.py

@@ -0,0 +1,19 @@
+from tensorflow.keras.callbacks import Callback
+import os
+
+class ManualCheckpoint(Callback):
+    def __init__(self, output, save_at=3, start_from=0):
+        super(Callback, self).__init__()
+
+        self.output= output
+        self.save_at= save_at
+        self.initial_epoch= start_from
+
+    def on_epoch_end(self, epoch, logs={}):
+        if (self.initial_epoch+1) % self.save_at==0:
+            save_path= os.path.sep.join([self.output,
+                                         "weights-epoch {}.hdf5".format(self.initial_epoch+1)])
+
+            self.model.save(save_path, overwrite=True)
+
+        self.initial_epoch+=1

sidekick/callbs/trainmonitor.py

@@ -0,0 +1,52 @@
+from tensorflow.keras import callbacks
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import numpy as np
+import json
+import os
+
+class TrainMonitor(callbacks.BaseLogger):
+    def __init__(self, figPath, jsonPath=None, startAt=0):
+        super(TrainMonitor, self).__init__()
+
+        self.figPath= figPath
+        self.jsonPath= jsonPath
+        self.startAt= startAt
+
+    def on_train_begin(self, logs={}):
+        self.H={}
+
+        if self.jsonPath is not None:
+            if os.path.exists(self.jsonPath):
+                self.H = json.loads(open(self.jsonPath).read())
+
+                if self.startAt > 0:
+                    for k in self.H.keys():
+                        self.H[k] = self.H[k][:self.startAt]
+
+    def on_epoch_end(self, epoch, logs={}):
+        for keys, values in logs.items():
+            l= self.H.get(keys, [])
+            l.append(float(values))
+            self.H[keys] = l
+
+        if self.jsonPath is not None:
+            with open(self.jsonPath, 'w') as f:
+                f.write(json.dumps(self.H))
+                f.close()
+
+        if len(self.H["loss"]) > 1:
+            N = np.arange(0, len(self.H["loss"]), 1)
+            plt.style.use("ggplot")
+            plt.figure()
+            plt.plot(N, self.H["loss"], label="train_loss")
+            plt.plot(N, self.H["val_loss"], label="val_loss")
+            plt.plot(N, self.H["accuracy"], label="train_acc")
+            plt.plot(N, self.H["val_accuracy"], label="val_acc")
+            plt.title("Training Loss and Accuracy [Epoch {}]".format(len(self.H["loss"])))
+            plt.xlabel("Epoch #")
+            plt.ylabel("Loss/Accuracy")
+            plt.legend()
+            plt.savefig(self.figPath)
+            plt.close()

sidekick/datah/__init__.py

@@ -0,0 +1 @@
+from .loader import Loader

sidekick/datah/loader.py

@@ -0,0 +1,30 @@
+import cv2
+import numpy as np
+import os
+
+class Loader:
+    def __init__(self,preprocessors=None):
+        self.preprocessors=preprocessors
+        if self.preprocessors is None:
+            self.preprocessors=[]
+
+    def load(self,imgpaths,verbose=-1):
+        data=[]
+        labels=[]
+
+        for (i,imgpath) in enumerate(imgpaths):
+            image=cv2.imread(imgpath)
+            label= imgpath.split(os.path.sep)[-2]
+
+            if self.preprocessors is not None:
+                for p in self.preprocessors:
+                    image=p.preprocess(image)
+
+            data.append(image)
+            labels.append(label)
+
+            if verbose > 0 and i >0 and (i+1)% verbose==0:
+                print("processed:{}/{}".format(i+1,len(imgpaths)))
+
+        print("Done!")
+        return (np.array(data),np.array(labels))

sidekick/eval/__init__.py

@@ -0,0 +1 @@
+from .calc_ranked import calculate_ranked

sidekick/eval/calc_ranked.py

@@ -0,0 +1,23 @@
+import numpy as np
+
+def calculate_ranked(preds, labels):
+    rank1=0
+    rank5=0
+
+    for p,l in zip(preds, labels):
+        #sort preds in descending order of their confidence and return the indices of these
+        p= np.argsort(p)[::-1]
+
+        # checking for rank5
+        if l in p[:5]:
+            rank5+=1
+        # checking rank1
+        if l==p[0]:
+            rank1+=1
+
+
+    # Final accuracies
+    rank1= rank1/len(labels)
+    rank5= rank5/len(labels)
+
+    return rank1,rank5

sidekick/io/__init__.py

@@ -0,0 +1 @@
+from .hdf5_writer import Hdf5Writer

sidekick/io/hdf5_writer.py

@@ -0,0 +1,53 @@
+import h5py
+import os
+
+class Hdf5Writer:
+    def __init__(self, dims, outputPath, dbName='Images', buffSize= 1000):
+
+        # throw an error if the file already exists
+        if os.path.exists(outputPath):
+            raise ValueError("PATH ALREADY PRESENT. PLEASE DELETE FILES"
+                             "BEFORE PROCEEDING.")
+
+        # database to store data
+        self.db= h5py.File(outputPath, 'w')
+        # define dataset containers to store data and labels
+        self.data= self.db.create_dataset(dbName, dims, dtype='float')
+        self.labels= self.db.create_dataset('Labels', shape=(dims[0],), dtype='int')
+
+        # defining a buffer and index variable for the buffer
+        self.buffSize= buffSize
+        self.buffer= {"data": [], "labels": []}
+        self.idx= 0
+
+    def add(self, values, labels):
+        self.buffer['data'].extend(values)
+        self.buffer['labels'].extend(labels)
+
+        if len(self.buffer['data'])>=self.buffSize:
+            self.flush()
+
+    def flush(self):
+        # When buffer size is reached flush data to dataset container
+        temp_idx= self.idx + len(self.buffer['data'])
+        # index from prev_idx to new_idx
+        self.data[self.idx:temp_idx]= self.buffer['data']
+        self.labels[self.idx:temp_idx]= self.buffer['labels']
+        # update new_idx
+        self.idx=temp_idx
+        # reinitialize the buffer
+        self.buffer={'data': [], 'labels': []}
+
+    def flushClassNames(self, classNames):
+
+        # Creating a special
+        labelNames= self.db.create_dataset('Label_Names', (len(classNames),),
+                                           dtype=h5py.special_dtype(vlen=unicode))
+        labelNames[:]= classNames
+
+    def close(self):
+
+        if len(self.buffer['data'])>0:
+            self.flush()
+
+        self.db.close()

sidekick/io/hdf5datagen.py

@@ -0,0 +1,50 @@
+from tensorflow.keras.utils import to_categorical
+import h5py
+import numpy as np
+
+class Hdf5DataGen:
+    def __init__(self, dbPath, batchSize, classes, encode=True, aug=None, preprocessors=None):
+
+        self.db= h5py.File(dbPath, 'r')
+        self.batchSize= batchSize
+        self.num_classes= classes
+        self.encode= encode
+        self.aug= aug
+        self.preprocessors= preprocessors
+
+        self.data_length= self.db['Images'].shape[0]
+
+    def generator(self, counter=np.inf):
+        start=0
+
+        while start< counter:
+            for i in np.arange(0, self.data_length, self.batchSize):
+                data = self.db['Images'][i:i+self.batchSize]
+                labels = self.db['Labels'][i:i + self.batchSize]
+
+                if self.encode:
+                    labels= to_categorical(labels, self.num_classes)
+
+                if self.preprocessors is not None:
+                    processed_data=[]
+
+                    for d in data:
+                        for p in self.preprocessors:
+                            d= p.preprocess(d)
+
+                        processed_data.append(d)
+
+                    data= np.array(processed_data)
+
+                if self.aug is not None:
+                    # Notice the next to get next value from generator
+                    data, labels= next(self.aug.flow(
+                        data, labels, batch_size= self.batchSize
+                    ))
+
+                yield (data, labels)
+
+            start+=1
+
+    def close(self):
+        self.db.close()