Octo-1.5-Small / test_api.py
Nirav Madhani
Flash server
0558e79
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)