Update app.py

app.py

@@ -10,13 +10,27 @@ from translate import Translator
 from langchain.prompts import PromptTemplate
 from langchain.chains import LLMChain
 from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
-from transformers import pipeline
 
 # Set API keys
 os.environ["HUGGINGFACEHUB_API_KEY"] = os.getenv("HF")
 os.environ["HF_TOKEN"] = os.getenv("HF")
 
-# Function to save text as an image
+# Split large response into smaller chunks (for translation)
+def split_text_into_chunks(text, max_length=450):
+    lines = text.split('\n')
+    chunks = []
+    current = ""
+    for line in lines:
+        if len(current) + len(line) + 1 <= max_length:
+            current += line + '\n'
+        else:
+            chunks.append(current.strip())
+            current = line + '\n'
+    if current:
+        chunks.append(current.strip())
+    return chunks
+
+# Save text to image
 def save_text_as_image(text, file_path):
     font = ImageFont.load_default()
     lines = text.split('\n')
@@ -34,7 +48,7 @@ def save_text_as_image(text, file_path):
     img.save(file_path)
     return file_path
 
-# Setup
+# Setup UI
 st.set_page_config(page_title="MediAssist 💊", layout="wide")
 
 st.markdown("""
@@ -59,7 +73,7 @@ if uploaded_file:
         temp_file.write(uploaded_file.read())
         orig_path = temp_file.name
 
-    # Preprocess Image
+    # Image preprocessing
     image = cv2.imread(orig_path)
     gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
     _, binary_inv = cv2.threshold(gray, 128, 255, cv2.THRESH_BINARY_INV)
@@ -74,16 +88,15 @@ if uploaded_file:
     text_list = reader.readtext(dilated, detail=0)
     text = "\n".join(text_list)
 
-    # Display
     col1, col2 = st.columns([1, 2])
     with col1:
-        st.image(dilated, caption="Preprocessed Prescription", channels="GRAY", use_container_width=True)
+        st.image(dilated, caption="🧾 Preprocessed Prescription", channels="GRAY", use_container_width=True)
     with col2:
         st.success("✅ Image Uploaded and Preprocessed")
-        st.markdown("#### 📝 Extracted Text from Image")
+        st.markdown("#### 📝 Extracted Text")
         st.code(text)
 
-    # Prompt
+    # Prompt LLM
     template = """
     You are a helpful medical assistant.
     Here is a prescription text extracted from an image:
@@ -117,51 +130,69 @@ if uploaded_file:
     
     chain = LLMChain(llm=llm, prompt=prompt)
 
+    response = ""
+    hindi_text = ""
+
     if st.button("🔍 Analyze Extracted Text"):
         with st.spinner("Analyzing with LLM..."):
             response = chain.run(prescription_text=text)
-        st.markdown("#### 💡 Analyzed Medicine Info")
-        st.text_area("Output", response, height=300)
+        st.markdown("#### 💡 AI-based Medicine Analysis")
+        st.text_area("LLM Output", response, height=300)
 
         # Save txt and image
-        txt_path = "prescription_output.txt"
+        txt_path = "medicine_analysis.txt"
         with open(txt_path, "w") as f:
             f.write(response)
 
-        img_path = "prescription_output.png"
+        img_path = "medicine_analysis.png"
         save_text_as_image(response, img_path)
 
-        # Target language code (like 'hi' for Hindi, 'mr' for Marathi, 'gu' for Gujarati)
-        target_lang = "hi"  
-        
-        translator = Translator(to_lang=target_lang)
-        hindi_text = translator.translate(response)
-
-        # # Translation to Hindi
-        # translator = pipeline("translation", model="Helsinki-NLP/opus-mt-en-hi")
-        # hindi_text = translator(response, max_length=400)[0]['translation_text']
-
-        st.markdown("#### 🌐 Translate to Hindi")
-        st.text_area("Translated (Hindi)", hindi_text, height=300)
-
-        st.markdown("#### 📥 Download Options")
-        colA, colB, colC, colD = st.columns(4)
-        with colA:
-            st.download_button("⬇️ Download TXT", data=response, file_name="medicine_analysis.txt")
-        with colB:
+        st.markdown("#### 📥 Download (English)")
+        col1, col2 = st.columns(2)
+        with col1:
+            st.download_button("⬇️ English TXT", data=response.encode(), file_name="medicine_analysis.txt")
+        with col2:
             with open(img_path, "rb") as img_file:
-                st.download_button("🖼️ Download Image", data=img_file, file_name="medicine_analysis.png", mime="image/png")
-        with colC:
-            st.download_button("⬇️ Hindi TXT", data=hindi_text, file_name="hindi_medicine_analysis.txt")
-        with colD:
-            hindi_img_path = "hindi_output.png"
-            save_text_as_image(hindi_text, hindi_img_path)
-            with open(hindi_img_path, "rb") as hindi_img_file:
-                st.download_button("🖼️ Hindi Image", data=hindi_img_file, file_name="hindi_output.png", mime="image/png")
+                st.download_button("🖼️ English Image", data=img_file, file_name="medicine_analysis.png", mime="image/png")
+
+    if response and st.button("🌐 Translate to Hindi"):
+        with st.spinner("Translating to Hindi..."):
+            target_lang = "hi"
+            translator = Translator(to_lang=target_lang)
+            chunks = split_text_into_chunks(response, max_length=450)
+
+            hindi_chunks = []
+            for chunk in chunks:
+                try:
+                    translated = translator.translate(chunk)
+                    hindi_chunks.append(translated)
+                except Exception as e:
+                    hindi_chunks.append("[Error translating chunk]")
+
+            hindi_text = "\n\n".join(hindi_chunks)
+
+        st.markdown("#### 🌐 Hindi Translation")
+        st.text_area("Translated Output (Hindi)", hindi_text, height=300)
+
+        hindi_img_path = "hindi_output.png"
+        save_text_as_image(hindi_text, hindi_img_path)
+
+        st.markdown("#### 📥 Download (Hindi)")
+        col3, col4 = st.columns(2)
+        with col3:
+            st.download_button("⬇️ Hindi TXT", data=hindi_text.encode(), file_name="hindi_medicine_analysis.txt")
+        with col4:
+            with open(hindi_img_path, "rb") as img_file:
+                st.download_button("🖼️ Hindi Image", data=img_file, file_name="hindi_medicine_analysis.png", mime="image/png")
 
     # Cleanup
-    os.remove(orig_path)
-    os.remove(dilated_path)
+    try:
+        os.remove(orig_path)
+        os.remove(dilated_path)
+        llm.clear()
+    except:
+        pass
+
 else:
     st.markdown("<center><i>📸 Upload a prescription image to get started</i></center>", unsafe_allow_html=True)