Update app.py

Model now returns [out_img, boxes, labels]

app.py

@@ -162,8 +162,9 @@ def run_detr(image:Image, detr_processor, detr_model, confidence_threshold: floa
                                                                     threshold=confidence_threshold)
     outputs = postprocessed_outputs[0]
     # dict{scores, logits, labels, boxes}
-    outputs = outputs['boxes']
-    return outputs
+    labels = [detr_model.config.id2label[label.item()] for label in outputs["labels"]]
+    boxes = outputs['boxes']
+    return boxes, labels
 
 
 def resize_dimensions(dimensions, target_size):
@@ -284,7 +285,6 @@ class ControlNetDepthDesignModelMulti:
         
         return design_image
 
-
 def create_demo(model):
     gr.Markdown("### demo")
     with gr.Row():
@@ -348,9 +348,9 @@ def create_demo(model):
             # -- run detr --
             # -----------------
             # clear_gpu()
-            bboxes = run_detr(out_img, detr_processor, detr_model, detr_confidence_threshold)
+            bboxes, labels = run_detr(out_img, detr_processor, detr_model, detr_confidence_threshold)
             
-        return out_img, bboxes.tolist()
+        return out_img, bboxes.tolist(), labels
 
     submit.click(on_submit, inputs=[input_image, input_text, num_steps, guidance_scale, seed, strength, detr_confidence_threshold, a_prompt, n_prompt, img_size], outputs=[design_image, bboxes])