Update Demo.py

Demo.py

@@ -1,127 +1,127 @@
-import streamlit as st
-import sparknlp
-import os
-import pandas as pd
-
-from sparknlp.base import *
-from sparknlp.annotator import *
-from pyspark.ml import Pipeline
-from sparknlp.pretrained import PretrainedPipeline
-from streamlit_tags import st_tags
-
-# Page configuration
-st.set_page_config(
-    layout="wide", 
-    initial_sidebar_state="auto"
-)
-
-# CSS for styling
-st.markdown("""
-    <style>
-        .main-title {
-            font-size: 36px;
-            color: #4A90E2;
-            font-weight: bold;
-            text-align: center;
-        }
-        .section {
-            background-color: #f9f9f9;
-            padding: 10px;
-            border-radius: 10px;
-            margin-top: 10px;
-        }
-        .section p, .section ul {
-            color: #666666;
-        }
-    </style>
-""", unsafe_allow_html=True)
-
-@st.cache_resource
-def init_spark():
-    return sparknlp.start()
-
-@st.cache_resource
-def create_pipeline(model):
-    image_assembler = ImageAssembler() \
-                .setInputCol("image") \
-                .setOutputCol("image_assembler")
-
-    image_classifier = ViTForImageClassification \
-        .pretrained(model) \
-        .setInputCols("image_assembler") \
-        .setOutputCol("class")
-
-    pipeline = Pipeline(stages=[
-        image_assembler,
-        image_classifier,
-    ])
-    return pipeline
-
-def fit_data(pipeline, data):
-    empty_df = spark.createDataFrame([['']]).toDF('text')
-    model = pipeline.fit(empty_df)
-    light_pipeline = LightPipeline(model)
-    annotations_result = light_pipeline.fullAnnotateImage(data)
-    return annotations_result[0]['class'][0].result
-
-def save_uploadedfile(uploadedfile):
-    filepath = os.path.join(IMAGE_FILE_PATH, uploadedfile.name)
-    with open(filepath, "wb") as f:
-        if hasattr(uploadedfile, 'getbuffer'):
-            f.write(uploadedfile.getbuffer())
-        else:
-            f.write(uploadedfile.read())
-        
-# Sidebar content
-model_list =   ['image_classifier_vit_base_cats_vs_dogs', 'image_classifier_vit_base_patch16_224', 'image_classifier_vit_CarViT', 'image_classifier_vit_base_beans_demo', 'image_classifier_vit_base_food101', 'image_classifier_vit_base_patch16_224_in21k_finetuned_cifar10']
-model = st.sidebar.selectbox(
-    "Choose the pretrained model",
-    model_list,
-    help="For more info about the models visit: https://sparknlp.org/models"
-)
-
-# Set up the page layout
-st.markdown(f'<div class="main-title">ViT for Image Classification</div>', unsafe_allow_html=True)
-# st.markdown(f'<div class="section"><p>{sub_title}</p></div>', unsafe_allow_html=True)
-
-# Reference notebook link in sidebar
-link = """
-<a href="https://github.com/JohnSnowLabs/spark-nlp/blob/master/examples/python/annotation/image/ViTForImageClassification.ipynb">
-    <img src="https://colab.research.google.com/assets/colab-badge.svg" style="zoom: 1.3" alt="Open In Colab"/>
-</a>
-"""
-st.sidebar.markdown('Reference notebook:')
-st.sidebar.markdown(link, unsafe_allow_html=True)
-
-# Load examples
-IMAGE_FILE_PATH = f"/content/sparknlp VIT Image Classification/inputs/{model}"
-image_files = sorted([file for file in os.listdir(IMAGE_FILE_PATH) if file.split('.')[-1]=='png' or file.split('.')[-1]=='jpg' or file.split('.')[-1]=='JPEG' or file.split('.')[-1]=='jpeg'])
-
-st.subheader("This model identifies image classes using the vision transformer (ViT).")
-
-img_options = st.selectbox("Select an image", image_files)
-uploadedfile = st.file_uploader("Try it for yourself!")
-
-if uploadedfile:
-    file_details = {"FileName":uploadedfile.name,"FileType":uploadedfile.type}
-    save_uploadedfile(uploadedfile)
-    selected_image = f"{IMAGE_FILE_PATH}/{uploadedfile.name}"
-elif img_options:
-    selected_image = f"{IMAGE_FILE_PATH}/{img_options}"
-
-st.subheader('Classified Image')
-
-image_size = st.slider('Image Size', 400, 1000, value=400, step = 100)
-
-try:
-    st.image(f"{IMAGE_FILE_PATH}/{selected_image}", width=image_size)
-except:
-    st.image(selected_image, width=image_size)
-
-st.subheader('Classification')
-
-spark = init_spark()
-Pipeline = create_pipeline(model)
-output = fit_data(Pipeline, selected_image)
-
+import streamlit as st
+import sparknlp
+import os
+import pandas as pd
+
+from sparknlp.base import *
+from sparknlp.annotator import *
+from pyspark.ml import Pipeline
+from sparknlp.pretrained import PretrainedPipeline
+from streamlit_tags import st_tags
+
+# Page configuration
+st.set_page_config(
+    layout="wide", 
+    initial_sidebar_state="auto"
+)
+
+# CSS for styling
+st.markdown("""
+    <style>
+        .main-title {
+            font-size: 36px;
+            color: #4A90E2;
+            font-weight: bold;
+            text-align: center;
+        }
+        .section {
+            background-color: #f9f9f9;
+            padding: 10px;
+            border-radius: 10px;
+            margin-top: 10px;
+        }
+        .section p, .section ul {
+            color: #666666;
+        }
+    </style>
+""", unsafe_allow_html=True)
+
+@st.cache_resource
+def init_spark():
+    return sparknlp.start()
+
+@st.cache_resource
+def create_pipeline(model):
+    image_assembler = ImageAssembler() \
+                .setInputCol("image") \
+                .setOutputCol("image_assembler")
+
+    image_classifier = ViTForImageClassification \
+        .pretrained(model) \
+        .setInputCols("image_assembler") \
+        .setOutputCol("class")
+
+    pipeline = Pipeline(stages=[
+        image_assembler,
+        image_classifier,
+    ])
+    return pipeline
+
+def fit_data(pipeline, data):
+    empty_df = spark.createDataFrame([['']]).toDF('text')
+    model = pipeline.fit(empty_df)
+    light_pipeline = LightPipeline(model)
+    annotations_result = light_pipeline.fullAnnotateImage(data)
+    return annotations_result[0]['class'][0].result
+
+def save_uploadedfile(uploadedfile):
+    filepath = os.path.join(IMAGE_FILE_PATH, uploadedfile.name)
+    with open(filepath, "wb") as f:
+        if hasattr(uploadedfile, 'getbuffer'):
+            f.write(uploadedfile.getbuffer())
+        else:
+            f.write(uploadedfile.read())
+        
+# Sidebar content
+model_list =   ['image_classifier_vit_base_cats_vs_dogs', 'image_classifier_vit_base_patch16_224', 'image_classifier_vit_CarViT', 'image_classifier_vit_base_beans_demo', 'image_classifier_vit_base_food101', 'image_classifier_vit_base_patch16_224_in21k_finetuned_cifar10']
+model = st.sidebar.selectbox(
+    "Choose the pretrained model",
+    model_list,
+    help="For more info about the models visit: https://sparknlp.org/models"
+)
+
+# Set up the page layout
+st.markdown(f'<div class="main-title">ViT for Image Classification</div>', unsafe_allow_html=True)
+# st.markdown(f'<div class="section"><p>{sub_title}</p></div>', unsafe_allow_html=True)
+
+# Reference notebook link in sidebar
+link = """
+<a href="https://github.com/JohnSnowLabs/spark-nlp/blob/master/examples/python/annotation/image/ViTForImageClassification.ipynb">
+    <img src="https://colab.research.google.com/assets/colab-badge.svg" style="zoom: 1.3" alt="Open In Colab"/>
+</a>
+"""
+st.sidebar.markdown('Reference notebook:')
+st.sidebar.markdown(link, unsafe_allow_html=True)
+
+# Load examples
+IMAGE_FILE_PATH = f"inputs/{model}"
+image_files = sorted([file for file in os.listdir(IMAGE_FILE_PATH) if file.split('.')[-1]=='png' or file.split('.')[-1]=='jpg' or file.split('.')[-1]=='JPEG' or file.split('.')[-1]=='jpeg'])
+
+st.subheader("This model identifies image classes using the vision transformer (ViT).")
+
+img_options = st.selectbox("Select an image", image_files)
+uploadedfile = st.file_uploader("Try it for yourself!")
+
+if uploadedfile:
+    file_details = {"FileName":uploadedfile.name,"FileType":uploadedfile.type}
+    save_uploadedfile(uploadedfile)
+    selected_image = f"{IMAGE_FILE_PATH}/{uploadedfile.name}"
+elif img_options:
+    selected_image = f"{IMAGE_FILE_PATH}/{img_options}"
+
+st.subheader('Classified Image')
+
+image_size = st.slider('Image Size', 400, 1000, value=400, step = 100)
+
+try:
+    st.image(f"{IMAGE_FILE_PATH}/{selected_image}", width=image_size)
+except:
+    st.image(selected_image, width=image_size)
+
+st.subheader('Classification')
+
+spark = init_spark()
+Pipeline = create_pipeline(model)
+output = fit_data(Pipeline, selected_image)
+
 st.markdown(f'This document has been classified as  : **{output}**')