Spaces:
Running
Running
File size: 3,151 Bytes
e3641b1 |
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 |
import cv2
import numpy as np
import json
rotation_map = {
0: None,
90: cv2.ROTATE_90_COUNTERCLOCKWISE,
180: cv2.ROTATE_180,
270: cv2.ROTATE_90_CLOCKWISE
}
class NumpyEncoder(json.JSONEncoder):
def default(self, obj):
if isinstance(obj, np.ndarray):
return obj.tolist()
return json.JSONEncoder.default(self, obj)
def draw_bboxes(image, bounding_boxes, boxes_id, scores):
image_with_boxes = image.copy()
for bbox, bbox_id, score in zip(bounding_boxes, boxes_id, scores):
x1, y1, x2, y2 = bbox
cv2.rectangle(image_with_boxes, (x1, y1), (x2, y2), (128, 128, 0), 2)
label = f'#{bbox_id}: {score:.2f}'
(label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
label_x = x1
label_y = y1 - 5 if y1 > 20 else y1 + 20
# Draw a filled rectangle as the background for the label
cv2.rectangle(image_with_boxes, (x1, label_y - label_height - 5),
(x1 + label_width, label_y + 5), (128, 128, 0), cv2.FILLED)
cv2.putText(image_with_boxes, label, (label_x, label_y),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
return image_with_boxes
def pad_image(image: np.ndarray, aspect_ratio: float) -> np.ndarray:
# Get the current aspect ratio of the image
image_height, image_width = image.shape[:2]
current_aspect_ratio = image_width / image_height
left_pad = 0
top_pad = 0
# Determine whether to pad horizontally or vertically
if current_aspect_ratio < aspect_ratio:
# Pad horizontally
target_width = int(aspect_ratio * image_height)
pad_width = target_width - image_width
left_pad = pad_width // 2
right_pad = pad_width - left_pad
padded_image = np.pad(image,
pad_width=((0, 0), (left_pad, right_pad), (0, 0)),
mode='constant')
else:
# Pad vertically
target_height = int(image_width / aspect_ratio)
pad_height = target_height - image_height
top_pad = pad_height // 2
bottom_pad = pad_height - top_pad
padded_image = np.pad(image,
pad_width=((top_pad, bottom_pad), (0, 0), (0, 0)),
mode='constant')
return padded_image, (left_pad, top_pad)
class VideoReader(object):
def __init__(self, file_name, rotate=0):
self.file_name = file_name
self.rotate = rotation_map[rotate]
try: # OpenCV needs int to read from webcam
self.file_name = int(file_name)
except ValueError:
pass
def __iter__(self):
self.cap = cv2.VideoCapture(self.file_name)
if not self.cap.isOpened():
raise IOError('Video {} cannot be opened'.format(self.file_name))
return self
def __next__(self):
was_read, img = self.cap.read()
if not was_read:
raise StopIteration
if self.rotate is not None:
img = cv2.rotate(img, self.rotate)
return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|