Create test_api.py

test_api.py

@@ -0,0 +1,66 @@
+import requests
+import sys
+import os
+
+def remove_background(image_path, output_path=None, api_url="http://localhost:7860/remove-bg"):
+    """
+    Remove the background of an image using the API.
+    
+    Args:
+        image_path (str): Path to the input image
+        output_path (str, optional): Path to save the output image. If None, will use input_name_nobg.png
+        api_url (str, optional): URL of the API endpoint
+    
+    Returns:
+        str: Path to the saved output image
+    """
+    # Check if image exists
+    if not os.path.exists(image_path):
+        print(f"Error: Image not found at {image_path}")
+        return None
+    
+    # Default output path if not specified
+    if output_path is None:
+        filename, _ = os.path.splitext(os.path.basename(image_path))
+        output_path = f"{filename}_nobg.png"
+    
+    # Prepare the file for upload
+    files = {"file": open(image_path, "rb")}
+    
+    print(f"Sending request to remove background from {image_path}...")
+    
+    try:
+        # Send the request
+        response = requests.post(api_url, files=files)
+        
+        # Check if request was successful
+        if response.status_code == 200:
+            # Save the output image
+            with open(output_path, "wb") as f:
+                f.write(response.content)
+            print(f"Background removed successfully! Saved to {output_path}")
+            return output_path
+        else:
+            print(f"Error: API returned status code {response.status_code}")
+            print(response.text)
+            return None
+    
+    except Exception as e:
+        print(f"Error: {str(e)}")
+        return None
+    
+    finally:
+        # Close the file
+        files["file"].close()
+
+if __name__ == "__main__":
+    # Check if image path is provided as command line argument
+    if len(sys.argv) < 2:
+        print("Usage: python test_api.py <image_path> [output_path] [api_url]")
+        sys.exit(1)
+    
+    image_path = sys.argv[1]
+    output_path = sys.argv[2] if len(sys.argv) > 2 else None
+    api_url = sys.argv[3] if len(sys.argv) > 3 else "http://localhost:7860/remove-bg"
+    
+    remove_background(image_path, output_path, api_url)