feat: 텍스트 마스킹 후 업로드 기능으로 변경

Author: brothergiven

.github/workflows/deploy.yml

@@ -179,6 +179,8 @@ jobs:
           task-definition: task-definition.json
           container-name: masking-service # Task Definition에 정의된 컨테이너 이름
           image: ${{ steps.build-image.outputs.image }}
+          environment-variables: |
+            S3_BUCKET_NAME=${{ secrets.S3_BUCKET_NAME }}
 
 
       - name: Deploy Amazon ECS task definition

masking-service/main.py

@@ -1,11 +1,12 @@
 import uvicorn
-from fastapi import FastAPI
-from pydantic import BaseModel
-import requests
+from fastapi import FastAPI, UploadFile, File, HTTPException
 from io import BytesIO
 from PIL import Image, ImageDraw
 import easyocr
 import numpy as np
+import boto3
+import os
+import uuid
 
 app=FastAPI(
     title="Image Masking Service",
@@ -16,8 +17,14 @@
     openapi_url="/mask/openapi.json"
 )
 
-class ImageReq(BaseModel):
-    image_url: str
+# Initialize S3 client
+s3_client = boto3.client(
+    's3',
+    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_REGION")
+)
+S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME")
 
 reader = easyocr.Reader(['en', 'ko'])
 
@@ -37,24 +44,35 @@ def mask_text(image: Image.Image) -> Image.Image:
 
     return masked
 
-@app.post("/mask/exec", summary="Mask text in image", description="Mask text in the given image URL.")
-async def mask_image(req: ImageReq):
+@app.post("/mask/exec", summary="Mask text in image and upload to S3", description="Upload an image, mask text, and upload the masked image to AWS S3.")
+async def mask_image(file: UploadFile = File(...)):
+    if not S3_BUCKET_NAME:
+        raise HTTPException(status_code=500, detail="S3_BUCKET_NAME environment variable not set.")
+
     try:
-        response = requests.get(req.image_url)
-        response.raise_for_status()
-    except requests.RequestException as e:
-        return {"error": f"Failed to fetch image from URL: {e}"}
+        contents = await file.read()
+        img = Image.open(BytesIO(contents)).convert("RGB")
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Error reading image file: {e}")
 
-    # try:
-    img = Image.open(BytesIO(response.content)).convert("RGB")
     masked_img = mask_text(img)
-    #     buffered = BytesIO()
-    #     masked_img.save(buffered, format="PNG")
-    #     img_str = "data:image/png;base64," + base64.b64encode(buffered.getvalue()).decode()
-    # except Exception as e:
-    #     return {"error": f"Error processing image: {e}"}
 
-    return {"masked_image": "masking completed"}
+    # Save masked image to a BytesIO object
+    masked_img_buffer = BytesIO()
+    masked_img.save(masked_img_buffer, format="PNG")
+    masked_img_buffer.seek(0) # Rewind the buffer to the beginning
+
+    # Generate a unique filename for S3
+    file_extension = file.filename.split(".")[-1] if "." in file.filename else "png"
+    s3_filename = f"masked_images/{uuid.uuid4()}.{file_extension}"
+
+    try:
+        s3_client.upload_fileobj(masked_img_buffer, S3_BUCKET_NAME, s3_filename)
+        s3_url = f"https://{S3_BUCKET_NAME}.s3.amazonaws.com/{s3_filename}"
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to upload masked image to S3: {e}")
+
+    return {"message": "Masking and upload completed successfully", "s3_url": s3_url}
 
 @app.get("/health", summary="Health Check", description="Check if the service is running.")
 async def health_check():

masking-service/requirements.txt

@@ -9,3 +9,4 @@ Pillow
 httpx
 --extra-index-url https://download.pytorch.org/whl/cpu
 torch
+boto3