File size: 4,118 Bytes
2b86da8 f6ed594 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 |
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, labels):
image_assembler = ImageAssembler() \
.setInputCol("image") \
.setOutputCol("image_assembler")
imageClassifier = CLIPForZeroShotClassification \
.pretrained() \
.setInputCols(["image_assembler"]) \
.setOutputCol("label") \
.setCandidateLabels(labels)
pipeline = Pipeline(stages=[
image_assembler,
imageClassifier,
])
return pipeline
def fit_data(pipeline, data):
model = pipeline.fit(data)
light_pipeline = LightPipeline(model)
annotations_result = light_pipeline.fullAnnotateImage(data)
return annotations_result[0]['label'][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 = st.sidebar.selectbox(
"Choose the pretrained model",
["CLIPForZeroShotClassification"],
help="For more info about the models visit: https://sparknlp.org/models"
)
# Set up the page layout
st.markdown(f'<div class="main-title">CLIPForZeroShotClassification</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/CLIPForZeroShotClassification.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 = "input"
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'])
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}"
candidateLabels = [
"a photo of a bird",
"a photo of a cat",
"a photo of a dog",
"a photo of a hen",
"a photo of a hippo",
"a photo of a room",
"a photo of a tractor",
"a photo of an ostrich",
"a photo of an ox"]
lables = st_tags(
label='Select labels',
text='Press enter to add more',
value=candidateLabels,
maxtags = -1)
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')
init_spark()
Pipeline = create_pipeline(model, lables)
output = fit_data(Pipeline, selected_image)
st.markdown(f'This document has been classified as : **{output}**') |