Update app.py

app.py

@@ -45,14 +45,19 @@ def analyze_image(image, prompt):
         with torch.no_grad():
             output = model.generate(**inputs, max_new_tokens=100)
         print(processor.decode(output[0]))
-        result = processor.decode(output[0], skip_special_tokens=True)
+        full_response = processor.decode(output[0])
+    
         
-        try:
-            return json.loads(result)
-        except json.JSONDecodeError:
-            return {"error": "Failed to parse model output as JSON", "raw_output": result}
-    except Exception as e:
-        return {"error": str(e), "traceback": traceback.format_exc()}
+        json_match = re.search(r'\{.*\}', full_response, re.DOTALL)
+        if json_match:
+            json_str = json_match.group(0)
+            try:
+                return json.loads(json_str)
+            except json.JSONDecodeError:
+                return {"error": "Failed to parse model output as JSON", "raw_output": json_str}
+        else:
+            return {"error": "No JSON found in the response", "raw_output": full_response}
+
 
 default_prompt = """Analyze this image and determine if it contains a data logger. A data logger is typically a small, black electronic device used to monitor and record data over time, such as voltage, temperature, or current, via external sensors.