Update app.py

app.py

@@ -76,9 +76,20 @@ def visualize(pred_mask, image_path, work_dir):
     cv2.imwrite(output_path, visual_result)
     return output_path
 
+from googletrans import Translator
+
+# 번역 함수 추가
+def translate_to_korean(text):
+    translator = Translator()
+    try:
+        result = translator.translate(text, dest='ko', src='en')
+        return result.text
+    except:
+        return text  # 번역 실패시 원본 텍스트 반환
+
 @spaces.GPU
 def image_vision(image_input_path, prompt):
-    # 입력된 프롬프트가 한글인지 확인
+    # 한글 입력 확인
     is_korean = any(ord('가') <= ord(char) <= ord('힣') for char in prompt)
     
     image_path = image_input_path
@@ -93,18 +104,21 @@ def image_vision(image_input_path, prompt):
     }
     return_dict = model.predict_forward(**input_dict)
     print(return_dict)
-    answer = return_dict["prediction"] # the text format answer
+    answer = return_dict["prediction"]
     
-    # 한글 프롬프트인 경우 응답을 한글로 변환
+    # 한글 프롬프트인 경우 응답을 한글로 번역
     if is_korean:
-        # 기본 응답 패턴을 한글로 변환
-        answer = answer.replace("Yes", "네")
-        answer = answer.replace("No", "아니오")
-        answer = answer.replace("[SEG]", "[분할]")
+        # [SEG]는 보존하면서 나머지 텍스트만 번역
+        if '[SEG]' in answer:
+            parts = answer.split('[SEG]')
+            translated_parts = [translate_to_korean(part) for part in parts]
+            answer = '[SEG]'.join(translated_parts)
+        else:
+            answer = translate_to_korean(answer)
     
     seg_image = return_dict["prediction_masks"]
     
-    if ('[SEG]' in answer or '[분할]' in answer) and Visualizer is not None:
+    if '[SEG]' in answer and Visualizer is not None:
         pred_masks = seg_image[0]
         temp_dir = tempfile.mkdtemp()
         pred_mask = pred_masks
@@ -116,10 +130,9 @@ def image_vision(image_input_path, prompt):
 
 @spaces.GPU(duration=80)
 def video_vision(video_input_path, prompt, video_interval):
-    # 입력된 프롬프트가 한글인지 확인
+    # 한글 입력 확인
     is_korean = any(ord('가') <= ord(char) <= ord('힣') for char in prompt)
     
-    # Open the original video
     cap = cv2.VideoCapture(video_input_path)
     original_fps = cap.get(cv2.CAP_PROP_FPS)
     frame_skip_factor = video_interval
@@ -135,13 +148,16 @@ def video_vision(video_input_path, prompt, video_interval):
     prediction = result['prediction']
     print(prediction)
 
-    # 한글 프롬프트인 경우 응답을 한글로 변환
+    # 한글 프롬프트인 경우 응답을 한글로 번역
     if is_korean:
-        prediction = prediction.replace("Yes", "네")
-        prediction = prediction.replace("No", "아니오")
-        prediction = prediction.replace("[SEG]", "[분할]")
-
-    if ('[SEG]' in prediction or '[분할]' in prediction) and Visualizer is not None:
+        if '[SEG]' in prediction:
+            parts = prediction.split('[SEG]')
+            translated_parts = [translate_to_korean(part) for part in parts]
+            prediction = '[SEG]'.join(translated_parts)
+        else:
+            prediction = translate_to_korean(prediction)
+
+    if '[SEG]' in prediction and Visualizer is not None:
         _seg_idx = 0
         pred_masks = result['prediction_masks'][_seg_idx]
         seg_frames = []