mirror of https://github.com/skygpu/skynet.git
Fix image getting logic
parent
a9b05b7ee7
commit
9fa5a01c34
|
@ -5,7 +5,7 @@ import random
|
||||||
import logging
|
import logging
|
||||||
import asyncio
|
import asyncio
|
||||||
|
|
||||||
from PIL import Image
|
from PIL import Image, UnidentifiedImageError
|
||||||
from json import JSONDecodeError
|
from json import JSONDecodeError
|
||||||
from decimal import Decimal
|
from decimal import Decimal
|
||||||
from hashlib import sha256
|
from hashlib import sha256
|
||||||
|
@ -244,16 +244,40 @@ class SkynetTelegramFrontend:
|
||||||
async def get_and_set_results(link: str):
|
async def get_and_set_results(link: str):
|
||||||
results[link] = await get_ipfs_file(link)
|
results[link] = await get_ipfs_file(link)
|
||||||
|
|
||||||
|
def get_image_from_resp(resp):
|
||||||
|
png_img = resp.raw
|
||||||
|
with Image.open(io.BytesIO(resp.raw)) as image:
|
||||||
|
w, h = image.size
|
||||||
|
|
||||||
|
if w > TG_MAX_WIDTH or h > TG_MAX_HEIGHT:
|
||||||
|
logging.warning(f'result is of size {image.size}')
|
||||||
|
image.thumbnail((TG_MAX_WIDTH, TG_MAX_HEIGHT))
|
||||||
|
tmp_buf = io.BytesIO()
|
||||||
|
image.save(tmp_buf, format='PNG')
|
||||||
|
png_img = tmp_buf.getvalue()
|
||||||
|
|
||||||
|
return png_img
|
||||||
|
|
||||||
tasks = [
|
tasks = [
|
||||||
get_and_set_results(ipfs_link),
|
get_and_set_results(ipfs_link),
|
||||||
get_and_set_results(ipfs_link_legacy)
|
get_and_set_results(ipfs_link_legacy)
|
||||||
]
|
]
|
||||||
await asyncio.gather(*tasks)
|
await asyncio.gather(*tasks)
|
||||||
|
|
||||||
|
png_img = None
|
||||||
|
|
||||||
resp = results[ipfs_link_legacy]
|
resp = results[ipfs_link_legacy]
|
||||||
if not resp or resp.status_code != 200:
|
if not resp or resp.status_code != 200:
|
||||||
logging.error(f'couldn\'t get ipfs hosted image at {ipfs_link_legacy}!')
|
logging.error(f'couldn\'t get ipfs hosted image at {ipfs_link_legacy}!')
|
||||||
|
|
||||||
|
else:
|
||||||
|
try:
|
||||||
|
png_img = get_image_from_resp(resp)
|
||||||
|
|
||||||
|
except UnidentifiedImageError:
|
||||||
|
logging.error(f'couldn\'t get ipfs hosted image at {ipfs_link_legacy}!')
|
||||||
|
|
||||||
|
if not png_img:
|
||||||
resp = results[ipfs_link]
|
resp = results[ipfs_link]
|
||||||
if not resp or resp.status_code != 200:
|
if not resp or resp.status_code != 200:
|
||||||
logging.error(f'couldn\'t get ipfs hosted image at {ipfs_link}!')
|
logging.error(f'couldn\'t get ipfs hosted image at {ipfs_link}!')
|
||||||
|
@ -265,16 +289,19 @@ class SkynetTelegramFrontend:
|
||||||
)
|
)
|
||||||
return True
|
return True
|
||||||
|
|
||||||
png_img = resp.raw
|
else:
|
||||||
with Image.open(io.BytesIO(resp.raw)) as image:
|
try:
|
||||||
w, h = image.size
|
png_img = get_image_from_resp(resp)
|
||||||
|
|
||||||
if w > TG_MAX_WIDTH or h > TG_MAX_HEIGHT:
|
except UnidentifiedImageError:
|
||||||
logging.warning(f'result is of size {image.size}')
|
logging.error(f'couldn\'t get ipfs hosted image at {ipfs_link}!')
|
||||||
image.thumbnail((TG_MAX_WIDTH, TG_MAX_HEIGHT))
|
await self.update_status_message(
|
||||||
tmp_buf = io.BytesIO()
|
status_msg,
|
||||||
image.save(tmp_buf, format='PNG')
|
caption,
|
||||||
png_img = tmp_buf.getvalue()
|
reply_markup=build_redo_menu(),
|
||||||
|
parse_mode='HTML'
|
||||||
|
)
|
||||||
|
return True
|
||||||
|
|
||||||
logging.info(f'success! sending generated image')
|
logging.info(f'success! sending generated image')
|
||||||
await self.bot.delete_message(
|
await self.bot.delete_message(
|
||||||
|
|
Loading…
Reference in New Issue