import requests
import base64
from PIL import Image
from io import BytesIO
import matplotlib.pyplot as plt
import argparse
parser = argparse.ArgumentParser(description="Client for FastAPI image inference")
parser.add_argument("--env", '-e', type=str, default="local", help="Environment: local or deployed")
args = parser.parse_args()

if args.env == "local":
    FASTAPI_URL = "http://localhost:8000"
elif args.env == "deployed":
    FASTAPI_URL = "https://riciii7-tumor-detection-fastapi.hf.space"
else: 
    raise ValueError("Invalid environment. Choose 'local' or 'deployed'.")

with open('glioma.jpg', 'rb') as f:
    files = {'file': ('glioma.jpg', f, 'image/jpeg')}
    response = requests.post(f'{FASTAPI_URL}/inference', files=files)

result = response.json()

img_data = base64.b64decode(result['image'])
img = Image.open(BytesIO(img_data))

plt.figure(figsize=(10, 8))
plt.imshow(img)
plt.title(f"Detections: {', '.join(result['detections'])}")
plt.axis('off')
plt.show()

img.save('result_with_detections.jpg')
print(f"Detections found: {result['detections']}")