Update app.py

app.py

@@ -1,40 +1,17 @@
-import streamlit as st
-import requests
-from PIL import Image
-from transformers import BlipProcessor, BlipForConditionalGeneration, pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
+im-port streamlit as st
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 import nltk
 from youtube_transcript_api import YouTubeTranscriptApi
 
 # Download NLTK data
 nltk.download('punkt')
 
-# Initialize the image captioning processor and model
-caption_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-caption_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
+# Initialize the image captioning pipeline
+captioner = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
 
-# Initialize the tokenizer and model for tag generation
-tag_tokenizer = AutoTokenizer.from_pretrained("fabiochiu/t5-base-tag-generation")
-tag_model = AutoModelForSeq2SeqLM.from_pretrained("fabiochiu/t5-base-tag-generation")
-
-# Function to generate captions for an image
-def generate_caption(img_url, text="a photography of"):
-    try:
-        raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')
-    except Exception as e:
-        st.error(f"Error loading image: {e}")
-        return None, None
-
-    # Conditional image captioning
-    inputs_conditional = caption_processor(raw_image, text, return_tensors="pt")
-    out_conditional = caption_model.generate(**inputs_conditional)
-    caption_conditional = caption_processor.decode(out_conditional[0], skip_special_tokens=True)
-
-    # Unconditional image captioning
-    inputs_unconditional = caption_processor(raw_image, return_tensors="pt")
-    out_unconditional = caption_model.generate(**inputs_unconditional)
-    caption_unconditional = caption_processor.decode(out_unconditional[0], skip_special_tokens=True)
-
-    return caption_conditional, caption_unconditional
+# Load the tokenizer and model for tag generation
+tokenizer = AutoTokenizer.from_pretrained("fabiochiu/t5-base-tag-generation")
+model = AutoModelForSeq2SeqLM.from_pretrained("fabiochiu/t5-base-tag-generation")
 
 # Function to fetch YouTube transcript
 def fetch_transcript(url):
@@ -47,51 +24,30 @@ def fetch_transcript(url):
         return str(e)
 
 # Streamlit app title
-st.title("Multi-purpose Machine Learning App")
+st.title("Multi-purpose Machine Learning App: WAVE_AI")
 
 # Create tabs for different functionalities
-tab1, tab2, tab3 = st.tabs(["Image Captioning", "Text Tag Generation", "YouTube Transcript"])
+tab1, tab2, tab3 = st.tabs(["Text Tag Generation", "Image Captioning", "YouTube Transcript"])
 
 # Image Captioning Tab
 with tab1:
-    st.header("Image Captioning")
-    
-    # Input for image URL
-    img_url = st.text_input("Enter Image URL:")
-    
-    # If an image URL is provided
-    if st.button("Generate Captions", key='caption_button'):
-        if img_url:
-            caption_conditional, caption_unconditional = generate_caption(img_url)
-            if caption_conditional and caption_unconditional:
-                st.success("Captions successfully generated!")
-                st.image(img_url, caption="Input Image", use_column_width=True)
-                st.write("### Conditional Caption")
-                st.write(caption_conditional)
-                st.write("### Unconditional Caption")
-                st.write(caption_unconditional)
-        else:
-            st.warning("Please enter an image URL.")
-
-# Text Tag Generation Tab
-with tab2:
     st.header("Text Tag Generation")
     
     # Text area for user input
     text = st.text_area("Enter the text for tag extraction:", height=200)
     
     # Button to generate tags
-    if st.button("Generate Tags", key='tag_button'):
+    if st.button("Generate Tags"):
         if text:
             try:
                 # Tokenize and encode the input text
-                inputs = tag_tokenizer([text], max_length=512, truncation=True, return_tensors="pt")
+                inputs = tokenizer([text], max_length=512, truncation=True, return_tensors="pt")
                 
                 # Generate tags
-                output = tag_model.generate(**inputs, num_beams=8, do_sample=True, min_length=10, max_length=64)
+                output = model.generate(**inputs, num_beams=8, do_sample=True, min_length=10, max_length=64)
                 
                 # Decode the output
-                decoded_output = tag_tokenizer.batch_decode(output, skip_special_tokens=True)[0]
+                decoded_output = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
                 
                 # Extract unique tags
                 tags = list(set(decoded_output.strip().split(", ")))
@@ -104,6 +60,28 @@ with tab2:
         else:
             st.warning("Please enter some text to generate tags.")
 
+# Text Tag Generation Tab
+with tab2:
+    st.header("Image Captioning Extractor")
+    
+    # Input for image URL
+    image_url = st.text_input("Enter the URL of the image:")
+    
+    # If an image URL is provided
+    if image_url:
+        try:
+            # Display the image
+            st.image(image_url, caption="Provided Image", use_column_width=True)
+            
+            # Generate the caption
+            caption = captioner(image_url)
+            
+            # Display the caption
+            st.write("**Generated Caption:**")
+            st.write(caption[0]['generated_text'])
+        except Exception as e:
+            st.error(f"An error occurred: {e}")
+
 # YouTube Transcript Tab
 with tab3:
     st.header("YouTube Video Transcript Extractor")
@@ -112,7 +90,7 @@ with tab3:
     youtube_url = st.text_input("Enter YouTube URL:")
     
     # Button to get transcript
-    if st.button("Get Transcript", key='transcript_button'):
+    if st.button("Get Transcript"):
         if youtube_url:
             transcript = fetch_transcript(youtube_url)
             if "error" not in transcript.lower():