kasper-boy commited on
Commit
744022c
·
verified ·
1 Parent(s): bd32079

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -1
app.py CHANGED
@@ -11,6 +11,10 @@ processor = DetrImageProcessor.from_pretrained('facebook/detr-resnet-101')
11
  model = DetrForObjectDetection.from_pretrained('facebook/detr-resnet-101')
12
 
13
  def object_detection(image, confidence_threshold):
 
 
 
 
14
  # Preprocess the image
15
  inputs = processor(images=image, return_tensors="pt")
16
 
@@ -56,7 +60,7 @@ def object_detection(image, confidence_threshold):
56
  # Define the Gradio interface
57
  demo = gr.Interface(
58
  fn=object_detection,
59
- inputs=[gr.Image(label="Upload an Image"), gr.Slider(minimum=0.0, maximum=1.0, label="Confidence Threshold")],
60
  outputs=[gr.Image(label="Detected Objects"), gr.Textbox(label="Detected Objects List")],
61
  title="Object Detection with DETR (ResNet-101)",
62
  description="Upload an image and get object detection results using the DETR model with a ResNet-101 backbone."
 
11
  model = DetrForObjectDetection.from_pretrained('facebook/detr-resnet-101')
12
 
13
  def object_detection(image, confidence_threshold):
14
+ # Convert the input to a PIL Image object if it's not already
15
+ if not isinstance(image, Image.Image):
16
+ image = Image.open(io.BytesIO(image))
17
+
18
  # Preprocess the image
19
  inputs = processor(images=image, return_tensors="pt")
20
 
 
60
  # Define the Gradio interface
61
  demo = gr.Interface(
62
  fn=object_detection,
63
+ inputs=[gr.Image(label="Upload an Image"), gr.Slider(minimum=0.0, maximum=1.0, label="Confidence Threshold", default=0.5)],
64
  outputs=[gr.Image(label="Detected Objects"), gr.Textbox(label="Detected Objects List")],
65
  title="Object Detection with DETR (ResNet-101)",
66
  description="Upload an image and get object detection results using the DETR model with a ResNet-101 backbone."