Spaces:
Sleeping
Sleeping
File size: 4,757 Bytes
5d92054 |
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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
"""
Created on Sat Apr 9 04:08:02 2022
@author: Admin_with ODD Team
Edited by our team : Sat Oct 5 10:00 2024
references: https://github.com/vinvino02/GLPDepth
"""
import io
import torch
import base64
from config import CONFIG
from torchvision import transforms
from matplotlib import pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from transformers import DetrForObjectDetection, DetrImageProcessor
class DETR:
def __init__(self):
self.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'
]
self.COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098],
[0.929, 0.694, 0.125], [0, 0, 1], [0.466, 0.674, 0.188],
[0.301, 0.745, 0.933]]
self.transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
self.model = DetrForObjectDetection.from_pretrained(CONFIG['detr_model_path'], revision="no_timm")
self.model.to(CONFIG['device'])
self.model.eval()
def box_cxcywh_to_xyxy(self, x):
x_c, y_c, w, h = x.unbind(1)
b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
(x_c + 0.5 * w), (y_c + 0.5 * h)]
return torch.stack(b, dim=1).to(CONFIG['device'])
def rescale_bboxes(self, out_bbox, size):
img_w, img_h = size
b = self.box_cxcywh_to_xyxy(out_bbox)
b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32).to(CONFIG['device'])
return b
def detect(self, im):
img = self.transform(im).unsqueeze(0).to(CONFIG['device'])
assert img.shape[-2] <= 1600 and img.shape[-1] <= 1600, 'Image too large'
outputs = self.model(img)
probas = outputs['logits'].softmax(-1)[0, :, :-1]
keep = probas.max(-1).values > 0.7
bboxes_scaled = self.rescale_bboxes(outputs['pred_boxes'][0, keep], im.size)
return probas[keep], bboxes_scaled
def visualize(self, im, probas, bboxes):
"""
Visualizes the detected bounding boxes and class probabilities on the image.
Parameters:
im (PIL.Image): The original input image.
probas (Tensor): Class probabilities for detected objects.
bboxes (Tensor): Bounding boxes for detected objects.
"""
# Convert image to RGB format for matplotlib
fig, ax = plt.subplots(figsize=(10, 6))
ax.imshow(im)
# Iterate over detections and draw bounding boxes and labels
for p, (xmin, ymin, xmax, ymax), color in zip(probas, bboxes, self.COLORS * 100):
# Detach tensors and convert to float
xmin, ymin, xmax, ymax = map(lambda x: x.detach().cpu().numpy().item(), (xmin, ymin, xmax, ymax))
ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,
fill=False, color=color, linewidth=3))
cl = p.argmax()
text = f'{self.CLASSES[cl]}: {p[cl].detach().cpu().numpy():0.2f}' # Detach probability as well
ax.text(xmin, ymin, text, fontsize=15, bbox=dict(facecolor='yellow', alpha=0.5))
ax.axis('off')
# Convert the plot to a PIL Image and then to bytes
canvas = FigureCanvas(fig)
buf = io.BytesIO()
canvas.print_png(buf)
buf.seek(0)
# Base64 encode the image
img_bytes = buf.getvalue()
img_base64 = base64.b64encode(img_bytes).decode('utf-8')
# Close the figure to release memory
plt.close(fig)
return img_base64
|