Create app_Jim20240322.py

app_Jim20240322.py

@@ -0,0 +1,78 @@
+import streamlit as st
+from transformers import pipeline
+#from diffusers import DiffusionPipeline
+from PIL import Image
+import requests
+import io
+from io import BytesIO
+
+# Load the image-to-text pipeline
+image_to_text = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
+
+# Load the text mask pipeline
+generate_mask = pipeline("fill-mask", model="google-bert/bert-base-uncased")
+
+# Load the text generation pipeline
+extend_text = pipeline("text-generation", model="pranavpsv/genre-story-generator-v2")
+
+# Load the text-to-image model
+#text_to_image = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
+
+def main():
+    st.title("SmartCart (Product Recommender)")
+
+    # User input for text or URL
+    input_option = st.radio("Select input option:", ("Text", "URL"))
+
+    # Input text
+    if input_option == "Text":
+        text_input = st.text_input("Enter the text:")
+        if st.button("Generate Story and Image") and text_input:
+            #generate_image(text_input)
+            generated_text = generate_mask_from_result(text_input)
+            st.success(f'Generated Caption: {text_input}')
+            st.success(f'Generated Text: {generated_text}')
+            
+
+    # Input URL
+    elif input_option == "URL":
+        image_url = st.text_input("Enter the image URL:")
+        if st.button("Generate Story and Image") and image_url:
+            image_text = image_to_text_from_url(image_url)
+            #generate_image(image_text)
+            generated_text = generate_mask_from_result(image_text)
+            st.success(f'Generated Caption: {image_text}')
+            st.success(f'Generated Text: {generated_text}')
+            
+
+def image_to_text_from_file(uploaded_file):
+    image_bytes = io.BytesIO(uploaded_file.read())
+    return image_to_text(image_bytes)[0]['generated_text']
+
+def image_to_text_from_url(image_url):
+    response = requests.get(image_url)
+    image_bytes = Image.open(BytesIO(response.content))
+    return image_to_text(image_bytes)[0]['generated_text']
+
+def generate_image(text):
+    rephrased_text = "I want to buy " + text + " and [MASK] for my children"
+    generated_image = text_to_image(rephrased_text)
+    st.image(generated_image, caption="Generated Image", use_column_width=True)
+
+def generate_mask_from_result(text):
+    output = generate_mask(f"I want to buy 2 toys for my children. I will buy {text} and [MASK].")
+    
+    if output and output[0]['token_str'] == text:
+        # If the first result matches the input, get the second output instead
+        second_output = output[1] if len(output) > 1 else None
+        result = second_output['token_str'] if second_output else None
+    else:
+        result = output[0]['token_str'] if output else None
+
+    extended_text = extend_text(f"A child with {text} and {result} ")
+    return extended_text[0]['generated_text']
+
+
+if __name__ == "__main__":
+    main()
+