Update app.py

app.py

@@ -2,6 +2,9 @@ import streamlit as st
 from io import BytesIO
 import PyPDF2
 import docx2txt
+from PIL import Image
+import torch
+from transformers import AutoProcessor, AutoModelForImageClassification
 
 st.set_page_config(page_title="📝 Note Input", layout="centered")
 st.title("📄 Type of Text Input")
@@ -22,10 +25,21 @@ def format_single_line(text: str) -> str:
 def extract_docx_text(uploaded_file):
     # docx2txt.process accepts a path or a file-like object
     return docx2txt.process(uploaded_file)
+# --- Image model setup ---
+MODEL_NAME = "google/vit-base-patch16-224"
 
+@st.cache_resource
+def load_image_model():
+    proc = AutoProcessor.from_pretrained(MODEL_NAME)
+    mdl  = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
+    return proc, mdl
+
+processor, model = load_image_model()
+
+# --- Main UI ---
 input_type = st.selectbox(
     "Select the type of input:",
-    ["Select...", "PDF", "Word Document", "Text"]
+    ["Select...", "PDF", "Word Document", "Text", "Image"]
 )
 
 if input_type == "PDF":
@@ -53,5 +67,24 @@ elif input_type == "Text":
         st.subheader("📝 Your Input Text")
         st.text_area("Content", notes, height=300)
 
+elif input_type == "Image":
+    uploaded_img = st.file_uploader("Upload a PNG image", type=["png"])
+    if uploaded_img is not None:
+        img = Image.open(uploaded_img).convert("RGB")
+        st.image(img, caption="🖼️ Uploaded Image", use_column_width=True)
+
+        # preprocess & inference
+        inputs = processor(images=img, return_tensors="pt")
+        with torch.no_grad():
+            outputs = model(**inputs)
+        probs = torch.softmax(outputs.logits, dim=-1)[0]
+        top5 = torch.topk(probs, k=5)
+
+        st.subheader("🔍 Top 5 Predictions")
+        for idx, score in zip(top5.indices.tolist(), top5.values.tolist()):
+            label = model.config.id2label[idx]
+            st.write(f"- **{label}**: {score*100:.1f}%")
+
 else:
     st.info("Please select an input type to get started.")
+