1
0
Fork 0
AutoPercenty2/img_trans/main.py

61 lines
3.0 KiB
Python

import cv2
from src.text_detection_with_paddle import detect_text
from src.removal import process_image_based_on_text_area
from src.translation import translate_texts_translatepy
from src.text_insertion import insert_text
from src.mini_crop_process import resize_and_crop_image
from src.rotate_img import flip_and_rotate_image
from src.inpainting import inpaint_text, inpaint_text_with_deepfill # 기존 인페인팅 메서드
import logging, sys, os
def main(image_path, img_name, thumb_or_trans, checkpoint_model_path):
# 이미지에서 텍스트 인식
detected_texts = detect_text(image_path)
print("Detected texts:", detected_texts)
# 처리된 이미지 경로 설정
Thumb_image_path = f'img/{img_name}_Thumb.jpg'
processed_image_path = f'img/{img_name}_processed.jpg'
translated_image_path = f'img/{img_name}_translated.jpg'
inpainted_image_path = f'img/{img_name}_inpainted.jpg'
# 텍스트 영역에 따라 이미지 처리: 크롭 또는 인페이팅
# crop_and_inpaint(image_path, detected_texts, processed_image_path)
# print(f"Image processed and saved to {processed_image_path}")
if thumb_or_trans: # 썸네일 처리
# 썸네일 생성 로직
print("썸네일 생성 처리 시작")
# 크롭과 인페인팅 결정 및 처리 함수 호출
removal_img = process_image_based_on_text_area(image_path, detected_texts, processed_image_path, threshold = 0.1)
# 추가 처리: 이미지 회전 및 좌우 반전
# removal_img = flip_and_rotate_image(removal_img)
cv2.imwrite(Thumb_image_path, removal_img)
print(f"썸네일 이미지 생성완료: {Thumb_image_path}")
else: # 이미지 번역
# 2단계: 인페이팅 기법으로 텍스트 제거
image = cv2.imread(image_path)
# inpainted_image = inpaint_text(image, detected_texts)
inpainted_image = inpaint_text_with_deepfill(image, detected_texts, checkpoint_model_path)
cv2.imwrite(inpainted_image_path, inpainted_image)
# 이미지 번역 처리
print("이미지 번역 처리 시작")
# 번역된 텍스트 추출 및 삽입
texts_to_translate = [text for _, text, _, _, _ in detected_texts]
print(f"원본 텍스트: {texts_to_translate}")
translated_texts = translate_texts_translatepy(texts_to_translate)
print(f"번역 텍스트: {texts_to_translate}")
translated_image = insert_text(inpainted_image_path, translated_texts, detected_texts)
cv2.imwrite(translated_image_path, translated_image)
print(f"번역된 이미지 생성완료 : {translated_image_path}")
if __name__ == "__main__":
base_path = getattr(sys, '_MEIPASS', os.path.dirname(os.path.abspath(__file__)))
checkpoint_model_path = os.path.join(base_path, "models", "checkpoint")
img_name = "1"
image_path = f'img/{img_name}.jpg'
thumb_or_trans = False # True for thumbnail, False for translation
main(image_path, img_name, thumb_or_trans, checkpoint_model_path)