first commit

app.py

@@ -0,0 +1,22 @@
+import streamlit as st
+from PIL import Image
+import kosmos
+
+kosmos.load_model()
+st.title('KOSMOS2 Model Evaluation')
+
+uploaded_file = st.file_uploader("Choose an image...", type="jpg")
+
+if uploaded_file is not None:
+    try:
+        image = Image.open(uploaded_file)
+        st.image(image, caption='Uploaded Image.', use_column_width=True)
+        st.write("")
+        prompt = st.text_input(f"Detect the main object in the image. The image is a")
+        if st.button('Classify'):
+            st.write(f"User: {prompt}")
+            with st.spinner('Processing...'):
+                label = kosmos.single_image_classification(image, prompt, 50)
+            st.write(f"Model: {label}")
+    except Exception as e:
+        st.error(f"An error occurred: {e}")

kosmos.py

@@ -0,0 +1,33 @@
+from transformers import AutoProcessor, AutoModelForVision2Seq
+from PIL import Image
+from os import path
+from torchvision.transforms import ToTensor
+
+model_id = "microsoft/kosmos-2-patch14-224"
+model = None
+processor = None
+
+def load_model():
+    global model, processor
+    model = AutoModelForVision2Seq.from_pretrained(model_id)
+    processor = AutoProcessor.from_pretrained(model_id)
+
+def single_image_classification(image, prompt="", max_new_tokens=30):
+    if model is None or processor is None:
+        load_model()
+    inputs = processor(text=prompt, images=image, return_tensors="pt")
+    
+    generated_ids = model.generate(
+        pixel_values=inputs["pixel_values"],
+        input_ids=inputs["input_ids"],
+        attention_mask=inputs["attention_mask"],
+        image_embeds=None,
+        image_embeds_position_mask=inputs["image_embeds_position_mask"],
+        use_cache=True,
+        max_new_tokens=max_new_tokens,
+    )
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+
+    processed_text = processor.post_process_generation(generated_text, cleanup_and_extract=True)
+
+    return processed_text

requirements.txt

@@ -0,0 +1,4 @@
+transformers
+torch
+torchvision
+streamlit