Update app.py

app.py

@@ -1,47 +1,25 @@
 import gradio as gr
-from transformers import AutoImageProcessor, ResNetForImageClassification, YolosFeatureExtractor, YolosForObjectDetection
+from transformers import YolosFeatureExtractor, YolosForObjectDetection
 import torch
 
+feature_extractor = YolosFeatureExtractor.from_pretrained('hustvl/yolos-small')
+model = YolosForObjectDetection.from_pretrained('hustvl/yolos-small')
 
-
-
-def detect(image1, image2):
-    ### Image 1, the object ### 
-    processor = AutoImageProcessor.from_pretrained("microsoft/resnet-50")
-    model = ResNetForImageClassification.from_pretrained("microsoft/resnet-50")
-    
-    inputs = processor(image1, return_tensors="pt")
-    
-    with torch.no_grad():
-        logits = model(**inputs).logits
-    
-    # model predicts one of the 1000 ImageNet classes
-    predicted_label = logits.argmax(-1).item()
-    object_label = model.config.id2label[predicted_label]
-    
-    ### Image 2, object detections ###
-    from PIL import Image
-    import requests
-    
-    feature_extractor = YolosFeatureExtractor.from_pretrained('hustvl/yolos-small')
-    model = YolosForObjectDetection.from_pretrained('hustvl/yolos-small')
-    
-    inputs = feature_extractor(images=image2, return_tensors="pt")
+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 object_label, bboxes
-
-
+    return outputs
 
 
 demo = gr.Interface(
     fn=detect,
-    inputs=[gr.inputs.Image(label="Object to detect"), gr.inputs.Image(label="Image to detect object in")],
-    outputs=["text", "text"],
+    inputs=[gr.inputs.Image(label="Input image")],
+    outputs=["text"],
     title="Object Counts in Image"
 )