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

product_diffusion_api/endpoints.py

@@ -5,6 +5,7 @@ from routers import painting
 
 
 
+
 app = FastAPI(openapi_url='/api/v1/product-diffusion/openapi.json',docs_url='/api/v1/product_diffusion/docs')
 app.add_middleware(
     CORSMiddleware,
@@ -38,4 +39,5 @@ async def root():
     
 @app.get("/health")
 def check_health():
-    return {"status": "ok"}
+    return {"status": "ok"}
+

scripts/products10k_captions.py

@@ -1,9 +1,10 @@
+import torch
 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
 
+# Assuming PRODUCTS_10k_DATASET and CAPTIONING_MODEL_NAME are defined in config.py
+from config import PRODUCTS_10k_DATASET, CAPTIONING_MODEL_NAME
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -25,12 +26,13 @@ class ImageCaptioner:
 
     Methods:
         process_dataset: Preprocesses the dataset.
-        generate_captions: Generates captions for the images in the dataset.
-
+        generate_caption: Generates a caption for a single image.
+        generate_captions: Generates captions for all images in the dataset.
     """
 
     def __init__(self, dataset: str, processor: str, model: str, prompt: str = "Product photo of"):
-        self.dataset = load_dataset(dataset, split="train")
+        self.dataset = load_dataset(dataset, split="test")
+        self.dataset = self.dataset.select(range(10000))  # For demonstration purposes
         self.processor = BlipProcessor.from_pretrained(processor)
         self.model = BlipForConditionalGeneration.from_pretrained(model).to(device)
         self.prompt = prompt
@@ -41,34 +43,48 @@ class ImageCaptioner:
 
         Returns:
             The preprocessed dataset.
-
         """
-        self.dataset = self.dataset.rename_column("pixel_values", "image")
+        # Check if 'image' column exists, otherwise use 'pixel_values' if it exists
+        image_column = "image" if "image" in self.dataset.column_names else "pixel_values"
+        self.dataset = self.dataset.rename_column(image_column, "image")
+
         if "label" in self.dataset.column_names:
             self.dataset = self.dataset.remove_columns(["label"])
+
+        # Add an empty 'text' column for captions if it doesn't exist
+        if "text" not in self.dataset.column_names:
+            new_column = [""] * len(self.dataset)
+            self.dataset = self.dataset.add_column("text", new_column)
+
         return self.dataset
 
-    def generate_captions(self):
+    def generate_caption(self, example):
         """
-        Generates captions for the images in the dataset.
+        Generates a caption for a single image.
 
-        Returns:
-            The dataset with captions.
+        Args:
+            example (dict): A dictionary containing the image data.
 
+        Returns:
+            dict: The dictionary with the generated caption.
         """
-        self.dataset = self.process_dataset()
-
-        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)
-            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]["text"] = blip_caption
-      
+        image = example["image"].convert("RGB")
+        inputs = self.processor(images=image, return_tensors="pt").to(device)
+        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)
+        example["text"] = blip_caption
+        return example
 
-       
+    def generate_captions(self):
+        """
+        Generates captions for all images in the dataset.
 
+        Returns:
+            Dataset: The dataset with generated captions.
+        """
+        self.dataset = self.process_dataset()
+        self.dataset = self.dataset.map(self.generate_caption, batched=False)
         return self.dataset
 
 # Initialize ImageCaptioner
@@ -76,10 +92,11 @@ ic = ImageCaptioner(
     dataset=PRODUCTS_10k_DATASET,
     processor=CAPTIONING_MODEL_NAME,
     model=CAPTIONING_MODEL_NAME,
-    prompt='Photography of '  # Adding the conditioning prompt
+    prompt='Commercial photography of'
 )
 
 # 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")
+
+# Save the dataset to the hub
+products10k_dataset.push_to_hub("VikramSingh178/Products-10k-BLIP-captions")