|
|
|
@ -15,7 +15,7 @@ import urllib |
|
|
|
|
import uuid |
|
|
|
|
from pathlib import Path |
|
|
|
|
from types import SimpleNamespace |
|
|
|
|
from typing import Union |
|
|
|
|
from typing import List, Union |
|
|
|
|
|
|
|
|
|
import cv2 |
|
|
|
|
import matplotlib.pyplot as plt |
|
|
|
@ -26,6 +26,9 @@ from tqdm import tqdm as tqdm_original |
|
|
|
|
|
|
|
|
|
from ultralytics import __version__ |
|
|
|
|
|
|
|
|
|
# Apply monkey patches |
|
|
|
|
from ultralytics.utils.patches import imread, imshow, imwrite, torch_load, torch_save |
|
|
|
|
|
|
|
|
|
# PyTorch Multi-GPU DDP Constants |
|
|
|
|
RANK = int(os.getenv("RANK", -1)) |
|
|
|
|
LOCAL_RANK = int(os.getenv("LOCAL_RANK", -1)) # https://pytorch.org/docs/stable/elastic/run.html |
|
|
|
@ -1073,10 +1076,40 @@ TESTS_RUNNING = is_pytest_running() or is_github_action_running() |
|
|
|
|
set_sentry() |
|
|
|
|
|
|
|
|
|
# Apply monkey patches |
|
|
|
|
from ultralytics.utils.patches import imread, imshow, imwrite, torch_load, torch_save |
|
|
|
|
|
|
|
|
|
torch.load = torch_load |
|
|
|
|
torch.save = torch_save |
|
|
|
|
if WINDOWS: |
|
|
|
|
# Apply cv2 patches for non-ASCII and non-UTF characters in image paths |
|
|
|
|
cv2.imread, cv2.imwrite, cv2.imshow = imread, imwrite, imshow |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def crop_and_pad(frame: np.ndarray, box: List[int], margin_percent: int = 10) -> np.ndarray: |
|
|
|
|
""" |
|
|
|
|
Crop box with margin and take square crop from frame. |
|
|
|
|
|
|
|
|
|
Args: |
|
|
|
|
frame (ndarray): The input frame. |
|
|
|
|
box (list): The bounding box coordinates. |
|
|
|
|
margin_percent (int, optional): The percentage [0-100] of margin to add around the detected object. Defaults to 10. |
|
|
|
|
|
|
|
|
|
Returns: |
|
|
|
|
ndarray: The cropped and resized frame. |
|
|
|
|
""" |
|
|
|
|
x1, y1, x2, y2 = map(int, box) |
|
|
|
|
w, h = x2 - x1, y2 - y1 |
|
|
|
|
|
|
|
|
|
# Add margin |
|
|
|
|
margin_x, margin_y = int(w * margin_percent / 100), int(h * margin_percent / 100) |
|
|
|
|
x1, y1 = max(0, x1 - margin_x), max(0, y1 - margin_y) |
|
|
|
|
x2, y2 = min(frame.shape[1], x2 + margin_x), min(frame.shape[0], y2 + margin_y) |
|
|
|
|
|
|
|
|
|
# Take square crop from frame |
|
|
|
|
size = max(y2 - y1, x2 - x1) |
|
|
|
|
center_y, center_x = (y1 + y2) // 2, (x1 + x2) // 2 |
|
|
|
|
half_size = size // 2 |
|
|
|
|
square_crop = frame[ |
|
|
|
|
max(0, center_y - half_size) : min(frame.shape[0], center_y + half_size), |
|
|
|
|
max(0, center_x - half_size) : min(frame.shape[1], center_x + half_size), |
|
|
|
|
] |
|
|
|
|
|
|
|
|
|
return cv2.resize(square_crop, (224, 224), interpolation=cv2.INTER_LINEAR) |
|
|
|
|