Update app.py

app.py

@@ -1,31 +1,36 @@
 import os
 import streamlit as st
-from transformers import AutoProcessor, AutoModelForPreTraining 
-import pytesseract
-import cv2
-import numpy as np
+import requests
 from PIL import Image
+import numpy as np
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 from huggingface_hub import login
 
 # Get the token from environment variables (set in Hugging Face Space)
 token = os.getenv("HF_Token")
 login(token)
 
-# Load the model
-processor = AutoProcessor.from_pretrained("google/paligemma-3b-mix-224")
-model = AutoModelForPreTraining.from_pretrained("google/paligemma-3b-mix-224")
-
-# Function to extract text from image using OCR
-def extract_text_from_image(image_file):
-    image = Image.open(image_file)
-    img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
-    text = pytesseract.image_to_string(img_cv)
-    return text
-
-# Function to get AI interpretation of the prescription
-def interpret_prescription(text):
-    response = model(text)
-    return response[0]['generated_text'].strip()
+# Initialize the Hugging Face model
+model_name = "google/paligemma-3b-mix-224"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
+# Function to transcribe handwritten notes using Hugging Face model
+def transcribe_handwriting(image):
+    # Convert image to array and preprocess
+    image = image.convert("RGB")
+    image = np.array(image)
+
+    # Prepare input for the model
+    inputs = tokenizer(image, return_tensors="pt")
+    
+    # Generate output
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_length=512)
+    
+    transcription = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return transcription
 
 # Set Streamlit page configuration
 st.set_page_config(
@@ -44,15 +49,8 @@ if uploaded_file is not None:
     # Display uploaded image
     st.image(uploaded_file, caption="Uploaded Prescription", use_column_width=True)
 
-    with st.spinner("Extracting text from prescription..."):
-        # Extract text from image using OCR
-        extracted_text = extract_text_from_image(uploaded_file)
-        st.subheader("Extracted Text from Prescription:")
-        st.text(extracted_text)
-
-        if extracted_text:
-            # Interpret extracted text using the model
-            with st.spinner("Interpreting the prescription..."):
-                ai_response = interpret_prescription(extracted_text)
-                st.subheader("AI Interpretation:")
-                st.text(ai_response)
+    with st.spinner("Transcribing handwriting..."):
+        # Transcribe handwritten notes
+        extracted_text = transcribe_handwriting(uploaded_file)
+        st.subheader("Transcribed Text from Prescription:")
+        st.text(extracted_text)