From 412086cb584ea35b46cfb55bfb35e290fe7357b1 Mon Sep 17 00:00:00 2001 From: semnisem Date: Sat, 16 Nov 2024 19:58:23 +0900 Subject: [PATCH 01/17] =?UTF-8?q?Feat:=20=EC=B6=94=EA=B0=80=20AI=EA=B8=B0?= =?UTF-8?q?=EB=B0=98=20=ED=95=84=EA=B8=B0=20=EC=A0=9C=EA=B1=B0=20=EB=B0=B1?= =?UTF-8?q?=EC=97=94=EB=93=9C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit s3에서 모델 다운로드 및 필기 제거 API --- app/AIProcessor.py | 142 +++++++++++++++++++++++++++++++++++++++++++++ app/main.py | 60 ++++++++++++++++++- 2 files changed, 201 insertions(+), 1 deletion(-) create mode 100644 app/AIProcessor.py diff --git a/app/AIProcessor.py b/app/AIProcessor.py new file mode 100644 index 0000000..8ee797e --- /dev/null +++ b/app/AIProcessor.py @@ -0,0 +1,142 @@ +import torch +import torchvision +import numpy as np +import cv2 +from ultralytics import YOLO +from segment_anything import sam_model_registry, SamPredictor +import io +import logging + + +class AIProcessor: + def __init__(self, yolo_path: str, sam_path: str, device: torch.device = torch.device("cpu")): + self.yolo_path = yolo_path + self.sam_path = sam_path + self.device = device + self.yolo_model = YOLO(self.yolo_path) # yolo 로드 + self.sam_model = None + self.predictor = None + + self.alpha_channel = None + self.size = 0 + self.height = 0 + self.width = 0 + + def remove_alpha(self, image): + self.height, self.width = image.shape[:2] + self.size = self.height * self.width + if image.shape[2] == 4: # RGBA or RGB + self.alpha_channel = image[:, :, 3] + image_rgb = image[:, :, :3] + else: + image_rgb = image + + return image_rgb + + def load_sam_model(self, model_type="vit_h"): # sam 로드 + self.sam_model = sam_model_registry[model_type](checkpoint=self.sam_path) + self.sam_model.to(self.device) + self.predictor = SamPredictor(self.sam_model) + + def object_detection(self, img): + results = self.yolo_model.predict(source=img, imgsz=640, device=self.device, + iou=0.3, conf=0.3) + bbox = results[0].boxes.xyxy.tolist() + logging.info(f'1차 마스킹 - 바운딩 박스 생성 완료, {len(bbox)}개') + return bbox + + def segment_from_points(self, user_points, user_labels, save_path=None): + input_point = np.array(user_points) + input_label = np.array(user_labels) + + masks_points, scores, logits = self.predictor.predict( + point_coords=input_point, + point_labels=input_label, + multimask_output=False, + ) + + mask_points_uint8 = (masks_points[0] * 255).astype(np.uint8) + # cv2.imwrite(save_path, mask_points_uint8) + logging.info('2차 마스킹 - 사용자 입력 점 세그먼트 완료.') + return mask_points_uint8 + + def segment_from_boxes(self, image, bbox, save_path=None): + input_boxes = torch.tensor(bbox, device=self.predictor.device) + transformed_boxes = self.predictor.transform.apply_boxes_torch(input_boxes, image.shape[:2]) + masks, _, _ = self.predictor.predict_torch( + point_coords=None, + point_labels=None, + boxes=transformed_boxes, + multimask_output=False, + ) + masks_np = np.zeros(masks.shape[-2:], dtype=np.uint8) + for mask in masks: + mask_np = mask.cpu().numpy().astype(np.uint8) * 255 + mask_np = mask_np.squeeze() + masks_np = cv2.bitwise_or(masks_np, mask_np) + + # cv2.imwrite(save_path, masks_np) + logging.info('1차 마스킹 - 바운딩 박스 세그먼트 완료.') + return masks_np + + def inpainting(self, image, mask_total): + inpainted_image = cv2.inpaint(image.copy(), mask_total, 10, cv2.INPAINT_TELEA) + final_image = cv2.convertScaleAbs(inpainted_image, alpha=1.5, beta=10) + # cv2.imwrite('test_images/inpainted_init.png', inpainted_image) + # cv2.imwrite('test_images/inpainted_final.png', final_image) + logging.info('인페인팅 및 후보정 완료.') + + return final_image + + def combine_alpha(self, image_rgb): + if self.alpha_channel is not None: # RGBA + image_rgba = cv2.merge([image_rgb, self.alpha_channel]) + return image_rgba + else: + return image_rgb + + def process(self, img_bytes, user_points, user_labels, extension='jpg'): + """ local test 용도 Vs. server test 용도 구분 """ + ### local 용도 + # img_path = img_bytes + # image = cv2.imread(img_path, cv2.IMREAD_COLOR) + + ### server 용도 + buffer = np.frombuffer(img_bytes, dtype=np.uint8) + image = cv2.imdecode(buffer, cv2.IMREAD_UNCHANGED) + + ### ready + self.load_sam_model() + self.predictor.set_image(image) + masks_total = np.zeros(image.shape[:2], dtype=np.uint8) + + ### 1차: Object Detection & Segment by Box + bbox = self.object_detection(image) + if len(bbox) > 0: + masks_by_box = self.segment_from_boxes(image, bbox, save_path=None) # 'test_images/seg_box.png' + masks_total = cv2.bitwise_or(masks_total, masks_by_box) + logging.info( + f"1차 마스킹 후 shape 점검: YOLOv11 감지된 영역 shape: {masks_by_box.shape}, 이미지 영역 shape: {image.shape}") # (1893, 1577, 3) (1893, 1577) + + ### 2차: points arguments by User & Segment by Points + if len(user_points) > 0: + mask_by_point = self.segment_from_points(user_points, user_labels, save_path=None) # save_path='test_images/seg_points.png' + masks_total = cv2.bitwise_or(masks_total, mask_by_point) + # cv2.imwrite('test_images/mask_total.png', masks_total) + if isinstance(masks_total, np.ndarray): + image_output = self.inpainting(image, masks_total) + logging.info('최종 마스킹 이미지 생성 완료') + else: + image_output = image + logging.info('최종 마스킹이 없거나, numpy 형식의 배열이 아님.') + + # image_input = self.combine_alpha(image_rgb) + # image_output = self.combine_alpha(image_output) + _, input_bytes = cv2.imencode("." + extension, image) + _, mask_bytes = cv2.imencode("." + extension, masks_total.astype(np.uint8)) + _, result_bytes = cv2.imencode("." + extension, image_output) + + # return 0, 0, 0 + return io.BytesIO(input_bytes), io.BytesIO(mask_bytes), io.BytesIO(result_bytes) + + diff --git a/app/main.py b/app/main.py index 5abd3dd..a4991df 100644 --- a/app/main.py +++ b/app/main.py @@ -65,10 +65,68 @@ def upload_image_to_s3(file_bytes, file_path): s3_client.upload_fileobj(file_bytes, BUCKET_NAME, file_path) +def download_model_from_s3(yolo_path: str = 'models/yolo11_best.pt', sam_path: str = 'models/sam_vit_h_4b8939.pth'): + dest_dir = f'../' # 모델을 저장할 컨테이너 내 경로 + try: + s3_client.download_file(BUCKET_NAME, yolo_path, dest_dir+yolo_path) + s3_client.download_file(BUCKET_NAME, sam_path, dest_dir+sam_path) + logger.info(f'YOLOv11 & SAM models downloaded successfully to {dest_dir}') + logger.info(f"files in 'models' Dir: {os.listdir(dest_dir)}") + except Exception as e: + print(f'Failed to download model: {e}') + + @app.get("/", status_code=status.HTTP_200_OK) def greeting(): - return JSONResponse(content={"message": "Hello! Let's start image processing"}) + return JSONResponse(content={"message": "Hello! Welcome to OnO's FastAPI Server!"}) + + +@app.get("/load-models", status_code=status.HTTP_200_OK) +async def get_models(): + try: + download_model_from_s3() + except Exception as e: + logger.error("Error with Download & Saving AIs: %s", e) + raise HTTPException(status_code=500, detail="Error with Download & Saving AIs") +@app.post("/process-shape") +async def processShape(request: Request): + """ AI handwriting detection & Telea Algorithm-based inpainting """ + data = await request.json() + full_url = data['fullUrl'] + point_list = data.get('points') + label_list = data.get('labels') # value or None + logger.info(f"사용자 입력 포인트: {point_list}") + logger.info(f"사용자 입력 라벨: {label_list}") + + try: + s3_key = parse_s3_url(full_url) + paths = create_file_path(s3_key, s3_key.split(".")[-1]) + img_bytes = download_image_from_s3(s3_key) # download from S3 + corrected_img_bytes = ImageManager.correct_rotation(img_bytes, paths['extension']) + logger.info(f"시용자 입력 이미지({s3_key}) 다운로드 및 전처리 완료") + + # aiProcessor = AIProcessor(yolo_path="./models/yolo11_best.pt", sam_path="./models/sam_vit_h_4b8939.pth") # local + aiProcessor = AIProcessor(yolo_path="../models/yolo11_best.pt", sam_path="../models/sam_vit_h_4b8939.pth") # server + img_input_bytes, img_mask_bytes, img_output_bytes = aiProcessor.process(img_bytes=corrected_img_bytes, + user_points=point_list, + user_labels=label_list) + logger.info("AI 필기 제거 프로세스 완료") + + upload_image_to_s3(img_input_bytes, paths["input_path"]) + upload_image_to_s3(img_mask_bytes, paths["mask_path"]) + upload_image_to_s3(img_output_bytes, paths["output_path"]) + + logger.info("AI 필기 제거 결과 이미지 업로드 완료") + return JSONResponse(content={"message": "File processed successfully", "path": paths}) + + except KeyError as e: + raise HTTPException(status_code=400, detail=f"Missing key: {e.args[0]}") + except ValueError as e: + raise HTTPException(status_code=422, detail=str(e)) + except Exception as e: + logger.error("Error during processing: %s", e) + raise HTTPException(status_code=500, detail="Error processing the image.") @app.post("/process-color") async def processColor(request: Request): From 4e626fa0b270e278666e88da4b803d079baa11ce Mon Sep 17 00:00:00 2001 From: semnisem Date: Sat, 16 Nov 2024 20:02:33 +0900 Subject: [PATCH 02/17] =?UTF-8?q?Feat:=20=EC=B6=94=EA=B0=80=20=ED=8C=A8?= =?UTF-8?q?=ED=82=A4=EC=A7=80=20=20&=20=EB=8F=84=EC=BB=A4=20=EB=94=94?= =?UTF-8?q?=EB=A0=89=ED=86=A0=EB=A6=AC?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit requirements.txt 및 import 추가 도커 작업 디렉토리 하위폴더 models 생성 --- Dockerfile | 1 + app/main.py | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/Dockerfile b/Dockerfile index 610fab7..6446684 100644 --- a/Dockerfile +++ b/Dockerfile @@ -14,6 +14,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ # 작업 디렉토리 설정 WORKDIR /test +RUN mkdir models # 의존성 파일 복사 및 설치 COPY ./requirements.txt /test/requirements.txt diff --git a/app/main.py b/app/main.py index a4991df..76c83c7 100644 --- a/app/main.py +++ b/app/main.py @@ -7,6 +7,7 @@ from starlette import status from starlette.responses import StreamingResponse, JSONResponse from ColorRemover import ColorRemover +from AIProcessor import AIProcessor import ImageFunctions as ImageManager import logging import os @@ -36,7 +37,7 @@ def create_file_path(obj_path, extension): - file_id = uuid4() # 각 클라이언트마다 고유한 파일 ID 생성 + file_id = uuid4()[:4] # 각 클라이언트마다 고유한 파일 ID 생성 dir_path = obj_path.rsplit('/', 1)[0] paths = {"input_path": f"{dir_path}/{file_id}.input.{extension}", "output_path": f"{dir_path}/{file_id}.output.{extension}", From ad64a37cb57992f2ef30bff4ceeced50d85dbb0d Mon Sep 17 00:00:00 2001 From: semnisem Date: Sat, 16 Nov 2024 20:02:45 +0900 Subject: [PATCH 03/17] =?UTF-8?q?Feat:=20=EC=B6=94=EA=B0=80=20=ED=8C=A8?= =?UTF-8?q?=ED=82=A4=EC=A7=80=20=20&=20=EB=8F=84=EC=BB=A4=20=EB=94=94?= =?UTF-8?q?=EB=A0=89=ED=86=A0=EB=A6=AC?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit requirements.txt 및 import 추가 도커 작업 디렉토리 하위폴더 models 생성 --- requirements.txt | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 2abd4e0..230ff78 100644 --- a/requirements.txt +++ b/requirements.txt @@ -44,4 +44,9 @@ botocore==1.34.141 requests==2.32.3 charset-normalizer==3.3.2 openai==1.46.0 -pymilvus==2.4.6 \ No newline at end of file +pymilvus==2.4.6 +torch==2.5.1 +torchvision==0.20.1 +matplotlib==3.9.1.post1 +ultralytics==8.3.31 +git+https://github.com/facebookresearch/segment-anything.git@dca509fe793f601edb92606367a655c15ac00fdf From 973def589895da5a7a503bee4477a0e7e6430808 Mon Sep 17 00:00:00 2001 From: semnisem Date: Sat, 16 Nov 2024 20:03:43 +0900 Subject: [PATCH 04/17] =?UTF-8?q?Chores:=20=EA=B0=9C=EC=84=A0=20=ED=94=84?= =?UTF-8?q?=EB=A1=AC=ED=94=84=ED=8C=85?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 프롬프팅 및 루트 인사말 --- app/main.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/app/main.py b/app/main.py index 76c83c7..11fa14e 100644 --- a/app/main.py +++ b/app/main.py @@ -40,8 +40,8 @@ def create_file_path(obj_path, extension): file_id = uuid4()[:4] # 각 클라이언트마다 고유한 파일 ID 생성 dir_path = obj_path.rsplit('/', 1)[0] paths = {"input_path": f"{dir_path}/{file_id}.input.{extension}", - "output_path": f"{dir_path}/{file_id}.output.{extension}", "mask_path": f"{dir_path}/{file_id}.mask.{extension}", + "output_path": f"{dir_path}/{file_id}.output.{extension}", "extension": extension} return paths @@ -517,8 +517,8 @@ async def retrieve(problem_text: str): async def augment(curriculum_context, query): prompt = ("너는 고등학생의 오답 문제를 통해 약점을 보완해주는 공책이야. \ 교육과정을 참고해서 오답 문제 핵심 의도를 바탕으로 문제에서 헷갈릴만한 요소, \ - 학생이 놓친 것 같은 중요한 개념을 찾아 그 개념에 대해 4줄 이내로 설명해주고, \ - 그 개념을 적용해서 풀이를 요점에 따라서 짧게(4줄 내외) 작성해줘. \ + 이 문제를 틀렸다면 놓칠 것 같은 같은 중요한 개념을 연관지어 그 개념에 대해 4줄 이내로 설명해주고, \ + 그 개념을 적용해서 풀이를 핵심적 원리에 집중하여 짧게(4줄 내외) 작성해줘. \ 만약 오답 문제와 교과과정이 관련이 없다고 판단되면, 교육과정은 참고하지 않으면 돼. \n\n\n") passage = f"오답 문제 : {query} \n\n\n" context = f"교과과정 : {curriculum_context} \n\n\n" From 01da7fce41aef91c765b1621895ecfd4eb4b0fe4 Mon Sep 17 00:00:00 2001 From: semnisem Date: Sat, 16 Nov 2024 20:15:11 +0900 Subject: [PATCH 05/17] =?UTF-8?q?Hot-Fix:=20=EC=88=98=EC=A0=95=20=EB=8F=84?= =?UTF-8?q?=EC=BB=A4=20=ED=8C=8C=EC=9D=BC=20=EA=B9=83=20=EC=84=A4=EC=B9=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit git Command가 없기 때문에 깃 액션에서 에러가 발생했음 --- Dockerfile | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/Dockerfile b/Dockerfile index 6446684..7af7b6a 100644 --- a/Dockerfile +++ b/Dockerfile @@ -10,7 +10,8 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ libglib2.0-0 \ gcc \ python3-dev && \ - rm -rf /var/lib/apt/lists/* + git && \ + rm -rf /var/lib/apt/lists/* \ # 작업 디렉토리 설정 WORKDIR /test From bc1992062ebeafa1ad401751af64aacb499c2b44 Mon Sep 17 00:00:00 2001 From: semnisem Date: Sat, 16 Nov 2024 20:20:23 +0900 Subject: [PATCH 06/17] =?UTF-8?q?Hot-Fix:=20=EC=88=98=EC=A0=95=20=EB=8F=84?= =?UTF-8?q?=EC=BB=A4=20=ED=8C=8C=EC=9D=BC=20=EA=B9=83=20=EC=84=A4=EC=B9=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit update 후 -y로 git 설치해야 함. --- Dockerfile | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/Dockerfile b/Dockerfile index 7af7b6a..8661e5f 100644 --- a/Dockerfile +++ b/Dockerfile @@ -5,12 +5,11 @@ FROM python:3.12-slim LABEL org.opencontainers.image.source="https://github.com/AI-SIP/MVP_CV" # 필요한 시스템 패키지 먼저 설치 -RUN apt-get update && apt-get install -y --no-install-recommends \ +RUN apt-get update && apt-get install -y git \ libgl1-mesa-glx \ libglib2.0-0 \ gcc \ python3-dev && \ - git && \ rm -rf /var/lib/apt/lists/* \ # 작업 디렉토리 설정 From afe22d99ff5832fa657ffe8fd1aa53c36d24645c Mon Sep 17 00:00:00 2001 From: semnisem Date: Sat, 16 Nov 2024 20:50:26 +0900 Subject: [PATCH 07/17] =?UTF-8?q?Hot-Fix:=20=EC=88=98=EC=A0=95=20AI=20?= =?UTF-8?q?=ED=95=84=EA=B8=B0=20=EB=B6=84=EC=84=9D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 수정: uuid는 자를 수 없음 이동: 디렉토리 생성 명령어는 도커파일 CMD에서 수정: 모델 없을 경우에만 다운로드 추가: 모델로드 기본 1회 실행 --- Dockerfile | 3 +-- app/main.py | 23 ++++++++++++++++++----- 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/Dockerfile b/Dockerfile index 8661e5f..91a6904 100644 --- a/Dockerfile +++ b/Dockerfile @@ -14,7 +14,6 @@ RUN apt-get update && apt-get install -y git \ # 작업 디렉토리 설정 WORKDIR /test -RUN mkdir models # 의존성 파일 복사 및 설치 COPY ./requirements.txt /test/requirements.txt @@ -27,4 +26,4 @@ COPY ./app /test/app/ WORKDIR /test/app # 컨테이너 실행 시 실행할 명령어 -CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"] \ No newline at end of file +CMD ["bash", "-c", "mkdir -p /test/models && uvicorn main:app --host 0.0.0.0 --port 8000"] \ No newline at end of file diff --git a/app/main.py b/app/main.py index 11fa14e..cf6b286 100644 --- a/app/main.py +++ b/app/main.py @@ -37,7 +37,7 @@ def create_file_path(obj_path, extension): - file_id = uuid4()[:4] # 각 클라이언트마다 고유한 파일 ID 생성 + file_id = uuid4() # 각 클라이언트마다 고유한 파일 ID 생성 dir_path = obj_path.rsplit('/', 1)[0] paths = {"input_path": f"{dir_path}/{file_id}.input.{extension}", "mask_path": f"{dir_path}/{file_id}.mask.{extension}", @@ -69,10 +69,21 @@ def upload_image_to_s3(file_bytes, file_path): def download_model_from_s3(yolo_path: str = 'models/yolo11_best.pt', sam_path: str = 'models/sam_vit_h_4b8939.pth'): dest_dir = f'../' # 모델을 저장할 컨테이너 내 경로 try: - s3_client.download_file(BUCKET_NAME, yolo_path, dest_dir+yolo_path) - s3_client.download_file(BUCKET_NAME, sam_path, dest_dir+sam_path) - logger.info(f'YOLOv11 & SAM models downloaded successfully to {dest_dir}') - logger.info(f"files in 'models' Dir: {os.listdir(dest_dir)}") + yolo_full_path = dest_dir+yolo_path + sam_full_path = dest_dir+sam_path + if not os.path.exists(yolo_full_path): + s3_client.download_file(BUCKET_NAME, yolo_path, yolo_full_path) + logger.info(f'YOLOv11 & SAM models downloaded successfully to {dest_dir}') + else: + logger.info(f'YOLOv11 already exists at {yolo_full_path}') + if not os.path.exists(sam_full_path): + s3_client.download_file(BUCKET_NAME, sam_path, sam_full_path) + logger.info(f'SAM models downloaded successfully to {dest_dir}') + else: + logger.info(f'SAM models already exists at {dest_dir}') + + + logger.info(f"Files in 'models' Dir: {os.listdir(dest_dir)}") except Exception as e: print(f'Failed to download model: {e}') @@ -90,6 +101,8 @@ async def get_models(): logger.error("Error with Download & Saving AIs: %s", e) raise HTTPException(status_code=500, detail="Error with Download & Saving AIs") +get_models() + @app.post("/process-shape") async def processShape(request: Request): """ AI handwriting detection & Telea Algorithm-based inpainting """ From ca972d44fe36539b5bfd6777c4802ec364569881 Mon Sep 17 00:00:00 2001 From: semnisem Date: Tue, 19 Nov 2024 21:23:33 +0900 Subject: [PATCH 08/17] =?UTF-8?q?Add:=20=EC=82=AC=EC=9A=A9=EC=9E=90=20?= =?UTF-8?q?=EC=9E=85=EB=A0=A5=20=ED=95=84=ED=84=B0=EB=A7=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 바운딩 박스 내 입력은 제외하기 --- app/AIProcessor.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/app/AIProcessor.py b/app/AIProcessor.py index 8ee797e..d6947a1 100644 --- a/app/AIProcessor.py +++ b/app/AIProcessor.py @@ -22,6 +22,21 @@ def __init__(self, yolo_path: str, sam_path: str, device: torch.device = torch.d self.height = 0 self.width = 0 + def remove_points_in_bboxes(point_list, label_list, bbox_list): + def is_point_in_bbox(p, b): + x, y = p + x_min, y_min, x_max, y_max = b + return x_min <= x <= x_max and y_min <= y <= y_max + + filtered_p = [] + filtered_l = [] + for p, l in zip(point_list, label_list): + if not any(is_point_in_bbox(p, b) for b in bbox_list): + filtered_p.append(p) + filtered_l.append(l) + + return filtered_p, filtered_l + def remove_alpha(self, image): self.height, self.width = image.shape[:2] self.size = self.height * self.width From 52e9eda8ddc497115f586ad1a0f250ff7001e90e Mon Sep 17 00:00:00 2001 From: semnisem Date: Tue, 19 Nov 2024 21:24:44 +0900 Subject: [PATCH 09/17] =?UTF-8?q?Fix:=20=EC=84=B8=EA=B7=B8=EB=A8=BC?= =?UTF-8?q?=ED=8A=B8=20=EB=AA=A8=EB=8D=B8=20=EB=B3=80=EA=B2=BD=20-=20?= =?UTF-8?q?=EA=B2=BD=EB=9F=89=ED=99=94?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 경량화 --- app/AIProcessor.py | 60 +++++++++++++++++++++++++++++++++------------- app/main.py | 3 +-- 2 files changed, 45 insertions(+), 18 deletions(-) diff --git a/app/AIProcessor.py b/app/AIProcessor.py index d6947a1..da2ab17 100644 --- a/app/AIProcessor.py +++ b/app/AIProcessor.py @@ -3,6 +3,7 @@ import numpy as np import cv2 from ultralytics import YOLO +from ultralytics import SAM from segment_anything import sam_model_registry, SamPredictor import io import logging @@ -14,7 +15,7 @@ def __init__(self, yolo_path: str, sam_path: str, device: torch.device = torch.d self.sam_path = sam_path self.device = device self.yolo_model = YOLO(self.yolo_path) # yolo 로드 - self.sam_model = None + self.sam_model = SAM(self.sam_path) # None self.predictor = None self.alpha_channel = None @@ -61,21 +62,33 @@ def object_detection(self, img): return bbox def segment_from_points(self, user_points, user_labels, save_path=None): - input_point = np.array(user_points) + '''input_point = np.array(user_points) input_label = np.array(user_labels) masks_points, scores, logits = self.predictor.predict( point_coords=input_point, point_labels=input_label, multimask_output=False, - ) + )''' + filtered_points, filtered_labels = self.remove_points_in_bboxes(user_points, user_labels, bbox) + logging.info(f'2차 마스킹 - 사용자 입력 필터링: from {len(user_points)}개 to {len(filtered_points)}개') + + results = self.sam_model.predict(points=filtered_points, labels=filtered_labels) + mask_points = results[0].masks.data + + masks_np = np.zeros(mask_points.shape[-2:], dtype=np.uint8) + for mask in mask_points: + mask_np = mask.cpu().numpy().astype(np.uint8) * 255 + mask_np = mask_np.squeeze() + masks_np = cv2.bitwise_or(masks_np, mask_np) - mask_points_uint8 = (masks_points[0] * 255).astype(np.uint8) + # mask_points_uint8 = (masks_points[0] * 255).astype(np.uint8) # cv2.imwrite(save_path, mask_points_uint8) - logging.info('2차 마스킹 - 사용자 입력 점 세그먼트 완료.') - return mask_points_uint8 + logging.info(f'2차 마스킹 - 사용자 입력에 대해 {len(mask_points)}개 영역으로 세그먼트 완료.') + return masks_np def segment_from_boxes(self, image, bbox, save_path=None): + ''' input_boxes = torch.tensor(bbox, device=self.predictor.device) transformed_boxes = self.predictor.transform.apply_boxes_torch(input_boxes, image.shape[:2]) masks, _, _ = self.predictor.predict_torch( @@ -84,14 +97,17 @@ def segment_from_boxes(self, image, bbox, save_path=None): boxes=transformed_boxes, multimask_output=False, ) - masks_np = np.zeros(masks.shape[-2:], dtype=np.uint8) - for mask in masks: + ''' + mask_boxes = self.sam_model.predict(bboxes=bbox) + + masks_np = np.zeros(mask_boxes.shape[-2:], dtype=np.uint8) + for mask in mask_boxes: mask_np = mask.cpu().numpy().astype(np.uint8) * 255 mask_np = mask_np.squeeze() masks_np = cv2.bitwise_or(masks_np, mask_np) # cv2.imwrite(save_path, masks_np) - logging.info('1차 마스킹 - 바운딩 박스 세그먼트 완료.') + logging.info(f'1차 마스킹 - 바운딩 박스 {mask_boxes}개 세그먼트 완료.') return masks_np def inpainting(self, image, mask_total): @@ -121,8 +137,8 @@ def process(self, img_bytes, user_points, user_labels, extension='jpg'): image = cv2.imdecode(buffer, cv2.IMREAD_UNCHANGED) ### ready - self.load_sam_model() - self.predictor.set_image(image) + #self.load_sam_model() + #self.predictor.set_image(image) masks_total = np.zeros(image.shape[:2], dtype=np.uint8) ### 1차: Object Detection & Segment by Box @@ -130,13 +146,15 @@ def process(self, img_bytes, user_points, user_labels, extension='jpg'): if len(bbox) > 0: masks_by_box = self.segment_from_boxes(image, bbox, save_path=None) # 'test_images/seg_box.png' masks_total = cv2.bitwise_or(masks_total, masks_by_box) - logging.info( - f"1차 마스킹 후 shape 점검: YOLOv11 감지된 영역 shape: {masks_by_box.shape}, 이미지 영역 shape: {image.shape}") # (1893, 1577, 3) (1893, 1577) - + #logging.info( f"1차 마스킹 후 shape 점검: YOLOv11 감지된 영역 shape: {masks_by_box.shape}, 이미지 영역 shape: {image.shape}") # (1893, 1577, 3) (1893, 1577) + else: + masks_by_box = None ### 2차: points arguments by User & Segment by Points if len(user_points) > 0: mask_by_point = self.segment_from_points(user_points, user_labels, save_path=None) # save_path='test_images/seg_points.png' masks_total = cv2.bitwise_or(masks_total, mask_by_point) + else: + mask_by_point = None # cv2.imwrite('test_images/mask_total.png', masks_total) if isinstance(masks_total, np.ndarray): image_output = self.inpainting(image, masks_total) @@ -151,7 +169,17 @@ def process(self, img_bytes, user_points, user_labels, extension='jpg'): _, mask_bytes = cv2.imencode("." + extension, masks_total.astype(np.uint8)) _, result_bytes = cv2.imencode("." + extension, image_output) - # return 0, 0, 0 - return io.BytesIO(input_bytes), io.BytesIO(mask_bytes), io.BytesIO(result_bytes) + if mask_by_point is not None and masks_by_box is not None: + _, mask_by_point_bytes = cv2.imencode("." + extension, mask_by_point.astype(np.uint8)) + _, mask_by_box_bytes = cv2.imencode("." + extension, masks_by_box.astype(np.uint8)) + return io.BytesIO(input_bytes), io.BytesIO(mask_bytes), io.BytesIO(result_bytes), io.BytesIO(mask_by_box_bytes), io.BytesIO(mask_by_point_bytes) + elif mask_by_point is not None: + _, mask_by_point_bytes = cv2.imencode("." + extension, mask_by_point.astype(np.uint8)) + return io.BytesIO(input_bytes), io.BytesIO(mask_bytes), io.BytesIO(result_bytes), None, io.BytesIO(mask_by_point_bytes) + elif masks_by_box is not None: + _, mask_by_box_bytes = cv2.imencode("." + extension, masks_by_box.astype(np.uint8)) + return io.BytesIO(input_bytes), io.BytesIO(mask_bytes), io.BytesIO(result_bytes), io.BytesIO(mask_by_box_bytes), None + else: + return io.BytesIO(input_bytes), io.BytesIO(mask_bytes), io.BytesIO(result_bytes), None, None diff --git a/app/main.py b/app/main.py index cf6b286..bdc21ec 100644 --- a/app/main.py +++ b/app/main.py @@ -66,7 +66,7 @@ def upload_image_to_s3(file_bytes, file_path): s3_client.upload_fileobj(file_bytes, BUCKET_NAME, file_path) -def download_model_from_s3(yolo_path: str = 'models/yolo11_best.pt', sam_path: str = 'models/sam_vit_h_4b8939.pth'): +def download_model_from_s3(yolo_path: str = 'models/yolo11_best.pt', sam_path: str = "models/mobile_sam.pt"): # models/sam_vit_h_4b8939.pth dest_dir = f'../' # 모델을 저장할 컨테이너 내 경로 try: yolo_full_path = dest_dir+yolo_path @@ -82,7 +82,6 @@ def download_model_from_s3(yolo_path: str = 'models/yolo11_best.pt', sam_path: s else: logger.info(f'SAM models already exists at {dest_dir}') - logger.info(f"Files in 'models' Dir: {os.listdir(dest_dir)}") except Exception as e: print(f'Failed to download model: {e}') From 8d856362e2b778924d6ec602e54086066975ca19 Mon Sep 17 00:00:00 2001 From: semnisem Date: Tue, 19 Nov 2024 21:44:19 +0900 Subject: [PATCH 10/17] =?UTF-8?q?HotFix:=20=EC=84=B8=EA=B7=B8=EB=A8=BC?= =?UTF-8?q?=ED=8A=B8=20=EB=AA=A8=EB=8D=B8=20=EB=B3=80=EA=B2=BD=20-=20?= =?UTF-8?q?=EA=B2=BD=EB=9F=89=ED=99=94?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 모델 경로 변경 --- app/main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/app/main.py b/app/main.py index bdc21ec..0b975d7 100644 --- a/app/main.py +++ b/app/main.py @@ -120,7 +120,7 @@ async def processShape(request: Request): logger.info(f"시용자 입력 이미지({s3_key}) 다운로드 및 전처리 완료") # aiProcessor = AIProcessor(yolo_path="./models/yolo11_best.pt", sam_path="./models/sam_vit_h_4b8939.pth") # local - aiProcessor = AIProcessor(yolo_path="../models/yolo11_best.pt", sam_path="../models/sam_vit_h_4b8939.pth") # server + aiProcessor = AIProcessor(yolo_path="../models/yolo11_best.pt", sam_path="../models/mobile_sam.pt") # server img_input_bytes, img_mask_bytes, img_output_bytes = aiProcessor.process(img_bytes=corrected_img_bytes, user_points=point_list, user_labels=label_list) From 8b141d57a4e12345760c6110fa515714b643b31f Mon Sep 17 00:00:00 2001 From: semnisem Date: Tue, 19 Nov 2024 22:20:27 +0900 Subject: [PATCH 11/17] =?UTF-8?q?HotFix:=20AI=20=EC=97=90=EB=9F=AC=20?= =?UTF-8?q?=EC=88=98=EC=A0=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 로직 에러 --- app/AIProcessor.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/app/AIProcessor.py b/app/AIProcessor.py index da2ab17..ea9e04e 100644 --- a/app/AIProcessor.py +++ b/app/AIProcessor.py @@ -23,7 +23,7 @@ def __init__(self, yolo_path: str, sam_path: str, device: torch.device = torch.d self.height = 0 self.width = 0 - def remove_points_in_bboxes(point_list, label_list, bbox_list): + def remove_points_in_bboxes(self, point_list, label_list, bbox_list): def is_point_in_bbox(p, b): x, y = p x_min, y_min, x_max, y_max = b @@ -61,7 +61,7 @@ def object_detection(self, img): logging.info(f'1차 마스킹 - 바운딩 박스 생성 완료, {len(bbox)}개') return bbox - def segment_from_points(self, user_points, user_labels, save_path=None): + def segment_from_points(self, image, user_points, user_labels, bbox, save_path=None): '''input_point = np.array(user_points) input_label = np.array(user_labels) @@ -73,7 +73,7 @@ def segment_from_points(self, user_points, user_labels, save_path=None): filtered_points, filtered_labels = self.remove_points_in_bboxes(user_points, user_labels, bbox) logging.info(f'2차 마스킹 - 사용자 입력 필터링: from {len(user_points)}개 to {len(filtered_points)}개') - results = self.sam_model.predict(points=filtered_points, labels=filtered_labels) + results = self.sam_model.predict(source=image, points=filtered_points, labels=filtered_labels) mask_points = results[0].masks.data masks_np = np.zeros(mask_points.shape[-2:], dtype=np.uint8) @@ -98,7 +98,8 @@ def segment_from_boxes(self, image, bbox, save_path=None): multimask_output=False, ) ''' - mask_boxes = self.sam_model.predict(bboxes=bbox) + results = self.sam_model.predict(source=image, bboxes=bbox) + mask_boxes = results[0].masks.data masks_np = np.zeros(mask_boxes.shape[-2:], dtype=np.uint8) for mask in mask_boxes: @@ -107,7 +108,7 @@ def segment_from_boxes(self, image, bbox, save_path=None): masks_np = cv2.bitwise_or(masks_np, mask_np) # cv2.imwrite(save_path, masks_np) - logging.info(f'1차 마스킹 - 바운딩 박스 {mask_boxes}개 세그먼트 완료.') + logging.info(f'1차 마스킹 - 바운딩 박스 {len(mask_boxes)}개 세그먼트 완료.') return masks_np def inpainting(self, image, mask_total): @@ -144,14 +145,14 @@ def process(self, img_bytes, user_points, user_labels, extension='jpg'): ### 1차: Object Detection & Segment by Box bbox = self.object_detection(image) if len(bbox) > 0: - masks_by_box = self.segment_from_boxes(image, bbox, save_path=None) # 'test_images/seg_box.png' + masks_by_box = self.segment_from_boxes(image, bbox, save_path=None) # 'test_images/seg_box.png' masks_total = cv2.bitwise_or(masks_total, masks_by_box) #logging.info( f"1차 마스킹 후 shape 점검: YOLOv11 감지된 영역 shape: {masks_by_box.shape}, 이미지 영역 shape: {image.shape}") # (1893, 1577, 3) (1893, 1577) else: masks_by_box = None ### 2차: points arguments by User & Segment by Points if len(user_points) > 0: - mask_by_point = self.segment_from_points(user_points, user_labels, save_path=None) # save_path='test_images/seg_points.png' + mask_by_point = self.segment_from_points(image, user_points, user_labels, bbox, save_path=None) # save_path='test_images/seg_points.png' masks_total = cv2.bitwise_or(masks_total, mask_by_point) else: mask_by_point = None From 53ea3dcc84d1d8f7e5d8deed2687ee693ac3a0c5 Mon Sep 17 00:00:00 2001 From: semnisem Date: Tue, 19 Nov 2024 22:21:58 +0900 Subject: [PATCH 12/17] =?UTF-8?q?Fix:=20=EC=B6=94=EA=B0=80=20AI=20?= =?UTF-8?q?=EC=A4=91=EA=B0=84=EA=B3=BC=EC=A0=95=20=EC=9D=B4=EB=AF=B8?= =?UTF-8?q?=EC=A7=80=20=EC=97=85=EB=A1=9C=EB=93=9C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- app/main.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/app/main.py b/app/main.py index 0b975d7..66482b8 100644 --- a/app/main.py +++ b/app/main.py @@ -42,6 +42,8 @@ def create_file_path(obj_path, extension): paths = {"input_path": f"{dir_path}/{file_id}.input.{extension}", "mask_path": f"{dir_path}/{file_id}.mask.{extension}", "output_path": f"{dir_path}/{file_id}.output.{extension}", + "one": f"{dir_path}/{file_id}.mask_b.{extension}", + "two": f"{dir_path}/{file_id}.mask_p.{extension}", "extension": extension} return paths @@ -119,9 +121,9 @@ async def processShape(request: Request): corrected_img_bytes = ImageManager.correct_rotation(img_bytes, paths['extension']) logger.info(f"시용자 입력 이미지({s3_key}) 다운로드 및 전처리 완료") - # aiProcessor = AIProcessor(yolo_path="./models/yolo11_best.pt", sam_path="./models/sam_vit_h_4b8939.pth") # local + # aiProcessor = AIProcessor(yolo_path='/Users/semin/models/yolo11_best.pt', sam_path='/Users/semin/models/mobile_sam.pt') # local aiProcessor = AIProcessor(yolo_path="../models/yolo11_best.pt", sam_path="../models/mobile_sam.pt") # server - img_input_bytes, img_mask_bytes, img_output_bytes = aiProcessor.process(img_bytes=corrected_img_bytes, + img_input_bytes, img_mask_bytes, img_output_bytes, one, two = aiProcessor.process(img_bytes=corrected_img_bytes, user_points=point_list, user_labels=label_list) logger.info("AI 필기 제거 프로세스 완료") @@ -129,6 +131,8 @@ async def processShape(request: Request): upload_image_to_s3(img_input_bytes, paths["input_path"]) upload_image_to_s3(img_mask_bytes, paths["mask_path"]) upload_image_to_s3(img_output_bytes, paths["output_path"]) + upload_image_to_s3(one, paths["one"]) + upload_image_to_s3(two, paths["two"]) logger.info("AI 필기 제거 결과 이미지 업로드 완료") return JSONResponse(content={"message": "File processed successfully", "path": paths}) From 9c5fe7bb71ffa5458a96c9403c5215a6eb781dc6 Mon Sep 17 00:00:00 2001 From: semnisem Date: Wed, 20 Nov 2024 01:24:59 +0900 Subject: [PATCH 13/17] =?UTF-8?q?Fix:=20=EB=B3=80=EA=B2=BD=20=EC=82=AC?= =?UTF-8?q?=EC=9A=A9=EC=9E=90=20=EC=9E=85=EB=A0=A5=20=EC=A0=90=20=EC=84=B8?= =?UTF-8?q?=EA=B7=B8=EB=A9=98=ED=8C=85?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 세그멘팅 방식을 사각형 바운더리박스화 --- app/AIProcessor.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/app/AIProcessor.py b/app/AIProcessor.py index ea9e04e..4136160 100644 --- a/app/AIProcessor.py +++ b/app/AIProcessor.py @@ -73,7 +73,8 @@ def segment_from_points(self, image, user_points, user_labels, bbox, save_path= filtered_points, filtered_labels = self.remove_points_in_bboxes(user_points, user_labels, bbox) logging.info(f'2차 마스킹 - 사용자 입력 필터링: from {len(user_points)}개 to {len(filtered_points)}개') - results = self.sam_model.predict(source=image, points=filtered_points, labels=filtered_labels) + # results = self.sam_model.predict(source=image, points=filtered_points, labels=filtered_labels) + results = self.sam_model.predict(source=image, bboxes=[user_points]) mask_points = results[0].masks.data masks_np = np.zeros(mask_points.shape[-2:], dtype=np.uint8) @@ -82,7 +83,6 @@ def segment_from_points(self, image, user_points, user_labels, bbox, save_path= mask_np = mask_np.squeeze() masks_np = cv2.bitwise_or(masks_np, mask_np) - # mask_points_uint8 = (masks_points[0] * 255).astype(np.uint8) # cv2.imwrite(save_path, mask_points_uint8) logging.info(f'2차 마스킹 - 사용자 입력에 대해 {len(mask_points)}개 영역으로 세그먼트 완료.') return masks_np From eaf80e0cc84d112c5538462b6c4ab9c5505bb1c9 Mon Sep 17 00:00:00 2001 From: semnisem Date: Wed, 20 Nov 2024 01:26:21 +0900 Subject: [PATCH 14/17] =?UTF-8?q?HotFix:=20=EB=B2=84=EA=B7=B8=20=EC=88=98?= =?UTF-8?q?=EC=A0=95=20422=20=EC=97=90=EB=9F=AC?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 중간 마스킹 과정 이미지 업로드 --- app/main.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/app/main.py b/app/main.py index 66482b8..98eeb41 100644 --- a/app/main.py +++ b/app/main.py @@ -131,8 +131,10 @@ async def processShape(request: Request): upload_image_to_s3(img_input_bytes, paths["input_path"]) upload_image_to_s3(img_mask_bytes, paths["mask_path"]) upload_image_to_s3(img_output_bytes, paths["output_path"]) - upload_image_to_s3(one, paths["one"]) - upload_image_to_s3(two, paths["two"]) + if one is not None: + upload_image_to_s3(one, paths["one"]) + if two is not None: + upload_image_to_s3(two, paths["two"]) logger.info("AI 필기 제거 결과 이미지 업로드 완료") return JSONResponse(content={"message": "File processed successfully", "path": paths}) From 03b3890a2dc28bad6d6c33b1bb48e747b8d265f2 Mon Sep 17 00:00:00 2001 From: semnisem Date: Wed, 20 Nov 2024 01:50:28 +0900 Subject: [PATCH 15/17] =?UTF-8?q?HotFix:=20=EB=B2=84=EA=B7=B8=20=EC=88=98?= =?UTF-8?q?=EC=A0=95=20=EC=9E=85=EB=A0=A5=20=ED=95=84=ED=84=B0=EB=A7=81=20?= =?UTF-8?q?=EC=A0=9C=EA=B1=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- app/AIProcessor.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/app/AIProcessor.py b/app/AIProcessor.py index 4136160..1462522 100644 --- a/app/AIProcessor.py +++ b/app/AIProcessor.py @@ -58,7 +58,6 @@ def object_detection(self, img): results = self.yolo_model.predict(source=img, imgsz=640, device=self.device, iou=0.3, conf=0.3) bbox = results[0].boxes.xyxy.tolist() - logging.info(f'1차 마스킹 - 바운딩 박스 생성 완료, {len(bbox)}개') return bbox def segment_from_points(self, image, user_points, user_labels, bbox, save_path=None): @@ -70,8 +69,8 @@ def segment_from_points(self, image, user_points, user_labels, bbox, save_path= point_labels=input_label, multimask_output=False, )''' - filtered_points, filtered_labels = self.remove_points_in_bboxes(user_points, user_labels, bbox) - logging.info(f'2차 마스킹 - 사용자 입력 필터링: from {len(user_points)}개 to {len(filtered_points)}개') + # filtered_points, filtered_labels = self.remove_points_in_bboxes(user_points, user_labels, bbox) + # logging.info(f'2차 마스킹 - 사용자 입력 필터링: from {len(user_points)}개 to {len(filtered_points)}개') # results = self.sam_model.predict(source=image, points=filtered_points, labels=filtered_labels) results = self.sam_model.predict(source=image, bboxes=[user_points]) From 099529d3647567b9d2d9dd088616e5d298938405 Mon Sep 17 00:00:00 2001 From: semnisem Date: Wed, 20 Nov 2024 03:10:23 +0900 Subject: [PATCH 16/17] =?UTF-8?q?Fix:=20=EC=B6=94=EA=B0=80=20=EC=9B=90?= =?UTF-8?q?=ED=98=95=20=EB=82=B4=EB=B6=80=20=EB=B9=84=EC=9A=B0=EA=B8=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- app/AIProcessor.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/app/AIProcessor.py b/app/AIProcessor.py index 1462522..fbbe80b 100644 --- a/app/AIProcessor.py +++ b/app/AIProcessor.py @@ -17,7 +17,7 @@ def __init__(self, yolo_path: str, sam_path: str, device: torch.device = torch.d self.yolo_model = YOLO(self.yolo_path) # yolo 로드 self.sam_model = SAM(self.sam_path) # None self.predictor = None - + self.indices = None self.alpha_channel = None self.size = 0 self.height = 0 @@ -58,6 +58,8 @@ def object_detection(self, img): results = self.yolo_model.predict(source=img, imgsz=640, device=self.device, iou=0.3, conf=0.3) bbox = results[0].boxes.xyxy.tolist() + self.indices = [index for index, value in enumerate(results[0].boxes.cls) if value == 1.0] + logging.info(f'객체 탐지 - {self.indices} 박스에 동그라미 존재') return bbox def segment_from_points(self, image, user_points, user_labels, bbox, save_path=None): @@ -101,10 +103,15 @@ def segment_from_boxes(self, image, bbox, save_path=None): mask_boxes = results[0].masks.data masks_np = np.zeros(mask_boxes.shape[-2:], dtype=np.uint8) - for mask in mask_boxes: - mask_np = mask.cpu().numpy().astype(np.uint8) * 255 + for i, mask in enumerate(mask_boxes): + mask_np = mask.cpu().numpy().astype(np.uint8) * 255 # True는 255, False는 0으로 변환 mask_np = mask_np.squeeze() + if i in self.indices: + # kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (13, 13)) # Adjust size as needed + kernel = np.ones((11, 11), np.uint8) + mask_np = cv2.morphologyEx(mask_np, cv2.MORPH_GRADIENT, kernel) masks_np = cv2.bitwise_or(masks_np, mask_np) + # cv2.imwrite(f'mask_box{i}.jpg', masks_np) # cv2.imwrite(save_path, masks_np) logging.info(f'1차 마스킹 - 바운딩 박스 {len(mask_boxes)}개 세그먼트 완료.') From 91d37cebe63ae3e3591ee7d9f5b674ab63cdd15b Mon Sep 17 00:00:00 2001 From: semnisem Date: Wed, 20 Nov 2024 03:10:34 +0900 Subject: [PATCH 17/17] =?UTF-8?q?Fix:=20=EB=B3=80=EA=B2=BD=20=EC=9D=B8?= =?UTF-8?q?=ED=8E=98=EC=9D=B8=ED=8C=85=20=EB=B0=A9=EC=8B=9D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- app/AIProcessor.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/app/AIProcessor.py b/app/AIProcessor.py index fbbe80b..fc80705 100644 --- a/app/AIProcessor.py +++ b/app/AIProcessor.py @@ -118,7 +118,12 @@ def segment_from_boxes(self, image, bbox, save_path=None): return masks_np def inpainting(self, image, mask_total): - inpainted_image = cv2.inpaint(image.copy(), mask_total, 10, cv2.INPAINT_TELEA) + # inpainted_image = cv2.inpaint(image.copy(), mask_total, 10, cv2.INPAINT_TELEA) + print(mask_total.shape) # (1893, 1577) with 0 or 255 + # print(image.shape) # (1893, 1577, 3) + + inpainted_image = image.copy() + inpainted_image[mask_total == 255] = [255, 255, 255] final_image = cv2.convertScaleAbs(inpainted_image, alpha=1.5, beta=10) # cv2.imwrite('test_images/inpainted_init.png', inpainted_image) # cv2.imwrite('test_images/inpainted_final.png', final_image)