Upload README.md

README.md

@@ -205,3 +205,86 @@ draw_entity_boxes_on_image(image, entities, show=True)
 Here is the annotated image:
 
 <a href="https://huggingface.co/ydshieh/kosmos-2-patch14-224/resolve/main/annotated_snowman.jpg" target="_blank"><img src="https://huggingface.co/ydshieh/kosmos-2-patch14-224/resolve/main/annotated_snowman.jpg" width="500"></a>
+
+
+## Running the Flask Server
+_flask_kosmos2.py_ shows the implementation of a Flask server for the model.
+It allowes the model to be approached as a REST API.
+
+After starting the server. You can send a POST request to `http://localhost:8005/process_prompt` with the following form data:
+- `prompt`: For example `<grounding> an image of`
+- `image`: The image file as binary data
+
+This in turn will produce a reply with the following JSON format:
+- `message`: The Kosmos-2 generated text
+- `entities`: The extracted entities
+
+An easy way to test this is through an application like Postman. Make sure the image field is set to `File`.
+
+```python
+
+from PIL import Image
+from transformers import AutoProcessor, AutoModelForVision2Seq
+from flask import Flask, request, jsonify
+import json
+
+app = Flask(__name__)
+
+model = AutoModelForVision2Seq.from_pretrained("ydshieh/kosmos-2-patch14-224", trust_remote_code=True)
+processor = AutoProcessor.from_pretrained("ydshieh/kosmos-2-patch14-224", trust_remote_code=True)
+
+
+@app.route('/process_prompt', methods=['POST'])
+def process_prompt():
+    try:
+        # Get the uploaded image data from the POST request
+        uploaded_file = request.files['image']
+        prompt = request.form.get('prompt')
+        image = Image.open(uploaded_file.stream)
+
+        print(image.size)
+
+        inputs = processor(text=prompt, images=image, return_tensors="pt")
+
+        generated_ids = model.generate(
+            pixel_values=inputs["pixel_values"],
+            input_ids=inputs["input_ids"][:, :-1],
+            attention_mask=inputs["attention_mask"][:, :-1],
+            img_features=None,
+            img_attn_mask=inputs["img_attn_mask"][:, :-1],
+            use_cache=True,
+            max_new_tokens=64,
+        )
+        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+        # By default, the generated  text is cleanup and the entities are extracted.
+        processed_text, entities = processor.post_process_generation(generated_text)
+        parsed_entities = entities_to_json(entities)
+        print(generated_text)
+        print(processed_text)
+        return jsonify({"message": processed_text, 'entities': parsed_entities})
+    except Exception as e:
+        return jsonify({"error": str(e)})
+
+
+def entities_to_json(entities):
+    result = []
+    for e in entities:
+        label = e[0]
+        box_coords = e[1]
+        box_size = e[2][0]
+        entity_result = {
+            "label": label,
+            "boundingBoxPosition": {"x": box_coords[0], "y": box_coords[1]},
+            "boundingBox": {"x_min": box_size[0], "y_min": box_size[1], "x_max": box_size[2], "y_max": box_size[3]}
+        }
+        print(entity_result)
+        result.append(entity_result)
+
+    return result
+
+
+if __name__ == '__main__':
+    app.run(host='localhost', port=8005)
+
+```