pred_classes infos

app.py

@@ -35,6 +35,16 @@ cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 # Step 4: define model
 predictor = DefaultPredictor(cfg)
 
+def get_bytes_shape_dtype(t):
+    """
+    input: tensor
+    output: 3 strings
+    """
+    t_numpy = t.cpu().numpy()
+    t_bytes = str(t_numpy.tobytes())
+    t_numpy_shape = str(t_numpy.shape)
+    t_numpy_dtype = str(t_numpy.dtype)
+    return t_bytes, t_numpy_shape, t_numpy_dtype
 
 def analyze_image(img):
     md = MetadataCatalog.get(cfg.DATASETS.TEST[0])
@@ -57,20 +67,19 @@ def analyze_image(img):
     for field in fields:
         if field == 'pred_boxes':
             boxes = output.get_fields()[field]
-            boxes_numpy = boxes.tensor.cpu().numpy()
-            boxes_bytes = str(boxes_numpy.tobytes())
-            boxes_numpy_shape = str(boxes_numpy.shape)
-            boxes_numpy_dtype = str(boxes_numpy.dtype)
+            boxes = boxes.tensor
+            boxes_bytes, boxes_numpy_shape, boxes_numpy_dtype = get_bytes_shape_dtype(boxes)
             # boxes_recover = torch.from_numpy(np.frombuffer(boxes_bytes, dtype=boxes_numpy_dtype).reshape(boxes_numpy_shape))
         elif field == 'scores':
             scores = output.get_fields()[field]
-            scores_numpy = scores.cpu().numpy()
-            scores_bytes = str(scores_numpy.tobytes())
-            scores_numpy_shape = str(scores_numpy.shape)
-            scores_numpy_dtype = str(scores_numpy.dtype)
+            scores_bytes, scores_numpy_shape, scores_numpy_dtype = get_bytes_shape_dtype(scores)
             # scores_recover = torch.from_numpy(np.frombuffer(scores_bytes, dtype=scores_numpy_dtype).reshape(scores_numpy_shape))
+        elif field == 'pred_classes':
+            pred_classes = output.get_fields()[field]
+            pred_classes_bytes, pred_classes_numpy_shape, pred_classes_numpy_dtype = get_bytes_shape_dtype(pred_classes)
+            # pred_classes_recover = torch.from_numpy(np.frombuffer(pred_classes_bytes, dtype=pred_classes_numpy_dtype).reshape(pred_classes_numpy_shape))
     
-    return result_image, num_instances, image_size, boxes_bytes, boxes_numpy_shape, boxes_numpy_dtype, scores_bytes, scores_numpy_shape, scores_numpy_dtype
+    return result_image, num_instances, image_size, boxes_bytes, boxes_numpy_shape, boxes_numpy_dtype, scores_bytes, scores_numpy_shape, scores_numpy_dtype, pred_classes_bytes, pred_classes_numpy_shape, pred_classes_numpy_dtype
     
 title = "Interactive demo: Document Layout Analysis with DiT"
 description = "Demo for Microsoft's DiT, the Document Image Transformer for state-of-the-art document understanding tasks. This particular model is fine-tuned on PubLayNet, a large dataset for document layout analysis (read more at the links below). To use it, simply upload an image or use the example image below and click 'Submit'. Results will show up in a few seconds. If you want to make the output bigger, right-click on it and select 'Open image in new tab'."
@@ -89,12 +98,15 @@ iface = gr.Interface(fn=analyze_image,
                          gr.outputs.Textbox(label="boxes numpy dtype"),
                          gr.outputs.Textbox(label="scores bytes"),
                          gr.outputs.Textbox(label="scores numpy shape"),
-                         gr.outputs.Textbox(label="scores numpy dtype")
+                         gr.outputs.Textbox(label="scores numpy dtype"),
+                         gr.outputs.Textbox(label="pred_classes bytes"),
+                         gr.outputs.Textbox(label="pred_classes numpy shape"),
+                         gr.outputs.Textbox(label="pred_classes numpy dtype")
                          ],
                      title=title,
                      description=description,
                      examples=examples,
                      article=article,
-                     css=css,
-                     enable_queue=True)
-iface.launch(debug=True, cache_examples=True)
+                     css=css
+                     )
+iface.launch(debug=True, cache_examples=True, enable_queue=True)