Spaces:
Sleeping
Sleeping
File size: 3,221 Bytes
51588cf |
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 84 85 86 87 |
from PIL import Image
import matplotlib.pyplot as plt
import io
# COCO classes
CLASSES = [
'N/A', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A',
'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse',
'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack',
'umbrella', 'N/A', 'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
'skateboard', 'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass',
'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich',
'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake',
'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table', 'N/A',
'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard',
'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A',
'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier',
'toothbrush'
]
COLORS = [
[0.000, 0.447, 0.741],
[0.850, 0.325, 0.098],
[0.929, 0.694, 0.125],
[0.494, 0.184, 0.556],
[0.466, 0.674, 0.188],
[0.301, 0.745, 0.933],
]
# Update JSON dictionary with rounded values and classes
def generate_output_json(json_dict):
json_dict['scores'] = [round(score, 3) for score in json_dict['scores']]
json_dict['boxes'] = [[round(coord, 3) for coord in box] for box in json_dict['boxes']]
json_dict['labels'] = [CLASSES[label] for label in json_dict['labels']]
return json_dict
# Generate matplotlib figure from prediction scores and boxes
def generate_output_figure(image_path, results, threshold):
pil_img = Image.open(image_path)
plt.figure(figsize=(16, 10))
plt.imshow(pil_img)
ax = plt.gca()
colors = COLORS * 100
print("\t Detailed information...")
for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
#box = [round(i, 2) for i in box]
print(
f"\t\t Detected {label} with confidence "
f"{score} at location {box}"
)
if score > threshold:
c = COLORS[hash(label) % len(COLORS)]
ax.add_patch(
plt.Rectangle((box[0], box[1]), box[2] - box[0], box[3] - box[1], fill=False, color=c, linewidth=3)
)
text = f"{label}: {score:0.2f}"
ax.text(box[0], box[1], text, fontsize=15, bbox=dict(facecolor="yellow", alpha=0.5))
plt.axis("off")
return plt.gcf()
# Generate PIL image from matplotlib figure
def generate_output_image(output_figure):
# Convert matplotlib figure to PIL image
#output_figure = plt.gcf()
buf = io.BytesIO()
output_figure.savefig(buf, bbox_inches="tight")
buf.seek(0)
output_pil_img = Image.open(buf)
return output_pil_img
def generate_gradio_outputs(image_path, response_dict, threshold):
output_json = generate_output_json(response_dict)
output_figure = generate_output_figure(image_path, output_json, threshold)
output_pil_img = generate_output_image(output_figure)
return output_json, output_pil_img
|