Update app.py

app.py

@@ -266,31 +266,40 @@ def process_image_detection(image, target_label, surprise_rating):
                 )
 
         plt.axis('off')
-
+    
         print("Saving final image...")  # Debug print
-        buf = io.BytesIO()
-        plt.savefig(buf,
-                    format='png',
-                    dpi=dpi,
-                    bbox_inches='tight',
-                    pad_inches=0,
-                    metadata={'dpi': original_dpi})
-        buf.seek(0)
-        plt.close()
-
-        output_image = Image.open(buf)
-        output_image = output_image.resize(original_size, Image.Resampling.LANCZOS)
-
-        final_buf = io.BytesIO()
-        output_image.save(final_buf, format='PNG', dpi=original_dpi)
-        final_buf.seek(0)
-
-        return final_buf
-
-    except Exception as e:
-        print(f"Process image detection error: {str(e)}")  # Debug print
-        print(f"Error occurred at line {e.__traceback__.tb_lineno}")  # Debug print
-        raise
+        try:
+            # Save to buffer
+            buf = io.BytesIO()
+            
+            # Force figure to be in a format we can save
+            fig.canvas.draw()
+            
+            # Get the image data from the figure
+            plot_data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+            plot_data = plot_data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+            
+            # Convert to PIL Image
+            output_image = Image.fromarray(plot_data)
+            
+            # Resize if needed
+            output_image = output_image.resize(original_size, Image.Resampling.LANCZOS)
+            
+            # Save to final buffer
+            final_buf = io.BytesIO()
+            output_image.save(final_buf, format='PNG', dpi=original_dpi)
+            final_buf.seek(0)
+            
+            # Cleanup
+            plt.close(fig)
+            
+            return final_buf
+            
+        except Exception as e:
+            print(f"Save error details: {str(e)}")  # Debug print
+            print(f"Figure type: {type(fig)}")
+            print(f"Canvas type: {type(fig.canvas)}")
+            raise
 
 
 def process_and_analyze(image):