import gradio as gr
from transformers import YolosFeatureExtractor, YolosForObjectDetection
import torch

feature_extractor = YolosFeatureExtractor.from_pretrained('hustvl/yolos-small')
model = YolosForObjectDetection.from_pretrained('hustvl/yolos-small')

def detect(image): 
    inputs = feature_extractor(images=image, return_tensors="pt")
    outputs = model(**inputs)
    
    # model predicts bounding boxes and corresponding COCO classes
    logits = outputs.logits
    bboxes = outputs.pred_boxes

    return outputs


demo = gr.Interface(
    fn=detect,
    inputs=[gr.inputs.Image(label="Input image")],
    outputs=["text"],
    title="Object Counts in Image"
)

demo.launch()