Spaces:
Running
Running
import os | |
from utils import process_image, run_model | |
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, output_dir, model): | |
print('Process start') | |
# Process the image | |
image = process_image(image_path, output_dir) | |
print('Process end') | |
print('Run start') | |
output_file_path = run_model(model, image, output_dir) | |
print('Run end') | |
# Upload the input image and generated mesh file to S3 | |
bucket_name = 'vasana-bkt1' | |
input_s3_key = f'input_images/{datetime.datetime.now().strftime("%Y%m%d%H%M%S")}-{os.path.basename(image_path)}' | |
output_s3_key = f'output_meshes/{datetime.datetime.now().strftime("%Y%m%d%H%M%S")}-{os.path.basename(output_file_path)}' | |
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) | |
print(f'Files uploaded to S3:\nInput Image: {input_s3_url}\nOutput Mesh: {output_s3_url}') | |
return output_file_path | |