Update app.py

app.py

@@ -1,9 +1,10 @@
-import os
+# import os
 
-os.system('pip install --upgrade pip')
+# os.system('pip install --upgrade pip')
 
 import streamlit as st
 from transformers import BlipForConditionalGeneration, AutoTokenizer
+import torchvision.transforms as transforms
 import torch
 from PIL import Image
 from io import BytesIO
@@ -21,17 +22,35 @@ model = BlipForConditionalGeneration.from_pretrained("MLInAi/CartoonCaptionGen")
 tokenizer = AutoTokenizer.from_pretrained("MLInAi/CartoonCaptionGen")
 
 # Function to generate caption for the uploaded image
+# def generate_caption(image):
+#     # Preprocess the image
+#     image = Image.open(image).convert("RGB")
+#     image = image.resize((224, 224))  # Resize the image to match model input size
+#     image_tensor = torch.tensor([torch.Tensor(image)]).permute(0, 3, 1, 2).to(device)
+
+#     # Generate caption
+#     output = model.generate(pixel_values=image_tensor)
+#     caption = tokenizer.decode(output[0], skip_special_tokens=True)
+#     return caption
+
+
 def generate_caption(image):
     # Preprocess the image
     image = Image.open(image).convert("RGB")
     image = image.resize((224, 224))  # Resize the image to match model input size
-    image_tensor = torch.tensor([torch.Tensor(image)]).permute(0, 3, 1, 2).to(device)
+    
+    # Convert the image to a tensor
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+    image_tensor = transform(image).unsqueeze(0).to(device)
 
     # Generate caption
     output = model.generate(pixel_values=image_tensor)
     caption = tokenizer.decode(output[0], skip_special_tokens=True)
     return caption
-
+   
 # Streamlit app
 st.title("Image Caption Generator")