processing.py 2.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051
  1. from boundingbox import BoundingBox
  2. import cv2
  3. import numpy as np
  4. def preprocess(img, input_shape, letter_box=True):
  5. if letter_box:
  6. img_h, img_w, _ = img.shape
  7. new_h, new_w = input_shape[0], input_shape[1]
  8. offset_h, offset_w = 0, 0
  9. if (new_w / img_w) <= (new_h / img_h):
  10. new_h = int(img_h * new_w / img_w)
  11. offset_h = (input_shape[0] - new_h) // 2
  12. else:
  13. new_w = int(img_w * new_h / img_h)
  14. offset_w = (input_shape[1] - new_w) // 2
  15. resized = cv2.resize(img, (new_w, new_h))
  16. img = np.full((input_shape[0], input_shape[1], 3), 127, dtype=np.uint8)
  17. img[offset_h:(offset_h + new_h), offset_w:(offset_w + new_w), :] = resized
  18. else:
  19. img = cv2.resize(img, (input_shape[1], input_shape[0]))
  20. img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
  21. img = img.transpose((2, 0, 1)).astype(np.float32)
  22. img /= 255.0
  23. return img
  24. def postprocess(num_dets, det_boxes, det_scores, det_classes, img_w, img_h, input_shape, letter_box=True):
  25. 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)
  26. scores = det_scores[0, :num_dets[0][0]]
  27. classes = det_classes[0, :num_dets[0][0]].astype(np.int)
  28. old_h, old_w = img_h, img_w
  29. offset_h, offset_w = 0, 0
  30. if letter_box:
  31. if (img_w / input_shape[1]) >= (img_h / input_shape[0]):
  32. old_h = int(input_shape[0] * img_w / input_shape[1])
  33. offset_h = (old_h - img_h) // 2
  34. else:
  35. old_w = int(input_shape[1] * img_h / input_shape[0])
  36. offset_w = (old_w - img_w) // 2
  37. boxes = boxes * np.array([old_w, old_h, old_w, old_h], dtype=np.float32)
  38. if letter_box:
  39. boxes -= np.array([offset_w, offset_h, offset_w, offset_h], dtype=np.float32)
  40. boxes = boxes.astype(np.int)
  41. detected_objects = []
  42. for box, score, label in zip(boxes, scores, classes):
  43. detected_objects.append(BoundingBox(label, score, box[0], box[2], box[1], box[3], img_w, img_h))
  44. return detected_objects