GPU and batching support for captioning

Author: guy-singer

fastdup/captions.py

@@ -4,7 +4,7 @@
 import cv2
 
 
-def generate_labels(filenames, modelname='automatic', batch_size=8):
+def generate_labels(filenames, model_name='automatic', device = -1, batch_size=8):
     '''
     This function generates captions for a given set of images, and takes the following arguments:
         - filenames: the list of images passed to the function
@@ -14,7 +14,9 @@ def generate_labels(filenames, modelname='automatic', batch_size=8):
             - BLIP-2: 'blip2'
             - BLIP: 'blip'
         - batch_size: the size of image batches to caption (default: 8)
+        - device: whether to use a GPU (default: -1, CPU only ; set to 0 for GPU)
     '''
+    # use GPU if device is specified
 
     # confirm necessary dependencies are installed, and import them
     try:
@@ -39,12 +41,11 @@ def generate_labels(filenames, modelname='automatic', batch_size=8):
         'blip': "Salesforce/blip-image-captioning-large"
     }
 
-    model = models[modelname]
+    model = models[model_name]
 
     # generate captions
     try:
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        captioner = pipeline("image-to-text", model=model, device=device, batch_size=batch_size)
+        captioner = pipeline("image-to-text", model=model, device=device, batch_size=batch_size, device=device)
 
         captions = []
         for image_path in tqdm(filenames):

fastdup/fastdup_controller.py

@@ -1260,7 +1260,7 @@ def _verify_fastdup_run_args(self, input_dir, work_dir, df_annot, subset, data_t
             else:
                 assert False, f"Wrong data type {data_type}"
 
-    def caption(self, model_name='automatic', subset: list = None, vqa_prompt: str = None, kwargs=None) -> pd.DataFrame:
+    def caption(self, model_name='automatic', device = -1, batch_size: int = 8, subset: list = None, vqa_prompt: str = None, kwargs=None) -> pd.DataFrame:
         if not self._fastdup_applied:
             raise RuntimeError('Fastdup was not applied yet, call run() first')
 
@@ -1272,7 +1272,7 @@ def caption(self, model_name='automatic', subset: list = None, vqa_prompt: str =
 
         if model_name in FD.CAPTION_MODEL_NAMES:
             from fastdup.captions import generate_labels
-            df['caption'] = generate_labels(df['filename'], model_name)
+            df['caption'] = generate_labels(df['filename'], model_name, device, batch_size)
         elif model_name == FD.VQA_MODEL1_NAME:
             from fastdup.captions import generate_vqa_labels
             df['caption'] = generate_vqa_labels(df['filename'], vqa_prompt, kwargs)

fastdup/galleries.py

@@ -74,7 +74,7 @@ def swap_dataframe(subdf, cols):
 
 
 
-def find_label(get_label_func, df, in_col, out_col, vqa_prompt: str = None, kwargs=None):
+def find_label(get_label_func, df, in_col, out_col, vqa_prompt: str = None, device = -1, kwargs=None):
 
 
     if (get_label_func is not None):
@@ -87,7 +87,7 @@ def find_label(get_label_func, df, in_col, out_col, vqa_prompt: str = None, kwar
                 df[out_col] = df['label']
             elif get_label_func in CAPTION_MODEL_NAMES:
                 from fastdup.captions import generate_labels
-                df[out_col] = generate_labels(df[in_col], get_label_func)
+                df[out_col] = generate_labels(df[in_col], get_label_func, device)
             elif get_label_func == VQA_MODEL1_NAME:
                 from fastdup.captions import generate_vqa_labels
                 df[out_col] = generate_vqa_labels(df[in_col], vqa_prompt, kwargs)