Update app.py

app.py

@@ -26,16 +26,16 @@ 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"
+OCR_MODEL = "microsoft/trocr-base-printed"
 
-@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()
+@st.cache_resource
+def load_ocr_model():
+    processor = TrOCRProcessor.from_pretrained(OCR_MODEL)
+    model     = VisionEncoderDecoderModel.from_pretrained(OCR_MODEL)
+    return processor, model
 
+ocr_processor, ocr_model = load_ocr_model()
 # --- Main UI ---
 input_type = st.selectbox(
     "Select the type of input:",
@@ -68,23 +68,23 @@ elif input_type == "Text":
         st.text_area("Content", notes, height=300)
 
 elif input_type == "Image":
-    uploaded_img = st.file_uploader("Upload a PNG image", type=["png"])
+    uploaded_img = st.file_uploader("Upload a PNG/JPG image", type=["png", "jpg", "jpeg"])
     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")
+        # 1. Preprocess for OCR
+        pixel_values = ocr_processor(images=img, return_tensors="pt").pixel_values
+
+        # 2. Generate and decode
         with torch.no_grad():
-            outputs = model(**inputs)
-        probs = torch.softmax(outputs.logits, dim=-1)[0]
-        top5 = torch.topk(probs, k=5)
+            generated_ids = ocr_model.generate(pixel_values)
+        extracted_text = ocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-        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}%")
+        st.subheader("🖋️ Extracted Text from Image")
+        st.text_area("OCR Result", extracted_text, height=300)
 
 else:
     st.info("Please select an input type to get started.")
 
+