data_utils/awss3.py

import boto3
import cv2
import os
from dotenv import load_dotenv

load_dotenv()

AWS_ACCESS_KEY_ID = os.environ.get("AWS_ACCESS_KEY_ID")
AWS_SECRET_ACCESS_KEY = os.environ.get("AWS_SECRET_ACCESS_KEY")


class AWSS3:
    def load_image(self, bucket, path, local_path):
        file = self.__s3.get_object(Bucket=bucket, Key=path)
        img_content = file["Body"].read()

        with open(local_path, "wb") as f:
            f.write(img_content)

        img = cv2.imread(local_path, cv2.IMREAD_COLOR)
        img = cv2.cvtColor(src=img, code=cv2.COLOR_BGR2RGB)

        return img

    def save_image(self, bucket, path, local_path) -> bool:
        with open(local_path, "rb") as f:
            image_content = f.read()

        if image_content:
            content_type = "image/" + local_path.rsplit(".", 1)[-1].lower().replace(
                "jpg", "jpeg"
            )
            self.__s3.put_object(
                Bucket=bucket,
                Key=path,
                Body=image_content,
                ACL="public-read",
                ContentDisposition="inline",
                ContentType=content_type,
            )
            return True
        else:
            return False

    def __init__(self):
        self.__s3 = boto3.client(
            "s3",
            aws_access_key_id=AWS_ACCESS_KEY_ID,
            aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
        )