adjust output for disease explanation

app.py

@@ -5,48 +5,70 @@ from PIL import Image
 import google.generativeai as genai 
 import os
 
+# Load TensorFlow model
 model_path = 'model'
 model = tf.saved_model.load(model_path)
 
+# Set up Gemini API
 api_key = os.getenv("GEMINI_API_KEY")
 genai.configure(api_key=api_key)
 
+# Labels for classification
 labels = ['cataract', 'diabetic_retinopathy', 'glaucoma', 'normal']
 
 def get_disease_detail(disease_name):
-  prompt = (
-  f"Diagnosis: {disease_name}\n\n"
-  "What is it?\n(Description about {disease_name})\n\n"
-  "What cause it?\n(Explain what causes {disease_name})\n\n"
-  "Suggestion\n(Suggestion to user)\n\n"
-  "Reminder: Always seek professional help, such as a doctor."
-  )
-  response = genai.GenerativeModel("gemini-1.5-flash").generate_content(prompt)
-  return response.text.strip()
+    prompt = (
+        f"Diagnosis: {disease_name}\n\n"
+        "What is it?\n(Description about the disease)\n\n"
+        "What causes it?\n(Explain what causes the disease)\n\n"
+        "Suggestions\n(Suggestion to user)\n\n"
+        "Reminder: Always seek professional help, such as a doctor."
+    )
+    response = genai.GenerativeModel("gemini-1.5-flash").generate_content(prompt)
+    return response.text.strip()
+
+def safe_extract_section(text, start_keyword, end_keyword):
+    """ Safely extract sections from the Gemini response based on start and end keywords."""
+    if start_keyword in text and end_keyword in text:
+        return text.split(start_keyword)[1].split(end_keyword)[0].strip()
+    elif start_keyword in text:
+        return text.split(start_keyword)[1].strip()
+    else:
+        return "Information not available."
 
 def predict_image(image):
-  image_resized = image.resize((224, 224))
-  image_array = np.array(image_resized).astype(np.float32) / 255.0
-  image_array = np.expand_dims(image_array, axis=0)
+    # Preprocess the image
+    image_resized = image.resize((224, 224))
+    image_array = np.array(image_resized).astype(np.float32) / 255.0
+    image_array = np.expand_dims(image_array, axis=0)
+
+    # Run prediction
+    predictions = model.signatures['serving_default'](tf.convert_to_tensor(image_array, dtype=tf.float32))['output_0']
+    top_index = np.argmax(predictions.numpy(), axis=1)[0]
+    top_label = labels[top_index]
+    top_probability = predictions.numpy()[0][top_index] * 100  # Convert to percentage
 
-  predictions = model.signatures['serving_default'](tf.convert_to_tensor(image_array, dtype=tf.float32))['output_0']
-    
-  # Highest prediction
-  top_index = np.argmax(predictions.numpy(), axis=1)[0]
-  top_label = labels[top_index]
-  top_probability = predictions.numpy()[0][top_index]
+    # Get explanation from Gemini API
+    explanation = get_disease_detail(top_label)
 
-  explanation = get_disease_detail(top_label)
+    # Extract relevant sections from the explanation
+    diagnosis_section = f"**Diagnosis:** {top_label}"
+    what_is_it = safe_extract_section(explanation, "What is it?", "What causes it?")
+    causes = safe_extract_section(explanation, "What causes it?", "Suggestions")
+    suggestions = safe_extract_section(explanation, "Suggestions", "Reminder")
+    reminder = "Always seek professional help, such as a doctor."
 
-  formatted_explanation = (
-    f"**Diagnosis:** {top_label}\n\n"
-        f"**What is it?** {explanation.split('What causes it?')[0].split('Suggestions')[0].strip()}\n\n"
-        f"**What causes it?** {explanation.split('What causes it?')[1].split('Suggestions')[0].strip()}\n\n"
-        f"**Suggestions:** {explanation.split('Suggestions')[1].split('Reminder')[0].strip()}\n\n"
-        f"**Reminder:** Always seek professional help, such as a doctor."
-  )
+    # Format explanation
+    formatted_explanation = (
+        f"{diagnosis_section}\n\n"
+        f"**What is it?** {what_is_it}\n\n"
+        f"**What causes it?** {causes}\n\n"
+        f"**Suggestions:** {suggestions}\n\n"
+        f"**Reminder:** {reminder}"
+    )
 
-  return {top_label: top_probability}, formatted_explanation  
+    # Return both the prediction and the explanation
+    return {top_label: top_probability}, formatted_explanation
 
 # Example images
 example_images = [
@@ -63,13 +85,14 @@ interface = gr.Interface(
     fn=predict_image,
     inputs=gr.Image(type="pil"),
     outputs=[
-      gr.Label(num_top_classes=1, label="Prediction"), 
-      gr.Textbox(label="Explanation")],
+        gr.Label(num_top_classes=1, label="Prediction"), 
+        gr.Textbox(label="Explanation")
+    ],
     examples=example_images,
     title="Eye Diseases Classifier",
     description=(
-      "Upload an image of an eye fundus, and the model will predict it.\n\n"
-      "**Disclaimer:** This model is intended as a form of learning process in the field of health-related machine learning and was trained with a limited amount and variety of data with a total of about 4000 data, so the prediction results may not always be correct. There is still a lot of room for improvisation on this model in the future."
+        "Upload an image of an eye fundus, and the model will predict it.\n\n"
+        "**Disclaimer:** This model is intended as a form of learning process in the field of health-related machine learning and was trained with a limited amount and variety of data with a total of about 4000 data, so the prediction results may not always be correct. There is still a lot of room for improvisation on this model in the future."
     ),
     allow_flagging="never"
 )