Update app.py

app.py

@@ -1,35 +1,35 @@
 import streamlit as st
+import cv2
+import pandas as pd
 from PIL import Image
 from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer, pipeline, AutoModelForSeq2SeqLM
-from datetime import datetime
-import pandas as pd
-import tempfile
-import base64
 import nltk
+import tempfile
+from nltk.corpus import wordnet
 import spacy
 from spacy.cli import download
-from streamlit_option_menu import option_menu
-import torch
-
-# Set device
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+import base64
+import numpy as np
+import datetime
 
-# Download necessary NLTK and spaCy data
+# Download necessary NLP models
 nltk.download('wordnet')
 nltk.download('omw-1.4')
 download("en_core_web_sm")
-
-# Load the models
 nlp = spacy.load("en_core_web_sm")
-model_name = "NourFakih/Vit-GPT2-COCO2017Flickr-115k-12"
+
+# Load the pre-trained models for image captioning and summarization
+model_name = "NourFakih/Vit-GPT2-COCO2017Flickr-85k-09"
 model = VisionEncoderDecoderModel.from_pretrained(model_name)
 feature_extractor = ViTImageProcessor.from_pretrained(model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# GPT2 only has bos/eos tokens but not decoder_start/pad tokens
 tokenizer.pad_token = tokenizer.eos_token
+# update the model config
 model.config.eos_token_id = tokenizer.eos_token_id
 model.config.decoder_start_token_id = tokenizer.bos_token_id
 model.config.pad_token_id = tokenizer.pad_token_id
-image_captioner = pipeline('image-to-text', model=model_name)
 
 model_sum_name = "google-t5/t5-base"
 tokenizer_sum = AutoTokenizer.from_pretrained("google-t5/t5-base")
@@ -40,8 +40,10 @@ if 'captured_images' not in st.session_state:
     st.session_state.captured_images = []
 
 def generate_caption(image):
-    caption = image_captioner(image)
-    return caption[0]['generated_text']
+    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
+    output_ids = model.generate(pixel_values)
+    caption = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    return caption
 
 def get_synonyms(word):
     synonyms = set()
@@ -88,9 +90,9 @@ def page_webcam_capture():
 
     if img_file:
         img = Image.open(img_file)
-        img_array = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+        img_array = np.array(img)
         caption = generate_caption(img)
-        capture_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        capture_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
         add_image_to_state(img_array, caption, capture_time)
         st.image(img, caption=f"Caption: {caption}")