refactor: Update image captioning script to use Salesforce/blip-image-captioning-base model

scripts/config.py

@@ -4,7 +4,7 @@ VAE_NAME= "madebyollin/sdxl-vae-fp16-fix"
 DATASET_NAME= "hahminlew/kream-product-blip-captions"
 PROJECT_NAME = "Product Photography"
 PRODUCTS_10k_DATASET = "amaye15/Products-10k"
-CAPTIONING_MODEL_NAME = "Salesforce/blip-image-captioning-large"
+CAPTIONING_MODEL_NAME = "Salesforce/blip-image-captioning-base"
 
 
 

scripts/products10k_captions.py

@@ -1,9 +1,10 @@
-from datasets import load_dataset
-from config import PRODUCTS_10k_DATASET, CAPTIONING_MODEL_NAME
+from datasets import load_dataset, Dataset
 from transformers import BlipProcessor, BlipForConditionalGeneration
 from tqdm import tqdm
+from config import PRODUCTS_10k_DATASET,CAPTIONING_MODEL_NAME
 import torch
 
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 class ImageCaptioner:
@@ -14,11 +15,13 @@ class ImageCaptioner:
         dataset (str): The path to the dataset.
         processor (str): The pre-trained processor model to use for image processing.
         model (str): The pre-trained model to use for caption generation.
+        prompt (str): The conditioning prompt to use for caption generation.
 
     Attributes:
         dataset: The loaded dataset.
         processor: The pre-trained processor model.
         model: The pre-trained caption generation model.
+        prompt: The conditioning prompt for generating captions.
 
     Methods:
         process_dataset: Preprocesses the dataset.
@@ -26,10 +29,11 @@ class ImageCaptioner:
 
     """
 
-    def __init__(self, dataset: str, processor: str, model: str):
+    def __init__(self, dataset: str, processor: str, model: str, prompt: str = "Product photo of"):
         self.dataset = load_dataset(dataset, split="train")
         self.processor = BlipProcessor.from_pretrained(processor)
         self.model = BlipForConditionalGeneration.from_pretrained(model).to(device)
+        self.prompt = prompt
 
     def process_dataset(self):
         """
@@ -57,19 +61,25 @@ class ImageCaptioner:
         for idx in tqdm(range(len(self.dataset))):
             image = self.dataset[idx]["image"].convert("RGB")
             inputs = self.processor(images=image, return_tensors="pt").to(device)
-            outputs = self.model.generate(**inputs)
+            prompt_inputs = self.processor(text=[self.prompt], return_tensors="pt").to(device)
+            outputs = self.model.generate(**inputs, **prompt_inputs)
             blip_caption = self.processor.decode(outputs[0], skip_special_tokens=True)
-            self.dataset[idx]["caption"] = blip_caption
-            print(f"Caption for image {idx}: {blip_caption}")
+            self.dataset[idx]["text"] = blip_caption
+      
 
-        # Optionally, you can save the dataset with captions to disk
-        # self.dataset.save_to_disk('path_to_save_dataset')
+       
 
         return self.dataset
 
+# Initialize ImageCaptioner
 ic = ImageCaptioner(
     dataset=PRODUCTS_10k_DATASET,
     processor=CAPTIONING_MODEL_NAME,
     model=CAPTIONING_MODEL_NAME,
+    prompt='Photography of '  # Adding the conditioning prompt
 )
-ic.generate_captions()
+
+# Generate captions for the dataset
+products10k_dataset = ic.generate_captions()
+new_dataset = Dataset.from_pandas(products10k_dataset.to_pandas())  # Convert to a `datasets` Dataset if necessary
+new_dataset.push_to_hub("VikramSingh178/Products-10k-BLIP-captions")