Update app.py

app.py

@@ -30,9 +30,7 @@ def process_class_list(classes_string: str):
 def model_inference(img, prob_threshold, classes_to_show):
     feature_extractor = AutoFeatureExtractor.from_pretrained(f"hustvl/yolos-small-dwr")
     model = YolosForObjectDetection.from_pretrained(f"hustvl/yolos-small-dwr")
-
     img = Image.fromarray(img)
-
     pixel_values = feature_extractor(img, return_tensors="pt").pixel_values
 
     with torch.no_grad():
@@ -40,11 +38,9 @@ def model_inference(img, prob_threshold, classes_to_show):
 
     probas = outputs.logits.softmax(-1)[0, :, :-1]
     keep = probas.max(-1).values > prob_threshold
-
     target_sizes = torch.tensor(img.size[::-1]).unsqueeze(0)
     postprocessed_outputs = feature_extractor.post_process(outputs, target_sizes)
     bboxes_scaled = postprocessed_outputs[0]["boxes"]
-
     classes_list = process_class_list(classes_to_show)
     res_img = plot_results(img, probas[keep], bboxes_scaled[keep], model, classes_list)
 
@@ -59,16 +55,10 @@ def plot_results(pil_img, prob, boxes, model, classes_list):
     for p, (xmin, ymin, xmax, ymax), c in zip(prob, boxes.tolist(), colors):
         cl = p.argmax()
         object_class = model.config.id2label[cl.item()]
-
         if len(classes_list) > 0:
             if object_class not in classes_list:
                 continue
-
-        ax.add_patch(
-            plt.Rectangle(
-                (xmin, ymin), xmax - xmin, ymax - ymin, fill=False, color=c, linewidth=3
-            )
-        )
+        ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, fill=False, color=c, linewidth=3))
         text = f"{object_class}: {p[cl]:0.2f}"
         ax.text(xmin, ymin, text, fontsize=15, bbox=dict(facecolor="yellow", alpha=0.5))
     plt.axis("off")
@@ -90,7 +80,7 @@ title = """Object Detection"""
 # example_list = [["examples/" + example] for example in os.listdir("examples")]
 example_list = [["carplane.webp"]]
 
-image_in = gr.components.Image(label="Upload an image")
+image_in = [gr.components.Image(label="Upload an image")]
 image_out = gr.components.Image()
 classes_to_show = gr.components.Textbox(placeholder="e.g. car, dog", label="Classes to filter (leave empty to detect all classes)")
 prob_threshold_slider = gr.components.Slider(minimum=0, maximum=1.0, step=0.01, value=0.7, label="Probability Threshold")