def handler(event, context):
for record in event['Records']:
message = json.loads(record["body"])
chat_id = str(message['chat_id'])
prefix = str(message['prefix'])
tg_client = TgClient(chat_id)
s3 = boto3.client('s3')
# do not process items with same prefix if
# already done - SQS could send duplicates
output_exists = S3Helper.is_file_exist(s3, BUCKET,
'{}_output.jpg'.format(prefix))
if (output_exists is False):
try:
tg_client.send_message('NST - Запустил обработку...')
# requesting EC2 instance to process the image
# NST content and style images
url = 'http://{}/api/?img_prefix={}'.format(
EC2_NST_SERVER_IP, prefix)
response = requests.get(url)
responseJson = json.loads(response.text)
file_name = responseJson['file_name']
file_url = 'https://{}.s3.amazonaws.com/{}'.format(
BUCKET, file_name)
# sending output image back
tg_client.send_photo(file_url)
except Exception as e:
tg_client.send_message('NST - Упс, ошибочка - {}'.format(e))
def is_style_image_exist(s3_client, chat_id):
return S3Helper.is_file_exist(s3_client, BUCKET,
'{}_style.jpg'.format(chat_id))
def handler(event, context):
s3 = boto3.client('s3')
imsize = 512
content_transform = transforms.Compose([
transforms.Resize(imsize),
transforms.ToTensor(),
transforms.Lambda(lambda x: x.mul(255))
])
for record in event['Records']:
message = json.loads(record["body"])
chat_id = str(message['chat_id'])
prefix = str(message['prefix'])
tg_client = TgClient(chat_id)
style_models = s3.list_objects(Bucket=GAN_STYLES_BUCKET)
styles = []
for style in style_models['Contents']:
if '.pth' in style['Key']:
styles.append(style['Key'])
# do not process items with same prefix if
# already done - SQS could send duplicates
output_exists = S3Helper.is_file_exist(
s3, BUCKET, '{}_output_{}.jpg'.format(prefix, styles[0]))
if (output_exists is False):
try:
tg_client.send_message('GAN - Нашел {} стиля, поехали!'.format(
len(styles)))
for i in range(len(styles)):
style = styles[i]
message_text = 'GAN - Запустил обработку стилем {}...'.format(
style)
if i == len(styles) - 1:
message_text = 'GAN - И последний - {}...'.format(
style)
tg_client.send_message(message_text)
content_img_file_name = '{}_content.jpg'.format(prefix)
content_img_file_path = f"/tmp/{content_img_file_name}"
s3.download_file(BUCKET, content_img_file_name,
content_img_file_path)
img = Image.open(content_img_file_path).convert('RGB')
img = content_transform(img)
img = img.unsqueeze(0).to(device)
model_file_name = style
model_file_path = f"/tmp/{model_file_name}"
s3.download_file(GAN_STYLES_BUCKET, model_file_name,
model_file_path)
style_model = TransformerNet()
state_dict = torch.load(model_file_path,
map_location=torch.device('cpu'))
for k in list(state_dict.keys()):
if re.search(r'in\d+\.running_(mean|var)$', k):
del state_dict[k]
style_model.load_state_dict(state_dict)
style_model.to(device)
with torch.no_grad():
output = style_model(img)
img = output[0].clone().clamp(0, 255).numpy()
img = img.transpose(1, 2, 0).astype("uint8")
print(img.shape, type(img))
img = Image.fromarray(img)
output_image_file_name = '{}_output_{}.jpg'.format(
prefix, style)
output_image_file_path = f"/tmp/{output_image_file_name}"
img.save(output_image_file_path, format="PNG")
with open(output_image_file_path, 'rb') as output_data:
s3 \
.put_object(
Bucket=BUCKET,
ACL='public-read',
Key=output_image_file_name,
Body=output_data)
file_url = 'https://{}.s3.amazonaws.com/{}'.format(
BUCKET, output_image_file_name)
# sending output image back
tg_client.send_photo(file_url)
except Exception as e:
tg_client.send_message('GAN - Упс, ошибочка - {}'.format(e))
