model option added for img2vec

Author: Stealeristaken

tasnif/calculations.py

@@ -6,15 +6,20 @@
 from .logger import info
 
 
-def get_embeddings(use_gpu=False, images=None):
+def get_embeddings(use_gpu=False, images=None, model='resnet-18'):
     """
     This Python function initializes an Img2Vec object, runs it on either GPU or CPU, and retrieves
     image embeddings.
+    :param use_gpu: The `use_gpu` parameter is a boolean that specifies whether to use GPU or CPU.
+    :param images: The `images` parameter is a list of image paths to be used for generating embeddings.
+    :param model: The `model` parameter is a string that specifies the model to use for generating.
+        For available models, see https://github.com/christiansafka/img2vec
+    :return: The function `get_embeddings` returns the embeddings of the images as np.ndarray.
     """
 
     info(f"Img2Vec is running on {'GPU' if use_gpu else 'CPU'}...")
-    img2vec = Img2Vec(cuda=use_gpu)
-
+    img2vec = Img2Vec(cuda=use_gpu, model=model)
+    print(f"Using model: {model}")
     embeddings = img2vec.get_vec(images, tensor=False)
     return embeddings
 

tasnif/tasnif.py

@@ -44,21 +44,23 @@ def read(self, folder_path):
         self.image_paths = read_images_from_directory(folder_path)
         self.images = read_with_pil(self.image_paths)
 
-    def calculate(self, pca=True, iter=10):
+    def calculate(self, pca=True, iter=10, model="resnet-18"):
         """
         The function calculates embeddings, performs PCA, and applies K-means clustering to the
         embeddings. It will not perform these operations if no images have been read.
 
         :param pca: The `pca` parameter is a boolean that specifies whether to perform PCA or not. Default is True
         :param iter: The `iter` parameter is an integer that specifies the number of iterations for the KMeans algorithm. Default is 10.
+        :param model: The `model` parameter is a string that specifies the model to use for generating embeddings. Default is 'resnet-18'.
+            For available models, see https://github.com/christiansafka/img2vec
         """
 
         if not self.images:
             raise ValueError(
                 "The images list can not be empty. Please call the read method before calculating."
             )
 
-        self.embeddings = get_embeddings(use_gpu=self.use_gpu, images=self.images)
+        self.embeddings = get_embeddings(use_gpu=self.use_gpu, images=self.images, model=model)
         if pca:
             self.pca_embeddings = calculate_pca(self.embeddings, self.pca_dim)
             self.centroid, self.labels, self.counts = calculate_kmeans(