Update src/streamlit_app.py

src/streamlit_app.py

@@ -3,6 +3,7 @@ import base64
 import requests
 import json
 import io
+import re
 from PIL import Image
 
 st.set_page_config(page_title="Solar Rooftop Analyzer", layout="centered")
@@ -19,17 +20,21 @@ def analyze_image_with_openrouter(image_file):
     buffer = io.BytesIO()
     img.save(buffer, format="JPEG")
     jpeg_bytes = buffer.getvalue()
+    
     # Base64 encode with content-type prefix
     encoded_image = "data:image/jpeg;base64," + base64.b64encode(jpeg_bytes).decode("utf-8")
+    
     prompt = (
-        "Analyze the rooftop in this image. Output JSON with: [Roof area (sqm), "
-        "Sunlight availability (%), Shading (Yes/No), Recommended solar panel type, "
-        "Estimated capacity (kW)]."
+        "Analyze the rooftop in this image. Output JSON with: "
+        "[Roof area (sqm), Sunlight availability (%), Shading (Yes/No), "
+        "Recommended solar panel type, Estimated capacity (kW)]."
     )
+    
     headers = {
         "Authorization": f"Bearer {OPENROUTER_API_KEY}",
         "Content-Type": "application/json"
     }
+    
     payload = {
         "model": VISION_MODEL_NAME,
         "messages": [
@@ -42,11 +47,23 @@ def analyze_image_with_openrouter(image_file):
             }
         ]
     }
+    
     response = requests.post("https://openrouter.ai/api/v1/chat/completions", json=payload, headers=headers)
+    
     if response.status_code == 200:
         return response.json()
+    
     return {"error": f"Failed to analyze image. Status code: {response.status_code}, Response: {response.text}"}
 
+def extract_json_from_response(content):
+    try:
+        match = re.search(r"\{.*\}", content, re.DOTALL)
+        if match:
+            return json.loads(match.group(0))
+    except Exception as e:
+        st.warning(f"Failed to parse JSON: {e}")
+    return None
+
 def estimate_roi(roof_area, capacity_kw, budget):
     cost_per_kw = 65000  # INR/kW
     estimated_cost = capacity_kw * cost_per_kw
@@ -74,23 +91,28 @@ if submitted:
         st.image(uploaded_file, caption="Uploaded Rooftop Image", use_column_width=True)
         with st.spinner("Analyzing rooftop image..."):
             ai_response = analyze_image_with_openrouter(uploaded_file)
+
         if "choices" in ai_response:
             try:
-                choice = ai_response["choices"][0]
-                # Some models may return content as a string, others as a dict
-                content = choice.get("message", {}).get("content", "")
-                content_json = json.loads(content)
-                st.success("Analysis complete!")
-                st.subheader("Rooftop Analysis")
-                st.json(content_json)
-                if "Roof area (sqm)" in content_json and "Estimated capacity (kW)" in content_json:
-                    roi = estimate_roi(
-                        roof_area=content_json["Roof area (sqm)"],
-                        capacity_kw=content_json["Estimated capacity (kW)"],
-                        budget=budget
-                    )
-                    st.subheader("ROI Estimation")
-                    st.json(roi)
+                content = ai_response["choices"][0]["message"]["content"]
+                content_json = extract_json_from_response(content)
+                
+                if content_json:
+                    st.success("Analysis complete!")
+                    st.subheader("Rooftop Analysis")
+                    st.json(content_json)
+
+                    if "Roof area (sqm)" in content_json and "Estimated capacity (kW)" in content_json:
+                        roi = estimate_roi(
+                            roof_area=content_json["Roof area (sqm)"],
+                            capacity_kw=content_json["Estimated capacity (kW)"],
+                            budget=budget
+                        )
+                        st.subheader("ROI Estimation")
+                        st.json(roi)
+                else:
+                    st.error("Could not extract structured data from the AI response.")
+                    st.text(content)  # Fallback: show raw content
             except Exception as e:
                 st.error(f"Error parsing analysis content: {e}")
                 st.json(ai_response)
@@ -98,3 +120,4 @@ if submitted:
             st.error("Failed to analyze the image. Please try again.")
     else:
         st.warning("Please upload an image and fill all fields.")
+