Update app.py

app.py

@@ -1,49 +1,31 @@
 import gradio as gr
 from openai import OpenAI
-import re
+import os
 
-def get_openrouter_client(api_key):
-    """Initialize OpenRouter client with user-provided API key"""
-    if not api_key or api_key.strip() == "":
-        return None, "Please enter your OpenRouter API key"
+def extract_medicine_names(api_key, image):
+    """Extract medicine names from a prescription image using OpenRouter and Gemini."""
+    if not api_key.strip():
+        return "Error: Please provide a valid OpenRouter API key."
     
     try:
         client = OpenAI(
             base_url="https://openrouter.ai/api/v1",
-            api_key=api_key
+            api_key=api_key,
         )
-        return client, None
-    except Exception as e:
-        return None, f"Error initializing client: {str(e)}"
-
-def extract_medicine_names(image, api_key):
-    """Extract medicine names from a prescription image using Gemini via OpenRouter"""
-    if not image:
-        return "Please upload a prescription image."
-    
-    # Get client with user-provided API key
-    client, error = get_openrouter_client(api_key)
-    if error:
-        return error
-    
-    try:
-        response = client.chat.completions.create(
+        
+        completion = client.chat.completions.create(
             extra_headers={
-                "HTTP-Referer": "https://medicine-extractor-app.com",
+                "HTTP-Referer": "gradio-medicine-extractor-app",
                 "X-Title": "Medicine Name Extractor",
             },
             model="google/gemini-2.0-flash-exp:free",
             messages=[
-                {
-                    "role": "system",
-                    "content": "You are an AI specialized in extracting medication names from prescription images. Only list the medication names, nothing else."
-                },
                 {
                     "role": "user",
                     "content": [
                         {
                             "type": "text",
-                            "text": "Extract ONLY the names of medications from this prescription image. Provide them as a numbered list. If this isn't a medical prescription, respond with 'No prescription detected'."
+                            "text": "This is a medical prescription image. Please analyze it and ONLY extract the medicine names. Return just a bulleted list of medicine names found, nothing else. If you can't identify any medicines or this isn't a prescription, please respond with 'No medicine names detected'."
                         },
                         {
                             "type": "image_url",
@@ -53,72 +35,50 @@ def extract_medicine_names(image, api_key):
                         }
                     ]
                 }
-            ],
-            max_tokens=300
+            ]
         )
         
-        result = response.choices[0].message.content.strip()
-        
-        # Check if no prescription was detected
-        if "No prescription detected" in result:
-            return "No prescription detected in the image."
-        
-        # Clean up the response to just include the medication names
-        # Remove any explanatory text that might appear before or after the list
-        medicines = []
-        for line in result.split('\n'):
-            # Look for numbered lines or lines starting with medication names
-            if re.match(r'^\d+\.', line.strip()):
-                # Extract text after the number and period
-                med_name = re.sub(r'^\d+\.\s*', '', line.strip())
-                medicines.append(med_name)
-        
-        if not medicines:
-            # If numbered list processing didn't work, return the raw output
-            return result
-        
-        return "\n".join([f"{i+1}. {med}" for i, med in enumerate(medicines)])
-        
+        return completion.choices[0].message.content
+    
     except Exception as e:
         return f"Error: {str(e)}"
 
 # Create the Gradio interface
-with gr.Blocks(title="Prescription Medicine Extractor") as app:
-    gr.Markdown("# Prescription Medicine Name Extractor")
-    gr.Markdown("Upload a prescription image to extract medication names.")
-    
-    api_key = gr.Textbox(
-        label="OpenRouter API Key", 
-        placeholder="Enter your OpenRouter API key here", 
-        type="password"
-    )
+with gr.Blocks(title="Medicine Name Extractor") as app:
+    gr.Markdown("# Medicine Name Extractor from Prescriptions")
+    gr.Markdown("Upload a prescription image to extract medicine names.")
     
     with gr.Row():
         with gr.Column():
-            image_input = gr.Image(type="filepath", label="Upload Prescription Image")
-            submit_btn = gr.Button("Extract Medicine Names", variant="primary")
+            api_key = gr.Textbox(
+                label="OpenRouter API Key", 
+                placeholder="Enter your OpenRouter API key", 
+                type="password"
+            )
+            img_input = gr.Image(
+                label="Upload Prescription Image",
+                type="filepath"
+            )
+            extract_btn = gr.Button("Extract Medicine Names")
         
         with gr.Column():
             output = gr.Textbox(label="Extracted Medicine Names", lines=10)
     
-    submit_btn.click(
-        fn=extract_medicine_names, 
-        inputs=[image_input, api_key], 
-        outputs=[output]
+    extract_btn.click(
+        fn=extract_medicine_names,
+        inputs=[api_key, img_input],
+        outputs=output
     )
     
     gr.Markdown("""
-    ## Usage Instructions
-    1. Enter your OpenRouter API key (get one from https://openrouter.ai)
-    2. Upload a clear image of a medical prescription
-    3. Click the "Extract Medicine Names" button
-    4. The names of medications will be displayed in the output box
+    ## How to use:
+    1. Enter your OpenRouter API key
+    2. Upload a prescription image
+    3. Click 'Extract Medicine Names'
     
-    **Note:** For best results, ensure the image is clear and the text is readable.
-    **Privacy Notice:** Your API key and images are processed only during the active session and are not stored.
+    The app will process the image and extract only the medicine names from the prescription.
     """)
 
 # Launch the app
 if __name__ == "__main__":
-    print("Starting Prescription Medicine Name Extractor application...")
     app.launch()