Spaces:
Running
Running
import requests | |
import base64 | |
from PIL import Image | |
import io | |
# API endpoint URL (adjust if running on a different host/port) | |
API_URL = "http://localhost:8000/predict" | |
# Example image URL from main.py | |
IMAGE_URL = "https://rail.eecs.berkeley.edu/datasets/bridge_release/raw/bridge_data_v2/datacol2_toykitchen7/drawer_pnp/01/2023-04-19_09-18-15/raw/traj_group0/traj0/images0/im_12.jpg" | |
TASK_TEXT = "pick up the fork" | |
def test_api(image_url=IMAGE_URL, task=TASK_TEXT): | |
try: | |
# Download image from URL | |
response = requests.get(image_url, stream=True) | |
response.raise_for_status() # Check for HTTP errors | |
img = Image.open(response.raw).resize((256, 256)) | |
# Convert image to base64 | |
img_byte_arr = io.BytesIO() | |
img.save(img_byte_arr, format="JPEG") # Save as JPEG (adjust if needed) | |
img_byte_arr = img_byte_arr.getvalue() | |
base64_string = base64.b64encode(img_byte_arr).decode("utf-8") | |
# Prepare payload for API | |
payload = { | |
"image_base64": base64_string, | |
"task": task | |
} | |
# Send POST request to API | |
api_response = requests.post(API_URL, json=payload) | |
api_response.raise_for_status() # Check for API errors | |
# Print the result | |
result = api_response.json() | |
print(f"Task: {task}") | |
print(f"Image URL: {image_url}") | |
print(f"Predicted Actions: {result['actions']}") | |
except requests.exceptions.RequestException as e: | |
print(f"Error fetching image or calling API: {e}") | |
except Exception as e: | |
print(f"Unexpected error: {e}") | |
if __name__ == "__main__": | |
# Test with default values (same as main.py) | |
test_api() | |
# Test with a different URL and task (optional) | |
# Replace with another valid URL if desired | |
print("\nTesting with another URL and task:") | |
test_api(IMAGE_URL, TASK_TEXT) |