Object_detection / app_json.py
Gabolozano's picture
Rename app.py to app_json.py
edeafe6 verified
raw
history blame
1.51 kB
import os
import gradio as gr
from transformers import pipeline, DetrForObjectDetection, DetrConfig, DetrImageProcessor
# Initialize the configuration for DetrForObjectDetection
config = DetrConfig.from_pretrained("facebook/detr-resnet-50")
# Create the model for object detection using the specified configuration
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", config=config)
# Initialize the image processor for DETR
image_processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
# Initialize the object detection pipeline with the model and image processor
od_pipe = pipeline(task='object-detection', model=model, image_processor=image_processor)
def get_pipeline_prediction(pil_image):
# Run the object detection pipeline on the input image
pipeline_output = od_pipe(pil_image)
# You might need to implement or adjust the rendering function based on the `pipeline_output`
# The `render_results_in_image` function is assumed here to draw bounding boxes and labels on the input image,
# but you'll need to define it according to your specific needs.
# For now, the output is directly returned since the question doesn't define `render_results_in_image`.
return pipeline_output
demo = gr.Interface(
fn=get_pipeline_prediction,
inputs=gr.Image(label="Input image",
type="pil"),
outputs=gr.JSON(label="Detected objects") # Adjusted to show JSON output if rendering function is not defined
)
demo.launch()