import os import re import time from typing import List, Dict, Any import torch import gradio as gr from PIL import Image from transformers import pipeline from datasets import load_dataset # Vocabulary dictionary covering Office-Home dataset classes + common COCO # household/office items DETR emits. Single-word keys are matched per-token in # captions and detection labels; multi-word keys (e.g. "dining table") are # matched as phrases. VOCAB_DICT = { # --- Furniture --- "chair": {"japanese": "いす", "romaji": "isu", "korean": "의자", "romanization": "uija"}, "table": {"japanese": "テーブル", "romaji": "teeburu", "korean": "테이블", "romanization": "teibeul"}, "dining table": {"japanese": "ダイニングテーブル", "romaji": "dainingu teeburu", "korean": "식탁", "romanization": "siktak"}, "desk": {"japanese": "机", "romaji": "tsukue", "korean": "책상", "romanization": "chaeksang"}, "bed": {"japanese": "ベッド", "romaji": "beddo", "korean": "침대", "romanization": "chimdae"}, "couch": {"japanese": "ソファ", "romaji": "sofa", "korean": "소파", "romanization": "sopa"}, "sofa": {"japanese": "ソファ", "romaji": "sofa", "korean": "소파", "romanization": "sopa"}, "shelf": {"japanese": "棚", "romaji": "tana", "korean": "선반", "romanization": "seonban"}, "curtain": {"japanese": "カーテン", "romaji": "kaaten", "korean": "커튼", "romanization": "keoteun"}, "file cabinet": {"japanese": "ファイルキャビネット", "romaji": "fairu kyabinetto", "korean": "파일 캐비닛", "romanization": "pail kaebinit"}, # --- Lighting / electrical --- "lamp": {"japanese": "ランプ", "romaji": "ranpu", "korean": "램프", "romanization": "raempeu"}, "desk lamp": {"japanese": "デスクランプ", "romaji": "desuku ranpu", "korean": "책상 램프", "romanization": "chaeksang raempeu"}, "lamp shade": {"japanese": "ランプシェード", "romaji": "ranpu sheedo", "korean": "램프 갓", "romanization": "raempeu gat"}, "fan": {"japanese": "扇風機", "romaji": "senpuuki", "korean": "선풍기", "romanization": "seonpunggi"}, "battery": {"japanese": "電池", "romaji": "denchi", "korean": "배터리", "romanization": "baeteori"}, "candle": {"japanese": "ろうそく", "romaji": "rousoku", "korean": "양초", "romanization": "yangcho"}, # --- Computing / electronics --- "laptop": {"japanese": "ノートパソコン", "romaji": "nooto pasokon", "korean": "노트북", "romanization": "noteubuk"}, "computer": {"japanese": "コンピュータ", "romaji": "konpyuuta", "korean": "컴퓨터", "romanization": "keompyuteo"}, "monitor": {"japanese": "モニター", "romaji": "monitaa", "korean": "모니터", "romanization": "moniteo"}, "keyboard": {"japanese": "キーボード", "romaji": "kiibodo", "korean": "키보드", "romanization": "kibodeu"}, "mouse": {"japanese": "マウス", "romaji": "mausu", "korean": "마우스", "romanization": "mauseu"}, "printer": {"japanese": "プリンター", "romaji": "purintaa", "korean": "프린터", "romanization": "peurinteo"}, "webcam": {"japanese": "ウェブカメラ", "romaji": "webu kamera", "korean": "웹캠", "romanization": "wepkaem"}, "speaker": {"japanese": "スピーカー", "romaji": "supiikaa", "korean": "스피커", "romanization": "seupikeo"}, "tv": {"japanese": "テレビ", "romaji": "terebi", "korean": "텔레비전", "romanization": "tellebijeon"}, "television": {"japanese": "テレビ", "romaji": "terebi", "korean": "텔레비전", "romanization": "tellebijeon"}, "remote": {"japanese": "リモコン", "romaji": "rimokon", "korean": "리모컨", "romanization": "rimokeon"}, "radio": {"japanese": "ラジオ", "romaji": "rajio", "korean": "라디오", "romanization": "radio"}, "phone": {"japanese": "電話", "romaji": "denwa", "korean": "전화", "romanization": "jeonhwa"}, "telephone": {"japanese": "電話", "romaji": "denwa", "korean": "전화", "romanization": "jeonhwa"}, "cell phone": {"japanese": "携帯電話", "romaji": "keitai denwa", "korean": "휴대폰", "romanization": "hyudaepon"}, "calculator": {"japanese": "電卓", "romaji": "dentaku", "korean": "계산기", "romanization": "gyesangi"}, "clock": {"japanese": "時計", "romaji": "tokei", "korean": "시계", "romanization": "sigye"}, "alarm clock": {"japanese": "目覚まし時計", "romaji": "mezamashi dokei", "korean": "알람 시계", "romanization": "allam sigye"}, # --- Stationery / office supplies --- "pen": {"japanese": "ペン", "romaji": "pen", "korean": "펜", "romanization": "pen"}, "pencil": {"japanese": "鉛筆", "romaji": "enpitsu", "korean": "연필", "romanization": "yeonpil"}, "marker": {"japanese": "マーカー", "romaji": "maakaa", "korean": "마커", "romanization": "makeo"}, "eraser": {"japanese": "消しゴム", "romaji": "keshigomu", "korean": "지우개", "romanization": "jiugae"}, "ruler": {"japanese": "定規", "romaji": "jougi", "korean": "자", "romanization": "ja"}, "scissors": {"japanese": "はさみ", "romaji": "hasami", "korean": "가위", "romanization": "gawi"}, "notebook": {"japanese": "ノート", "romaji": "nooto", "korean": "공책", "romanization": "gongchaek"}, "book": {"japanese": "本", "romaji": "hon", "korean": "책", "romanization": "chaek"}, "folder": {"japanese": "フォルダ", "romaji": "foruda", "korean": "폴더", "romanization": "poldeo"}, "clipboard": {"japanese": "クリップボード", "romaji": "kurippu boodo", "korean": "클립보드", "romanization": "keullipbodeu"}, "calendar": {"japanese": "カレンダー", "romaji": "karendaa", "korean": "달력", "romanization": "dallyeok"}, "paper clip": {"japanese": "クリップ", "romaji": "kurippu", "korean": "종이 클립", "romanization": "jongi keullip"}, "push pin": {"japanese": "画びょう", "romaji": "gabyou", "korean": "압정", "romanization": "apjeong"}, "exit sign": {"japanese": "出口表示", "romaji": "deguchi hyouji", "korean": "출구 표지", "romanization": "chulgu pyoji"}, # --- Kitchen / dining --- "mug": {"japanese": "マグカップ", "romaji": "magu kappu", "korean": "머그컵", "romanization": "meogeukeop"}, "cup": {"japanese": "カップ", "romaji": "kappu", "korean": "컵", "romanization": "keop"}, "wine glass": {"japanese": "ワイングラス", "romaji": "wain gurasu", "korean": "와인 잔", "romanization": "wain jan"}, "bottle": {"japanese": "ボトル", "romaji": "botoru", "korean": "병", "romanization": "byeong"}, "bowl": {"japanese": "ボウル", "romaji": "bouru", "korean": "그릇", "romanization": "geureut"}, "fork": {"japanese": "フォーク", "romaji": "fooku", "korean": "포크", "romanization": "pokeu"}, "spoon": {"japanese": "スプーン", "romaji": "supuun", "korean": "숟가락", "romanization": "sutgarak"}, "knife": {"japanese": "ナイフ", "romaji": "naifu", "korean": "칼", "romanization": "kal"}, "kettle": {"japanese": "やかん", "romaji": "yakan", "korean": "주전자", "romanization": "jujeonja"}, "pan": {"japanese": "フライパン", "romaji": "furaipan", "korean": "팬", "romanization": "paen"}, "oven": {"japanese": "オーブン", "romaji": "oobun", "korean": "오븐", "romanization": "obeun"}, "microwave": {"japanese": "電子レンジ", "romaji": "denshi renji", "korean": "전자레인지", "romanization": "jeonjareinji"}, "toaster": {"japanese": "トースター", "romaji": "toosutaa", "korean": "토스터", "romanization": "toseuteo"}, "refrigerator": {"japanese": "冷蔵庫", "romaji": "reizouko", "korean": "냉장고", "romanization": "naengjanggo"}, "sink": {"japanese": "流し", "romaji": "nagashi", "korean": "싱크대", "romanization": "singkeudae"}, "soda": {"japanese": "ソーダ", "romaji": "sooda", "korean": "탄산음료", "romanization": "tansaneumnyo"}, # --- Bathroom --- "toothbrush": {"japanese": "歯ブラシ", "romaji": "ha burashi", "korean": "칫솔", "romanization": "chitsol"}, "toilet": {"japanese": "トイレ", "romaji": "toire", "korean": "화장실", "romanization": "hwajangsil"}, # --- Tools / hardware --- "hammer": {"japanese": "ハンマー", "romaji": "hanmaa", "korean": "망치", "romanization": "mangchi"}, "drill": {"japanese": "ドリル", "romaji": "doriru", "korean": "드릴", "romanization": "deuril"}, "screwdriver": {"japanese": "ドライバー", "romaji": "doraibaa", "korean": "드라이버", "romanization": "deuraibeo"}, "bucket": {"japanese": "バケツ", "romaji": "baketsu", "korean": "양동이", "romanization": "yangdongi"}, "mop": {"japanese": "モップ", "romaji": "moppu", "korean": "대걸레", "romanization": "daegeolle"}, "trash can": {"japanese": "ゴミ箱", "romaji": "gomibako", "korean": "쓰레기통", "romanization": "sseuregitong"}, # --- Personal items / clothing --- "backpack": {"japanese": "リュックサック", "romaji": "ryukku sakku", "korean": "백팩", "romanization": "baekpaek"}, "handbag": {"japanese": "ハンドバッグ", "romaji": "hando baggu", "korean": "핸드백", "romanization": "haendeubaek"}, "suitcase": {"japanese": "スーツケース", "romaji": "suutsu keesu", "korean": "여행 가방", "romanization": "yeohaeng gabang"}, "umbrella": {"japanese": "傘", "romaji": "kasa", "korean": "우산", "romanization": "usan"}, "glasses": {"japanese": "眼鏡", "romaji": "megane", "korean": "안경", "romanization": "angyeong"}, "tie": {"japanese": "ネクタイ", "romaji": "nekutai", "korean": "넥타이", "romanization": "nektai"}, "helmet": {"japanese": "ヘルメット", "romaji": "herumetto", "korean": "헬멧", "romanization": "helmet"}, "sneakers": {"japanese": "スニーカー", "romaji": "suniikaa", "korean": "운동화", "romanization": "undonghwa"}, "flipflops": {"japanese": "ビーチサンダル", "romaji": "biichi sandaru", "korean": "슬리퍼", "romanization": "seullipeo"}, "bike": {"japanese": "自転車", "romaji": "jitensha", "korean": "자전거", "romanization": "jajeongeo"}, # --- Decor / misc --- "flower": {"japanese": "花", "romaji": "hana", "korean": "꽃", "romanization": "kkot"}, "plant": {"japanese": "植物", "romaji": "shokubutsu", "korean": "식물", "romanization": "singmul"}, "potted plant": {"japanese": "鉢植え", "romaji": "hachi-ue", "korean": "화분", "romanization": "hwabun"}, "vase": {"japanese": "花瓶", "romaji": "kabin", "korean": "꽃병", "romanization": "kkotbyeong"}, "toy": {"japanese": "おもちゃ", "romaji": "omocha", "korean": "장난감", "romanization": "jangnangam"}, "teddy bear": {"japanese": "テディベア", "romaji": "tedi bea", "korean": "곰인형", "romanization": "gominhyeong"}, "postit": {"japanese": "付箋", "romaji": "fusen", "korean": "포스트잇", "romanization": "poseuteuit"}, "hairdryer": {"japanese": "ドライヤー", "romaji": "doraiyaa", "korean": "드라이어", "romanization": "deuraieo"}, } # Pre-split single-word vs multi-word keys for efficient matching _SINGLE_WORD_KEYS = {k for k in VOCAB_DICT if " " not in k} _MULTI_WORD_KEYS = [k for k in VOCAB_DICT if " " in k] # Device setup USE_GPU = torch.cuda.is_available() DEVICE = 0 if USE_GPU else -1 TORCH_DTYPE = torch.float16 if USE_GPU else None # Load models globally as pipelines caption_pipeline = pipeline( "image-to-text", model="Salesforce/blip-image-captioning-base", device=DEVICE, ) def generate_caption(image: Image.Image) -> str: """Generate caption using BLIP image-to-text pipeline.""" out = caption_pipeline(image, max_new_tokens=50) if isinstance(out, list) and out and "generated_text" in out[0]: return out[0]["generated_text"] return "" detection_pipeline = pipeline( "object-detection", model="facebook/detr-resnet-50", device=DEVICE, ) # Load up to 10 sample images from flwrlabs/office-home for one-click testing. # Filter to Office-Home classes whose label matches a key in VOCAB_DICT, so the # samples are guaranteed to produce vocab the app can actually translate. Dedupe # by class to maximize variety. Streaming mode avoids downloading the full dataset. SAMPLE_DIR = "sample_images" MAX_STREAM_SCAN = 2000 # safety cap so we don't iterate forever def load_sample_images(n: int = 10) -> List[str]: paths: List[str] = [] try: os.makedirs(SAMPLE_DIR, exist_ok=True) ds = load_dataset("flwrlabs/office-home", split="train", streaming=True) class_names = ds.features["label"].names if "label" in ds.features else [] seen_classes: set = set() for i, example in enumerate(ds): if len(paths) >= n or i >= MAX_STREAM_SCAN: break img = example.get("image") label_idx = example.get("label") if img is None or label_idx is None or not class_names: continue raw_label = class_names[label_idx] normalized = raw_label.lower().replace("_", "") if not any(vocab_key in normalized for vocab_key in VOCAB_DICT): continue if raw_label in seen_classes: continue seen_classes.add(raw_label) path = os.path.join(SAMPLE_DIR, f"sample_{len(paths):02d}_{raw_label}.jpg") img.convert("RGB").save(path, "JPEG") paths.append(path) except Exception as e: print(f"Could not load sample images from flwrlabs/office-home: {e}") return paths SAMPLE_PATHS = load_sample_images(10) def clean_text(text: str) -> str: """Clean and normalize text.""" return re.sub(r"[^a-zA-Z\s]", "", text.lower()).strip() def extract_vocab_from_caption(caption: str) -> List[str]: """Extract vocab from caption text. Single-word keys match per-token; multi-word keys are matched as phrases.""" cleaned = clean_text(caption) tokens = set(cleaned.split()) matches = {k for k in _SINGLE_WORD_KEYS if k in tokens} matches.update(k for k in _MULTI_WORD_KEYS if k in cleaned) return list(matches) def extract_vocab_from_detection(detection_results: List[Dict]) -> List[str]: """Extract vocab from detection labels (often multi-word, e.g. 'dining table').""" matches = set() for res in detection_results: if res.get("score", 0) <= 0.5: continue label = res.get("label", "").lower() if label in VOCAB_DICT: matches.add(label) continue for token in label.split(): if token in _SINGLE_WORD_KEYS: matches.add(token) return list(matches) def translate_term(term: str, lang: str) -> Dict[str, str]: """Translate term using dictionary.""" if term not in VOCAB_DICT: return {"translation": "translation unavailable", "romanization": "N/A"} entry = VOCAB_DICT[term] if lang == "Japanese": return {"translation": entry["japanese"], "romanization": entry["romaji"]} elif lang == "Korean": return {"translation": entry["korean"], "romanization": entry["romanization"]} return {"translation": term, "romanization": "N/A"} def generate_flashcard_table(vocab_list: List[str], lang: str) -> List[List[str]]: """Generate flashcard table.""" table = [["English", f"{lang} Translation", "Romanization", "Source"]] for term in vocab_list: trans = translate_term(term, lang) table.append([term, trans["translation"], trans["romanization"], "extracted"]) return table def compute_comparison_stats( caption_vocab: List[str], detection_vocab: List[str], caption_time: float, detection_time: float, detection_results: List[Dict], ) -> str: """Compute comparison statistics.""" overlap = set(caption_vocab) & set(detection_vocab) avg_conf = sum(r["score"] for r in detection_results) / len(detection_results) if detection_results else 0.0 stats = f""" Captioning Vocab Terms: {len(caption_vocab)} Detection Vocab Terms: {len(detection_vocab)} Overlapping Terms: {len(overlap)} Caption Output Length: {len(' '.join(caption_vocab))} Detection Output Length: {len(detection_vocab)} Average Detection Confidence: {avg_conf:.2f} Captioning Time: {caption_time:.2f}s Detection Time: {detection_time:.2f}s Conclusion: {'Captioning' if len(caption_vocab) > len(detection_vocab) else 'Detection'} provided more vocabulary terms. """ return stats.strip() def process_image(image: Image.Image, language: str): """Main processing function.""" if image is None: return "No image uploaded.", [], [], "No image." # Algorithm 1: Captioning start = time.time() try: caption = generate_caption(image) except Exception as e: caption = f"Captioning failed: {e}" caption_time = time.time() - start # Algorithm 2: Detection start = time.time() try: detection_results = detection_pipeline(image) except Exception as e: detection_results = [] detection_time = time.time() - start # NLP: Extract vocab caption_vocab = extract_vocab_from_caption(caption) detection_vocab = extract_vocab_from_detection(detection_results) all_vocab = list(set(caption_vocab + detection_vocab)) # Flashcard table flashcard_table = generate_flashcard_table(all_vocab, language) # Comparison stats stats = compute_comparison_stats(caption_vocab, detection_vocab, caption_time, detection_time, detection_results) return caption, detection_results, flashcard_table, stats # Gradio Interface with gr.Blocks(title="Multimodal Language Flashcard Generator") as demo: gr.Markdown("# Multimodal Language Flashcard Generator") gr.Markdown("Upload an image, select a language, and generate flashcards with captioning and object detection.") with gr.Row(): image_input = gr.Image(type="pil", label="Upload Image") lang_input = gr.Dropdown(["Japanese", "Korean"], label="Target Language", value="Japanese") if SAMPLE_PATHS: gr.Examples( examples=[[p] for p in SAMPLE_PATHS], inputs=[image_input], label="Sample images from flwrlabs/office-home (click one to load)", ) generate_btn = gr.Button("Generate Flashcards") with gr.Row(): caption_output = gr.Textbox(label="Image Caption", lines=2) detection_output = gr.Dataframe(label="Object Detection Results", headers=["Label", "Score", "Box"]) flashcard_output = gr.Dataframe(label="Flashcard Table", headers=["English", "Translation", "Romanization", "Source"]) stats_output = gr.Textbox(label="Comparison Statistics", lines=8) generate_btn.click( fn=process_image, inputs=[image_input, lang_input], outputs=[caption_output, detection_output, flashcard_output, stats_output], ) if __name__ == "__main__": demo.launch()