Update app.py

app.py

@@ -1,33 +1,30 @@
+import torch
 import gradio as gr
-from transformers import TrOCRProcessor, VisionEncoderDecoderModel
-from PIL import Image, ImageOps
+from promptcap import PromptCap
 
-# OCRモデルとプロセッサの準備
-processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-stage1")
-model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-stage1")
+# モデルの読み込み
+model = PromptCap("tifa-benchmark/promptcap-coco-vqa")
 
-# OCR実行関数
-def ocr_image(image, invert_image):
-    # チェックボックスがオンなら画像を反転
-    if invert_image:
-        image = ImageOps.mirror(image)
-    
-    # 画像をモデルに入力し、文字起こしを実行
-    pixel_values = processor(images=image, return_tensors="pt").pixel_values
-    generated_ids = model.generate(pixel_values)
-    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    return generated_text
+# CUDAが使用可能ならGPUを使用
+if torch.cuda.is_available():
+    model.cuda()
 
-# Gradioインターフェースの構築
+# 画像と質問を入力として、キャプションを生成する関数を定義
+def generate_caption(image, question):
+    prompt = f"please describe this image according to the given question: {question}"
+    # PromptCapモデルでキャプションを生成
+    caption = model.caption(prompt, image)
+    return caption
+
+# Gradioインターフェースの定義
 interface = gr.Interface(
-    fn=ocr_image,                               # 実行関数
+    fn=generate_caption,  # キャプション生成関数
     inputs=[
-        gr.Image(type="pil"),                   # カメラからの画像入力
-        gr.Checkbox(label="画像を反転")          # 画像反転のチェックボックス
+        gr.inputs.Image(type="filepath", label="Input Image"),  # 画像入力
+        gr.inputs.Textbox(label="Question")  # 質問入力
     ],
-    outputs="text",                             # 出力形式（テキスト）
-    live=False                                  # リアルタイム処理を無効化
+    outputs=gr.outputs.Textbox(label="Generated Caption")  # キャプション出力
 )
 
-# 実行ボタンの設定
+# インターフェースを起動
 interface.launch()