Download utils.py from torchFlow/detect-it: direct link, hf CLI and curl.
- Browser
- Download file 2.38 kB
-
https://huggingface.co/spaces/torchFlow/detect-it/resolve/main/utils.py
- Command line
-
hf download hf://spaces/torchFlow/detect-it/utils.py
-
curl -L -o utils.py https://huggingface.co/spaces/torchFlow/detect-it/resolve/main/utils.py
2.38 kB
| # Utilities | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image | |
| def load_img(img_path, input_shape): | |
| # Loading image | |
| image = cv2.imread(img_path) | |
| image_height, image_width = image.shape[:2] | |
| Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) | |
| input_height, input_width = input_shape[2:] | |
| image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| resized = cv2.resize(image_rgb, (input_width, input_height)) # image resized as req by onnx model | |
| input_image = resized / 255.0 # scaling image | |
| input_image = input_image.transpose(2,0,1) # dim rearranged as req by onnx (batch_size, channel, height, width) | |
| input_tensor = input_image[np.newaxis, :, :, :].astype(np.float32) | |
| return image, image_height, image_width, input_height, input_width, input_tensor | |
| def nms(boxes, scores, iou_threshold): | |
| # Sort by score | |
| sorted_indices = np.argsort(scores)[::-1] | |
| keep_boxes = [] | |
| while sorted_indices.size > 0: | |
| # Pick the last box | |
| box_id = sorted_indices[0] | |
| keep_boxes.append(box_id) | |
| # Compute IoU of the picked box with the rest | |
| ious = compute_iou(boxes[box_id, :], boxes[sorted_indices[1:], :]) | |
| # Remove boxes with IoU over the threshold | |
| keep_indices = np.where(ious < iou_threshold)[0] | |
| # print(keep_indices.shape, sorted_indices.shape) | |
| sorted_indices = sorted_indices[keep_indices + 1] | |
| return keep_boxes | |
| def compute_iou(box, boxes): | |
| # Compute xmin, ymin, xmax, ymax for both boxes | |
| xmin = np.maximum(box[0], boxes[:, 0]) | |
| ymin = np.maximum(box[1], boxes[:, 1]) | |
| xmax = np.minimum(box[2], boxes[:, 2]) | |
| ymax = np.minimum(box[3], boxes[:, 3]) | |
| # Compute intersection area | |
| intersection_area = np.maximum(0, xmax - xmin) * np.maximum(0, ymax - ymin) | |
| # Compute union area | |
| box_area = (box[2] - box[0]) * (box[3] - box[1]) | |
| boxes_area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) | |
| union_area = box_area + boxes_area - intersection_area | |
| # Compute IoU | |
| iou = intersection_area / union_area | |
| return iou | |
| def xywh2xyxy(x): | |
| # Convert bounding box (x, y, w, h) to bounding box (x1, y1, x2, y2) | |
| y = np.copy(x) | |
| y[..., 0] = x[..., 0] - x[..., 2] / 2 | |
| y[..., 1] = x[..., 1] - x[..., 3] / 2 | |
| y[..., 2] = x[..., 0] + x[..., 2] / 2 | |
| y[..., 3] = x[..., 1] + x[..., 3] / 2 | |
| return y |