Create app.py

app.py

@@ -0,0 +1,125 @@
+import streamlit as st
+import tensorflow as tf
+import numpy as np
+from PIL import Image
+import requests
+from io import BytesIO
+from selenium import webdriver
+from selenium.webdriver.chrome.service import Service
+from selenium.webdriver.common.by import By
+from webdriver_manager.chrome import ChromeDriverManager
+import time
+import pandas as pd
+import base64
+
+def load_model(model_path):
+    interpreter = tf.lite.Interpreter(model_path=model_path)
+    interpreter.allocate_tensors()
+    return interpreter
+
+def preprocess_image(image, input_size):
+    image = image.convert('RGB')
+    image = image.resize((input_size, input_size))
+    image_np = np.array(image, dtype=np.float32)
+    image_np = np.expand_dims(image_np, axis=0)
+    image_np = image_np / 255.0  # Normalize to [0, 1]
+    return image_np
+
+def run_inference(interpreter, input_data):
+    input_details = interpreter.get_input_details()
+    output_details = interpreter.get_output_details()
+    
+    interpreter.set_tensor(input_details[0]['index'], input_data)
+    interpreter.invoke()
+    
+    output_data_shopping_intent = interpreter.get_tensor(output_details[0]['index'])
+    output_data_sensitive = interpreter.get_tensor(output_details[1]['index'])
+    
+    return output_data_shopping_intent, output_data_sensitive
+
+def fetch_images_from_url(url):
+    options = webdriver.ChromeOptions()
+    options.add_argument('--headless')
+    options.add_argument('--no-sandbox')
+    options.add_argument('--disable-dev-shm-usage')
+    options.add_argument('--disable-gpu')
+
+    service = Service(ChromeDriverManager().install())
+    driver = webdriver.Chrome(service=service, options=options)
+    driver.get(url)
+
+    # Give the page some time to load and execute JavaScript
+    time.sleep(10)
+    
+    images = driver.find_elements(By.TAG_NAME, 'img')
+    img_urls = [img.get_attribute('src') for img in images if img.get_attribute('src')]
+
+    driver.quit()
+    return img_urls
+
+def image_to_base64(image):
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode()
+
+def main():
+    st.set_page_config(layout="wide")
+    st.title("Image Classification with TFLite")
+    st.write("Enter a URL to fetch and classify all images on the page.")
+
+    model_path = "model.tflite"
+    url = st.text_input("Enter URL")
+
+    if url:
+        img_urls = fetch_images_from_url(url)
+        if img_urls:
+            st.write(f"Found {len(img_urls)} images")
+            interpreter = load_model(model_path)
+            input_details = interpreter.get_input_details()
+            input_shape = input_details[0]['shape']
+            input_size = input_shape[1]  # assuming square input
+
+            data = []
+            errors = []
+
+            for img_url in img_urls:
+                try:
+                    response = requests.get(img_url)
+                    image = Image.open(BytesIO(response.content))
+
+                    input_data = preprocess_image(image, input_size)
+                    output_data_shopping_intent, output_data_sensitive = run_inference(interpreter, input_data)
+
+                    # Convert image to Base64
+                    image.thumbnail((100, 100))
+                    thumbnail_base64 = image_to_base64(image)
+                    thumbnail_data_url = f"data:image/png;base64,{thumbnail_base64}"
+
+                    data.append({
+                        'Thumbnail': thumbnail_data_url,
+                        'URL': img_url,
+                        'Shopping Intent': output_data_shopping_intent.flatten().tolist(),
+                        'Sensitivity': output_data_sensitive.flatten().tolist()
+                    })
+                except Exception as e:
+                    errors.append(f"Could not process image {img_url}: {e}")
+
+            # Convert data to DataFrame
+            df = pd.DataFrame(data)
+
+            # Configure DataFrame display with images, URLs, and classifications
+            st.data_editor(df, column_config={
+                "Thumbnail": st.column_config.ImageColumn("Thumbnail", help="Image thumbnails"),
+                "URL": st.column_config.LinkColumn("URL"),
+                "Shopping Intent": st.column_config.BarChartColumn("Shopping Intent", width="small"),
+                "Sensitivity": st.column_config.BarChartColumn("Sensitivity", width="small")
+            })
+
+            # Display errors in an expandable section
+            if errors:
+                with st.expander(f"Could not process {len(errors)} images"):
+                    for error in errors:
+                        st.write(error)
+
+if __name__ == "__main__":
+    main()