Update app.py

app.py

@@ -81,24 +81,33 @@ def submit_query(session_id, query):
 
 
 def extract_json_from_answer(answer, image_analysis):
-    """Extract and clean JSON from the LLM response"""
+    """Extract and clean JSON from the LLM response and append image analysis results."""
     try:
-        # First try to parse the answer directly
-
-        return json.loads(answer)
+        # Try to parse the JSON answer directly
+        json_data = json.loads(answer)
     except json.JSONDecodeError:
-
         try:
             # If that fails, try to find JSON content and parse it
             start_idx = answer.find('{')
-            end_idx = answer.rfind('}')
+            end_idx = answer.rfind('}') + 1
             if start_idx != -1 and end_idx != 0:
                 json_str = answer[start_idx:end_idx]
-                json_str += f",'image_analysis': ['prediction': {image_analysis['prediction']}, 'confidence': {image_analysis['confidence']}%.]"
-                return json.loads(json_str)
-        except (json.JSONDecodeError, ValueError):
-            logger.error("Failed to parse JSON from response")
+                json_data = json.loads(json_str)
+            else:
+                raise ValueError("Failed to locate JSON in the answer")
+        except (json.JSONDecodeError, ValueError) as e:
+            logger.error(f"Failed to parse JSON from response: {str(e)}")
             raise
+    
+    # Append the image analysis data
+    if image_analysis:
+        json_data["image_analysis"] = {
+            "prediction": image_analysis["prediction"],
+            "confidence": f"{image_analysis['confidence']:.2f}%"  # Format confidence as percentage
+        }
+
+    return json_data
+
 
 def load_model():
     try:
@@ -194,8 +203,8 @@ def gradio_interface(patient_info, image):
             logger.info(f"Prediction: {prediction}")
             # Format prediction results
             image_analysis = {
-                "prediction": classes[int(prediction[0][0])],
-                "confidence": float(prediction[0][0]) * 100
+                "prediction": classes[np.argmax(prediction[0])],
+                "confidence": np.max(predicition[0]) * 100
             }
             logger.info(f"Image analysis results: {image_analysis}")