vanhai123 commited on
Commit
6a4da27
·
verified ·
1 Parent(s): a283bc6

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +36 -0
app.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import BlipProcessor, BlipForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM
3
+ from PIL import Image
4
+ import torch
5
+
6
+ # Load model 1: English image captioning
7
+ blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
8
+ blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
9
+
10
+ # Load model 2: Translate EN → VI
11
+ translator_tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
12
+ translator_model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
13
+ lang_code = "vie_Latn"
14
+
15
+ def caption_translate(image):
16
+ # Step 1: Get English caption
17
+ inputs = blip_processor(image, return_tensors="pt")
18
+ out = blip_model.generate(**inputs)
19
+ eng_caption = blip_processor.decode(out[0], skip_special_tokens=True)
20
+
21
+ # Step 2: Translate to Vietnamese
22
+ inputs = translator_tokenizer(eng_caption, return_tensors="pt", src_lang="eng_Latn", tgt_lang=lang_code)
23
+ translated = translator_model.generate(**inputs, max_length=100)
24
+ vi_caption = translator_tokenizer.decode(translated[0], skip_special_tokens=True)
25
+
26
+ return f"📷 Mô tả: {vi_caption}\n\n(English: {eng_caption})"
27
+
28
+ iface = gr.Interface(
29
+ fn=caption_translate,
30
+ inputs=gr.Image(type="pil"),
31
+ outputs="text",
32
+ title="🧠 AI Mô Tả Hình Ảnh Bằng Tiếng Việt",
33
+ description="Upload ảnh, hệ thống sẽ mô tả nội dung bằng tiếng Việt bằng cách kết hợp 2 mô hình: caption → translate"
34
+ )
35
+
36
+ iface.launch()