55 lines
1.7 KiB
Python
55 lines
1.7 KiB
Python
|
import requests
|
||
|
|
import base64
|
||
|
|
import os
|
||
|
|
from PIL import Image
|
||
|
|
import io
|
||
|
|
from datetime import datetime
|
||
|
|
|
||
|
|
def image_to_base64(image_path):
|
||
|
|
with open(image_path, "rb") as image_file:
|
||
|
|
return base64.b64encode(image_file.read()).decode()
|
||
|
|
|
||
|
|
def base64_to_image(base64_str):
|
||
|
|
image_data = base64.b64decode(base64_str)
|
||
|
|
return Image.open(io.BytesIO(image_data))
|
||
|
|
|
||
|
|
def save_results(generated_images, output_dir="results"):
|
||
|
|
os.makedirs(output_dir, exist_ok=True)
|
||
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
|
|
|
||
|
|
saved_paths = []
|
||
|
|
for idx, img_base64 in enumerate(generated_images):
|
||
|
|
img = base64_to_image(img_base64)
|
||
|
|
output_path = os.path.join(output_dir, f"result_{timestamp}_{idx}.png")
|
||
|
|
img.save(output_path)
|
||
|
|
saved_paths.append(output_path)
|
||
|
|
print(f"Saved image to {output_path}")
|
||
|
|
|
||
|
|
return saved_paths
|
||
|
|
|
||
|
|
data = {
|
||
|
|
"model_image": image_to_base64("imgs/person.jpg"),
|
||
|
|
"garment_image": image_to_base64("imgs/garment.jpg"),
|
||
|
|
"category": "Upper-body",
|
||
|
|
"resolution": "768x1024",
|
||
|
|
"n_steps": 30,
|
||
|
|
"image_scale": 2.0,
|
||
|
|
"num_images": 1
|
||
|
|
}
|
||
|
|
|
||
|
|
try:
|
||
|
|
response = requests.post("http://localhost:8001/try-on", json=data)
|
||
|
|
response.raise_for_status()
|
||
|
|
result = response.json()
|
||
|
|
|
||
|
|
if result["status"] == "success":
|
||
|
|
saved_files = save_results(result["generated_images"])
|
||
|
|
print(f"Successfully generated {len(saved_files)} images")
|
||
|
|
print(f"Seed used: {result['seed']}")
|
||
|
|
else:
|
||
|
|
print("Generation failed:", result.get("detail", "Unknown error"))
|
||
|
|
|
||
|
|
except requests.exceptions.RequestException as e:
|
||
|
|
print(f"Error making request: {e}")
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Error processing results: {e}")
|