Spaces:
Sleeping
Sleeping
| import os | |
| from utils import process_image, run_model | |
| from typing import Tuple, Optional | |
| from PIL import Image | |
| from boto3 import Session | |
| import torch | |
| import pickle | |
| import datetime | |
| import gzip | |
| # Retrieve credentials from environment variables | |
| session = Session( | |
| aws_access_key_id=os.getenv('AWS_ACCESS_KEY_ID'), | |
| aws_secret_access_key=os.getenv('AWS_SECRET_ACCESS_KEY'), | |
| region_name=os.getenv('AWS_DEFAULT_REGION') | |
| ) | |
| s3 = session.client('s3') | |
| def load_model(): | |
| with gzip.open('model_quantized_compressed.pkl.gz', 'rb') as f_in: | |
| model_data = f_in.read() | |
| model = pickle.loads(model_data) | |
| print("Model Loaded") | |
| return model | |
| def upload_to_s3(file_path, bucket_name, s3_key): | |
| with open(file_path, 'rb') as f: | |
| s3.upload_fileobj(f, bucket_name, s3_key) | |
| s3_url = f's3://{bucket_name}/{s3_key}' | |
| return s3_url | |
| def generate_mesh(image_path:str, | |
| output_dir:str ='tmp/output/', | |
| no_remove_bg:bool =True, | |
| foreground_ratio:float =0.85 , | |
| render:bool =False , | |
| mc_resolution:int =256 , | |
| bake_texture_flag:bool =False , | |
| texture_resolution:int =2048, | |
| model=None, | |
| bucket_name:str=None, | |
| input_folder:str=None, | |
| output_folder:str=None, | |
| input_s3_id:str='input_image.png', | |
| output_s3_id:str='output_mesh.obj', | |
| output_video_s3_id:str=None | |
| ) -> Tuple[Optional[str], Optional[str]] : | |
| print('Process start') | |
| image = process_image(image_path=image_path, | |
| output_dir=output_dir , | |
| no_remove_bg=no_remove_bg , | |
| foreground_ratio=foreground_ratio) | |
| print('Process end') | |
| print('Run start') | |
| output_file_path ,output_video_path = run_model(model=model, | |
| image=image, | |
| output_dir=output_dir , | |
| device="cuda:0" if torch.cuda.is_available() else "cpu", | |
| render=render , | |
| mc_resolution=mc_resolution , | |
| model_save_format='obj', | |
| bake_texture_flag=bake_texture_flag , | |
| texture_resolution=texture_resolution) | |
| print('Run end') | |
| print('Uploading to bucket...') | |
| # Upload the input image and generated mesh file to S3 | |
| if input_folder != None: | |
| input_s3_key = input_folder + '/' + input_s3_id | |
| else: | |
| input_s3_key = input_s3_id | |
| if output_folder != None: | |
| output_s3_key = output_folder + '/' + output_s3_id | |
| else: | |
| output_s3_key = output_s3_id | |
| output_video_s3_key = output_video_s3_id | |
| input_s3_url = upload_to_s3(image_path, bucket_name, input_s3_key) | |
| output_s3_url = upload_to_s3(output_file_path, bucket_name, output_s3_key) | |
| if output_video_path != None: | |
| if output_folder != None: | |
| output_video_s3_key = output_folder + '/' + output_video_s3_id | |
| else: | |
| output_video_s3_key = output_video_s3_id | |
| output_video_s3_url = upload_to_s3(output_video_path, bucket_name, output_video_s3_key) | |
| print(f'Files uploaded to S3:\nInput Image: {input_s3_url}\nOutput Mesh: {output_s3_url}\nOutput Video: {output_video_s3_url}') | |
| return output_file_path ,output_video_path |