Update app.py

output label - korean by checkbox

app.py

@@ -12,6 +12,13 @@ owl_processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16-ensemb
 # dino_processor = AutoProcessor.from_pretrained("IDEA-Research/grounding-dino-base")
 # dino_model = AutoModelForZeroShotObjectDetection.from_pretrained("IDEA-Research/grounding-dino-base").to("cuda")
 
+english_candidate_labels = ["hat", "sunglass", "hair band", "glove", "arm sleeve", "watch", "singlet", "t-shirts", "energy gel", "half pants", "socks", "shoes", "ear phone"]
+korean_candidate_labels = ["모자", "썬글라스", "헤어밴드", "장갑", "팔토시", "시계", "싱글렛", "티셔츠", "에너지젤", "쇼츠바지", "양말", "신발", "이어폰"]
+english_candidate_labels_string = ",".join(english_candidate_labels)
+
+# 영문 레이블을 한글 레이블로 매칭하는 딕셔너리 생성
+label_mapping = dict(zip(english_candidate_labels, korean_candidate_labels))
+
 @spaces.GPU
 def infer(img, text_queries, score_threshold, model):
   
@@ -65,13 +72,19 @@ def query_image(img, text_queries, owl_threshold, flag_output_korean):
     owl_output = infer(img, text_queries, owl_threshold, "owl")
     # dino_output = infer(img, text_queries, dino_threshold, "dino")
 
-
+    # add - check flag output korean
+    owl_output_final = []
+    if flag_output_korean:
+        for box, label in owl_output:
+            kor_label = label_mapping[label]
+            owl_output_final.append(box, kor_label)
+        
+    else:
+        owl_output_final = owl_output
+        
     # return (img, owl_output), (img, dino_output)
-    return (img, owl_output)
+    return (img, owl_output_final)
 
-english_candidate_labels = ["hat", "sunglass", "hair band", "glove", "arm sleeve", "watch", "singlet", "t-shirts", "energy gel", "half pants", "socks", "shoes", "ear phone"]
-korean_candidate_labels = ["모자", "썬글라스", "헤어밴드", "장갑", "팔토시", "시계", "싱글렛", "티셔츠", "에너지젤", "쇼츠바지", "양말", "신발", "이어폰"]
-english_candidate_labels_string = ",".join(english_candidate_labels)
 
 owl_threshold = gr.Slider(0, 1, value=0.16, label="OWL Threshold")
 # dino_threshold = gr.Slider(0, 1, value=0.12, label="Grounding DINO Threshold")