123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051 |
- from boundingbox import BoundingBox
- import cv2
- import numpy as np
- def preprocess(img, input_shape, letter_box=True):
- if letter_box:
- img_h, img_w, _ = img.shape
- new_h, new_w = input_shape[0], input_shape[1]
- offset_h, offset_w = 0, 0
- if (new_w / img_w) <= (new_h / img_h):
- new_h = int(img_h * new_w / img_w)
- offset_h = (input_shape[0] - new_h) // 2
- else:
- new_w = int(img_w * new_h / img_h)
- offset_w = (input_shape[1] - new_w) // 2
- resized = cv2.resize(img, (new_w, new_h))
- img = np.full((input_shape[0], input_shape[1], 3), 127, dtype=np.uint8)
- img[offset_h:(offset_h + new_h), offset_w:(offset_w + new_w), :] = resized
- else:
- img = cv2.resize(img, (input_shape[1], input_shape[0]))
- img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
- img = img.transpose((2, 0, 1)).astype(np.float32)
- img /= 255.0
- return img
- def postprocess(num_dets, det_boxes, det_scores, det_classes, img_w, img_h, input_shape, letter_box=True):
- boxes = det_boxes[0, :num_dets[0][0]] / np.array([input_shape[0], input_shape[1], input_shape[0], input_shape[1]], dtype=np.float32)
- scores = det_scores[0, :num_dets[0][0]]
- classes = det_classes[0, :num_dets[0][0]].astype(np.int)
- old_h, old_w = img_h, img_w
- offset_h, offset_w = 0, 0
- if letter_box:
- if (img_w / input_shape[1]) >= (img_h / input_shape[0]):
- old_h = int(input_shape[0] * img_w / input_shape[1])
- offset_h = (old_h - img_h) // 2
- else:
- old_w = int(input_shape[1] * img_h / input_shape[0])
- offset_w = (old_w - img_w) // 2
- boxes = boxes * np.array([old_w, old_h, old_w, old_h], dtype=np.float32)
- if letter_box:
- boxes -= np.array([offset_w, offset_h, offset_w, offset_h], dtype=np.float32)
- boxes = boxes.astype(np.int)
- detected_objects = []
- for box, score, label in zip(boxes, scores, classes):
- detected_objects.append(BoundingBox(label, score, box[0], box[2], box[1], box[3], img_w, img_h))
- return detected_objects
|