Update app.py

app.py

@@ -1,24 +1,33 @@
 import streamlit as st
+import cv2
+import numpy as np
 from PIL import Image
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel, pipeline
 import re
 
-# Load TrOCR Model for Handwritten OCR
+# Install TrOCR Model for Handwritten OCR
 processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
 model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
 
 # Load pre-trained QA model
 qa_pipeline = pipeline("question-answering", model="distilbert-base-uncased-distilled-squad")
 
-# Function to extract text using TrOCR
-def extract_text_from_handwriting(image_file):
-    image = Image.open(image_file).convert("RGB")
+# Preprocess image for better OCR performance
+def preprocess_image(image_file):
+    image = np.array(Image.open(image_file).convert("RGB"))
+    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)  # Convert to grayscale
+    blurred = cv2.GaussianBlur(gray, (5, 5), 0)  # Remove noise
+    thresh = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)[1]  # Increase contrast
+    return Image.fromarray(thresh)  # Convert back to PIL format
+
+# Extract text using TrOCR
+def extract_text_from_handwriting(image):
     pixel_values = processor(images=image, return_tensors="pt").pixel_values
     generated_ids = model.generate(pixel_values)
     extracted_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
     return extracted_text
 
-# Extract student name and roll number
+# Extract student info
 def extract_student_info(text):
     name = re.search(r"NAME\s*=\s*([\w\s]+)", text, re.IGNORECASE)
     roll_no = re.search(r"Roll\s*NO\s*=\s*(\d+)", text, re.IGNORECASE)
@@ -26,12 +35,12 @@ def extract_student_info(text):
     roll_number = roll_no.group(1).strip() if roll_no else "Unknown"
     return student_name, roll_number
 
-# Extract questions from the text
+# Extract questions
 def extract_questions_from_text(text):
-    questions = re.findall(r'(?:[^\n]*\?)', text)  # Extract sentences ending with "?"
+    questions = re.findall(r'(?:[^\n]*\?)', text)
     return questions
 
-# Grading function using QA model
+# Grade answers
 def grade_answer(question, context):
     result = qa_pipeline(question=question, context=context)
     return result['score'], "Correct" if result['score'] > 0.5 else "Incorrect"
@@ -44,10 +53,13 @@ st.write("Upload an image or handwritten file to process.")
 uploaded_image = st.file_uploader("Upload Handwritten Image", type=["png", "jpg", "jpeg"])
 
 if uploaded_image:
-    st.image(uploaded_image, caption="Uploaded Handwritten File", use_container_width=True)
+    # Preprocess the image
+    st.image(uploaded_image, caption="Original Image", use_container_width=True)
+    preprocessed_image = preprocess_image(uploaded_image)
+    st.image(preprocessed_image, caption="Preprocessed Image", use_container_width=True)
 
     # Extract text using TrOCR
-    extracted_text = extract_text_from_handwriting(uploaded_image)
+    extracted_text = extract_text_from_handwriting(preprocessed_image)
     st.subheader("Extracted Text")
     st.text(extracted_text)