examples/rt_nto_example.py

"""
project @ CTO_TCP_ZERO_GPU
created @ 2024-11-18
author  @ github.com/ishworrsubedii
"""
import asyncio
import aiohttp
import json
from pathlib import Path


async def test_rt_nto(
        image_path: str,
        necklace_ids: list,
        categories: list,
        storename: str,
        api_url: str = "http://localhost:8000/rt_nto"
):
    image_content = None
    with open(image_path, 'rb') as f:
        image_content = f.read()

    data = aiohttp.FormData()
    data.add_field('necklace_id_list', ','.join(necklace_ids))
    data.add_field('category_list', ','.join(categories))
    data.add_field('storename', storename)
    data.add_field('image',
                   image_content,
                   filename=Path(image_path).name,
                   content_type='image/jpeg')

    async with aiohttp.ClientSession() as session:
        try:
            async with session.post(api_url, data=data) as response:
                if response.status != 200:
                    print(f"Server returned status code: {response.status}")
                    try:
                        error_content = await response.text()
                        print(f"Error content: {error_content}")
                    except Exception as e:
                        print(f"Could not read error content: {e}")
                    return

                async for line in response.content:
                    if line:
                        try:
                            result = json.loads(line.decode('utf-8'))
                            if result.get('code') == 200:
                                print(f"Success - Progress: {result['progress']}")
                                print(f"Necklace ID: {result['necklace_id']}")
                                print(f"Category: {result['category']}")
                                print(f"Output URL: {result['output']}")
                                print(f"Mask URL: {result['mask']}")
                                print(f"Inference Time: {result['inference_time']}s")
                            else:
                                print(f"Error - Progress: {result.get('progress')}")
                                print(f"Error Details: {result.get('error')}")
                            print("-" * 50)
                        except json.JSONDecodeError as e:
                            print(f"Error decoding response: {e}")
                            print(f"Raw line: {line}")
        except aiohttp.ClientError as e:
            print(f"Connection error: {e}")
        except Exception as e:
            print(f"Unexpected error: {e}")


async def main():
    image_path = "/home/ishwor/Downloads/download (7).png"
    necklace_ids = ["ST01SHO001", "ST01SHO002", "ST01SHO003"]
    categories = ["Short Necklaces", "Short Necklaces", "Short Necklaces"]
    storename = "Store1"

    try:
        await test_rt_nto(
            image_path=image_path,
            necklace_ids=necklace_ids,
            categories=categories,
            storename=storename
        )
    except Exception as e:
        print(f"Error in main: {e}")


if __name__ == "__main__":
    asyncio.run(main())