Update app.py

app.py

@@ -266,44 +266,45 @@ def process_image_detection(image, target_label, surprise_rating):
                 )
 
         plt.axis('off')
-
-        print("Saving final image...")
-        try:
-            # Force figure to be rendered
-            fig.canvas.draw()
-            
-            # Get the RGBA buffer from the figure
-            w, h = fig.canvas.get_width_height()
-            buf = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
-            buf.shape = (h, w, 3)
-            
-            # Create PIL Image from buffer
-            output_image = Image.fromarray(buf)
-            
-            # Resize to original size if needed
-            if output_image.size != original_size:
-                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)}")
-            print(f"Figure type: {type(fig)}")
-            print(f"Canvas type: {type(fig.canvas)}")
-            raise
-
-    except Exception as e:
-        print(f"Process image detection error: {str(e)}")
-        print(f"Error occurred at line {e.__traceback__.tb_lineno}")
-        raise
+        
+                print("Saving final image...")
+                try:
+                    # Save directly to buffer using savefig
+                    buf = io.BytesIO()
+                    fig.savefig(buf, 
+                               format='png',
+                               dpi=dpi,
+                               bbox_inches='tight',
+                               pad_inches=0)
+                    buf.seek(0)
+                    
+                    # Open as PIL Image
+                    output_image = Image.open(buf)
+                    
+                    # Convert to RGB if needed
+                    if output_image.mode != 'RGB':
+                        output_image = output_image.convert('RGB')
+                    
+                    # Resize to original size if needed
+                    if output_image.size != original_size:
+                        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)
+                    buf.close()
+                    
+                    return final_buf
+        
+                except Exception as e:
+                    print(f"Save error details: {str(e)}")
+                    print(f"Figure type: {type(fig)}")
+                    print(f"Canvas type: {type(fig.canvas)}")
+                    raise
 
 def process_and_analyze(image):
     if image is None: