Update pages/model.py

pages/model.py

@@ -18,62 +18,90 @@ st.set_page_config(page_title="🧠 MediAssist", layout="centered")
 st.markdown("<h1 style='text-align: center; color: #4A90E2;'>🧠 MediAssist</h1>", unsafe_allow_html=True)
 st.markdown("<h4 style='text-align: center;'>Upload a doctor's prescription and get detailed medicine analysis</h4><br>", unsafe_allow_html=True)
 
-uploaded_file = st.file_uploader("📤 Upload Prescription Image", type=["jpg", "jpeg", "png"])
+uploaded_file = st.file_uploader("📤 Upload Prescription Image (JPG/PNG)", type=["jpg", "jpeg", "png"])
 
 if uploaded_file:
-
+    # Save uploaded image temporarily
     with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file:
         temp_file.write(uploaded_file.read())
-        img_path = temp_file.name
+        orig_path = temp_file.name
 
-    image = cv2.imread(img_path)
+    # Step 1: Preprocess the image
+    image = cv2.imread(orig_path)
     gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    _, binary = cv2.threshold(gray, 130, 255, cv2.THRESH_BINARY_INV)
-    dilated = cv2.dilate(binary, np.ones((2, 2), np.uint8), iterations=1)
+    _, binary_inv = cv2.threshold(gray, 128, 255, cv2.THRESH_BINARY_INV)
+    kernel = np.ones((3, 3), np.uint8)
+    dilated = cv2.dilate(binary_inv, kernel, iterations=1)
 
-    reader = easyocr.Reader(['en'], gpu=False)
-    extracted_text = "\n".join(reader.readtext(dilated, detail=0))
+    # Save the dilated image temporarily for reference
+    dilated_path = orig_path.replace(".png", "_dilated.png")
+    cv2.imwrite(dilated_path, dilated)
 
-    st.image(dilated, caption="🧾 Preprocessed Image", use_column_width=True)
-    st.markdown("### 📜 Extracted Text from Image")
-    st.code(extracted_text)
+    # Step 2: OCR using EasyOCR
+    reader = easyocr.Reader(['en'])
+    text_list = reader.readtext(dilated, detail=0)
+    text = "\n".join(text_list)
 
+    # Step 3: Prompt for the LLM
     template = """
-    You're a medical assistant AI. Below is a doctor's handwritten prescription text:
+    You are a helpful medical assistant.
+
+    Here is a prescription text extracted from an image:
 
     {prescription_text}
 
-    Based on the text, please do the following:
-    1. Extract all medicine names (ignore other notes).
-    2. For each medicine, mention:
-       - When to take it (morning/night, before/after food)
-       - Dosage
+    Please do the following:
+
+    1. Extract only the medicine names mentioned in the prescription (ignore any other text).
+    2. For each medicine, provide:
+       - When to take it (timing and dosage)
        - Possible side effects
-       - Any precautions or special instructions
-    Return results in clear bullet points.
-    """
+       - Any special instructions
 
+    Format your answer as bullet points, listing only medicines and their details.
+    """
     prompt = PromptTemplate(input_variables=["prescription_text"], template=template)
 
-    llm = ChatHuggingFace(
-        llm=HuggingFaceEndpoint(
-            repo_id="meta-llama/Llama-3.1-8B-Instruct",
-            provider="nebius",
-            temperature=0.5,
-            max_new_tokens=500,
-            task="conversational"
-        )
+    # Step 4: Load LLM
+    llm_model = HuggingFaceEndpoint(
+        repo_id="meta-llama/Llama-3.1-8B-Instruct",
+        provider="nebius",
+        temperature=0.6,
+        max_new_tokens=300,
+        task="conversational"
     )
 
-    chain = LLMChain(llm=llm, prompt=prompt)
+    model = ChatHuggingFace(
+        llm=llm_model,
+        repo_id="meta-llama/Llama-3.1-8B-Instruct",
+        provider="nebius",
+        temperature=0.6,
+        max_new_tokens=300,
+        task="conversational"
+    )
+
+    chain = LLMChain(llm=model, prompt=prompt)
+
+    # Step 5: Layout for output
+    col1, col2 = st.columns([1, 2])
+
+    with col1:
+        st.image(dilated, caption="🧾 Preprocessed Prescription", channels="GRAY", use_container_width=True)
+
+    with col2:
+        st.success("✅ Prescription Uploaded & Preprocessed Successfully")
+        st.markdown("### 📜 Extracted Text")
+        st.code(text)
 
-    if st.button("🔍 Analyze Prescription"):
-        with st.spinner("Analyzing the prescription..."):
-            response = chain.run(prescription_text=extracted_text)
-        st.markdown("### 💊 Medicine Summary")
-        st.success(response)
+        if st.button("🔍 Analyze Text"):
+            with st.spinner("Analyzing with LLM..."):
+                response = chain.run(prescription_text=text)
+            st.markdown("### 💡 AI-Powered Summary")
+            st.success(response)
 
-    os.remove(img_path)
+    # Cleanup temp files
+    os.remove(orig_path)
+    os.remove(dilated_path)
 
 else:
-    st.info("Upload a prescription image to start the analysis.")
+    st.markdown("<center><i>Upload a prescription image to begin analysis.</i></center>", unsafe_allow_html=True)