Spaces:
Sleeping
Sleeping
File size: 3,152 Bytes
d4c3acc 246f207 ceb95cf c21a752 ceb95cf 52b4046 ceb95cf 52b4046 ceb95cf 6564a3a ceb95cf d4c3acc 246f207 dba19bd f6b7865 a0333be f6b7865 a0333be f6b7865 e14364d ceb95cf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
import os
import gradio as gr
from transformers import pipeline, DetrForObjectDetection, DetrConfig, DetrImageProcessor
import numpy as np
import cv2
from PIL import Image
def draw_detections(image, detections):
# Convert PIL image to a numpy array
np_image = np.array(image)
# Convert RGB to BGR for OpenCV
np_image = cv2.cvtColor(np_image, cv2.COLOR_RGB2BGR)
for detection in detections:
# Extract scores, labels, and bounding boxes properly
score = detection['score']
label = detection['labels']
box = detection['boxes'] # Make sure 'boxes' data structure matches expected in terms of naming and indexing
x_min, y_min, x_max, y_max = map(int, [box[0], box[1], box[2], box[3]])
cv2.rectangle(np_image, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
cv2.putText(np_image, f'{label} {score:.2f}', (x_min, max(y_min - 10, 0)),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
# Convert BGR to RGB for displaying
final_image = cv2.cvtColor(np_image, cv2.COLOR_BGR2RGB)
# Convert the numpy array to PIL Image
final_pil_image = Image.fromarray(final_image)
return final_pil_image
# Initialize objects from transformers
config = DetrConfig.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", config=config)
image_processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
od_pipe = pipeline(task='object-detection', model=model, image_processor=image_processor)
##def get_pipeline_prediction(pil_image):
## try:
# Run the object detection pipeline
## pipeline_output = od_pipe(pil_image)
# Draw the detection results on the image
## processed_image = draw_detections(pil_image, pipeline_output)
# Provide both the image and the JSON detection results
## return processed_image, pipeline_output
## except Exception as e:
# Log the error
## print(f"An error occurred: {str(e)}")
# Return a message and an empty JSON
## return pil_image, {"error": str(e)}
def get_pipeline_prediction(pil_image):
try:
# Run the object detection pipeline
pipeline_output = od_pipe(pil_image)
# Debugging: print the keys in the output dictionary
if pipeline_output:
print("Keys available in the detection output:", pipeline_output[0].keys())
# Draw the detection results on the image
processed_image = draw_detections(pil_image, pipeline_output)
# Provide both the image and the JSON detection results
return processed_image, pipeline_output
except Exception as e:
# Log the error
print(f"An error occurred: {str(e)}")
# Return a message and an empty JSON
return pil_image, {"error": str(e)}
demo = gr.Interface(
fn=get_pipeline_prediction,
inputs=gr.Image(label="Input image", type="pil"),
outputs=[
gr.Image(label="Annotated Image"),
gr.JSON(label="Detected Objects")
]
)
demo.launch()
|