Save output results as json

app.py

@@ -1,10 +1,5 @@
 import os
-# os.system("sudo apt-get update && sudo apt-get install -y git")
-# os.system("sudo apt-get -y install pybind11-dev")
-# os.system("git clone https://github.com/facebookresearch/detectron2.git")
-# os.system("pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 torchaudio==0.9.0 -f https://download.pytorch.org/whl/torch_stable.html")
 os.system("cd detectron2 && pip install detectron2-0.6-cp310-cp310-linux_x86_64.whl")
-# os.system("pip3 install torch torchvision torchaudio")
 os.system("pip install deepspeed==0.7.0")
 
 import site
@@ -12,10 +7,11 @@ from importlib import reload
 reload(site)
 
 from PIL import Image
+from io import BytesIO
 import argparse
 import sys
 import numpy as np
-import cv2
+import torch
 import gradio as gr
 
 from detectron2.config import get_cfg
@@ -80,11 +76,37 @@ def setup_cfg(args):
 
 def predict(image_file):
     image_array = np.array(image_file)[:, :, ::-1]  # BGR
-    _, visualized_output = dense_captioning_demo.run_on_image(image_array)
-    visualized_output.save(os.path.join(os.getcwd(), "output.jpg"))
-    output_image = cv2.imread(os.path.join(os.getcwd(), "output.jpg"))
-    output_image = cv2.cvtColor(output_image, cv2.COLOR_BGR2RGB)
-    return Image.fromarray(output_image)
+    predictions, visualized_output = dense_captioning_demo.run_on_image(image_array)
+    buffer = BytesIO()
+    visualized_output.fig.savefig(buffer, format='png')
+    buffer.seek(0)
+    detections = {}
+    predictions = predictions["instances"].to(torch.device("cpu"))
+
+    for box, description, score in zip(
+        predictions.pred_boxes,
+        predictions.pred_object_descriptions.data,
+        predictions.scores,
+    ):
+        if description not in detections:
+            detections[description] = []
+        detections[description].append(
+            {
+                "xmin": float(box[0]),
+                "ymin": float(box[1]),
+                "xmax": float(box[2]),
+                "ymax": float(box[3]),
+                "score": float(score),
+            }
+        )
+
+    output = {
+        "dense_captioning_results": {
+            "detections": detections,
+        }
+    }
+
+    return Image.open(buffer), output
 
 
 
@@ -102,7 +124,7 @@ demo = gr.Interface(
     title="Dense Captioning - GRiT",
     fn=predict,
     inputs=gr.Image(type='pil', label="Original Image"),
-    outputs=gr.Image(type="pil",label="Output Image"),
+    outputs=[gr.Image(type="pil",label="Output Image"), "json"],
     examples=["example_1.jpg", "example_2.jpg"],
 )