Update app.py

app.py

@@ -1,25 +1,33 @@
 import gradio as gr
-from transformers import YolosFeatureExtractor, YolosForObjectDetection
+from transformers import AutoImageProcessor, AutoModelForObjectDetection
 import torch
 
-feature_extractor = YolosFeatureExtractor.from_pretrained('hustvl/yolos-small')
-model = YolosForObjectDetection.from_pretrained('hustvl/yolos-small')
+image_processor = AutoImageProcessor.from_pretrained('hustvl/yolos-small')
+model = AutoModelForObjectDetection.from_pretrained('hustvl/yolos-small')
 
 def detect(image): 
     inputs = feature_extractor(images=image, return_tensors="pt")
     outputs = model(**inputs)
+
+    # convert outputs to COCO API
+    target_sizes = torch.tensor([image.size[::-1]])
+    results = image_processor.post_process_object_detection(outputs,
+                                                            threshold=0.9,
+                                                            target_sizes=target_sizes)[0]
     
     # model predicts bounding boxes and corresponding COCO classes
-    logits = outputs.logits
-    bboxes = outputs.pred_boxes
+    #logits = outputs.logits
+    #bboxes = outputs.pred_boxes
 
-    return outputs
+    # label and the count
+    counts = {}
 
+    return results
 
 demo = gr.Interface(
     fn=detect,
     inputs=[gr.inputs.Image(label="Input image")],
-    outputs=["text"],
+    outputs=["text"]#, gr.Label(num_top_classes=10)],
     title="Object Counts in Image"
 )