from io import BytesIO

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 from extract import take_screenshot
 from PIL import Image
+from io import BytesIO
 
 def main():
     st.title("Website Visualizer")

extract.py

@@ -1,6 +1,7 @@
 from selenium import webdriver
 from selenium.common.exceptions import WebDriverException
 from PIL import Image
+from io import BytesIO
 
 def take_screenshot(url):
     options = webdriver.ChromeOptions()