Spaces:
Sleeping
Sleeping
import streamlit as st | |
import os | |
from PIL import Image | |
import google.generativeai as genai | |
from dotenv import load_dotenv | |
from mysrap import search_medicine_supertails | |
from google.generativeai.types import HarmCategory, HarmBlockThreshold | |
def presc_analyze(): | |
# Load environment variables | |
load_dotenv() | |
# Configure the API key | |
genai.configure(api_key=os.getenv("GOOGLE_API_KEY")) | |
# Function to load Gemini Vision Pro model and get response | |
def get_gemini_response(input_prompt, image_data, user_prompt): | |
# model = genai.GenerativeModel("gemini-pro-vision") | |
safety_settings={ | |
HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE, | |
HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE, | |
# HarmCategory.HARM_CATEGORY_DANGEROUS: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE, | |
HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE, | |
} | |
model = genai.GenerativeModel("gemini-pro-vision",safety_settings=safety_settings) | |
response = model.generate_content([input_prompt, image_data[0], user_prompt]) | |
return response.text | |
# Function to extract data from uploaded image | |
def input_image_setup(uploaded_file): | |
if uploaded_file is not None: | |
bytes_data = uploaded_file.getvalue() | |
image_parts = [ | |
{ | |
"mime_type": uploaded_file.type, | |
"data": bytes_data, | |
} | |
] | |
return image_parts | |
else: | |
raise FileNotFoundError("No file uploaded") | |
def extract_medicine_names_heuristic(text): | |
import re | |
# Split based on common delimiters, numbers, or new lines | |
potential_medicines = re.split(r'\d+\.\s*|\n|\r|,', text) | |
# Clean up and filter out non-medicines | |
potential_medicines = [med.strip() for med in potential_medicines if med.strip() and med.lower() not in ["the", "to", "of", "and", "is"]] | |
return potential_medicines | |
# Streamlit app configuration | |
# st.set_page_config(page_title="Prescription & Medicine Information Extractor") | |
st.header("Prescription & Medicine Information Extractor") | |
user_prompt = "Give me the details about the Prescription and along with basic personal details and tell about medications and frequency in a proper tabular form" | |
uploaded_file = st.file_uploader("Choose a Prescription Image...", type=["jpg", "jpeg", "png"]) | |
image = "" | |
if uploaded_file is not None: | |
image = Image.open(uploaded_file) | |
st.image(image, caption="Uploaded Image", use_column_width=True) | |
submit = st.button("Extract Information") | |
# System prompt for understanding prescriptions | |
info_prompt = """You are an expert in understanding prescriptions. | |
You will receive prescription images and answer questions based on them. | |
Please consider the following information: {user_prompt}""" | |
# System prompt for extracting medicine names (heuristic) | |
medicine_prompt = "Please list all the medications mentioned in the prescription.Make sure that you just give medicine name with strength." | |
if submit: | |
image_data = input_image_setup(uploaded_file) | |
# Extract information based on user prompt | |
info_response = get_gemini_response(info_prompt.format(user_prompt=user_prompt), image_data, "") | |
st.subheader("Extracted Information:") | |
st.write(info_response) | |
# Extract medicine names using a separate query | |
medicine_response = get_gemini_response(medicine_prompt, image_data, "") | |
medicine_names = extract_medicine_names_heuristic(medicine_response) | |
if medicine_names: | |
st.subheader("Medicine Names:") | |
# Store medicine names in a list | |
medicine_list = ", ".join(medicine_names) | |
st.write(medicine_list) | |
# Display at least 6-7 medicines with clickable links for each extracted medicine name | |
st.subheader("Medicines with Links:") | |
for medicine in medicine_names: | |
st.markdown(f"### {medicine}") | |
results_df = search_medicine_supertails(medicine) | |
if results_df is not None and not results_df.empty: | |
for index, row in results_df.head(7).iterrows(): | |
st.markdown(f"[{row['Product Name']}]({row['Link']}) - {row['Price']} ₹") | |
else: | |
st.write(f"No results found for {medicine}") | |
else: | |
st.write("No medicine names found using the heuristic approach.") | |
st.balloons() | |
if __name__ == '__main__': | |
presc_analyze() | |