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))