Added to handle 3 different image formats

somasundaram1702 · Aug 17, 2020 · 7f0a499 · 7f0a499
1 parent c26d7eb
commit 7f0a499
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/extract_embeds.py b/extract_embeds.py
@@ -48,6 +48,9 @@ def crop_resize(self):
         return(self.face)
 
     def check_folders(self):
+        if 'outputs' not in os.listdir():
+            os.makedirs('./outputs')
+
         y,x = sorted(os.listdir(self.root_path),key=len)
         assert(x == 'train' or x == 'Train' or x == 'TRAIN')
         assert(y == 'test' or y == 'Test' or y == 'TEST')
@@ -86,7 +89,7 @@ def extract_faces(self):
 
             print('Images of '+label + ' under process')
             for pics in os.listdir(self.test_path+'/'+label):
-              if pics.endswith('.jpg'):
+              if pics.endswith('.jpg') or pics.endswith('.jpeg') or pics.endswith('.png') or pics.endswith('.tif'):
                 #print('Name of the pic:',pics)
                 try:
                   face_details = self.identify_face(self.test_path+'/'+label+'/'+pics)
@@ -97,7 +100,7 @@ def extract_faces(self):
                   print('Pic {} in {} is not processed'.format(pics,label))
                   continue
         print('time taken to process {} secs'.format(time.time()-start2))
-        savez_compressed('Extracted_faces.npz', self.X_train_faces, self.y_train_faces, self.X_test_faces, self.y_test_faces)
+        savez_compressed('./outputs/Extracted_faces.npz', self.X_train_faces, self.y_train_faces, self.X_test_faces, self.y_test_faces)
         return(self.X_train_faces, self.y_train_faces, self.X_test_faces, self.y_test_faces)
 
     def load_facenet(self):
@@ -133,7 +136,7 @@ def extract_embeddings(self):
         self.newTestX = asarray(self.newTestX)
         print(self.newTestX.shape)                                      
         # save arrays to one file in compressed format
-        savez_compressed('Extracted_embeddings.npz', self.newTrainX, self.y_train_faces, self.newTestX, self.y_test_faces)
+        savez_compressed('./outputs/Extracted_embeddings.npz', self.newTrainX, self.y_train_faces, self.newTestX, self.y_test_faces)
         print('Time taken from reading images to extracting embeds: {} s'.format(time.time()-beginn))
 
 if __name__ == '__main__':
@@ -147,4 +150,4 @@ def extract_embeddings(self):
     _extract_embeds.check_folders()
     x,y,a,b = _extract_embeds.extract_faces()
     _extract_embeds.load_facenet()
-    _extract_embeds.extract_embeddings()
+    _extract_embeds.extract_embeddings()