def __getitem__(self, idx: int) -> Optional[Dict[str, Any]]:
image_path = self.file_paths[idx]
image = cv2.imread(str(image_path))
if image is None:
return None
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image_height, image_width = image.shape[:2]
image = self.resize(image=image)["image"]
paded = pad_to_size(target_size=(self.max_size, self.max_size),
image=image)
image = paded["image"]
pads = paded["pads"]
image = self.transform(image=image)["image"]
return {
"torched_image": tensor_from_rgb_image(image),
"image_path": str(image_path),
"pads": np.array(pads),
"image_height": image_height,
"image_width": image_width,
}
def test_pad_to_size(image, bboxes, keypoints, target_height, target_width):
target_size = (target_height, target_width)
padded_dict = pad_to_size(target_size, image, bboxes, keypoints)
unpadded_dict = unpad_from_size(**padded_dict)
assert np.array_equal(image, unpadded_dict["image"])
assert np.array_equal(
bboxes, unpadded_dict["bboxes"]), f"{bboxes} {unpadded_dict['bboxes']}"
assert np.array_equal(keypoints, unpadded_dict["keypoints"])
def __getitem__(self, idx: int) -> Optional[Dict[str, Any]]:
image_path = self.file_paths[idx]
image = load_rgb(image_path)
height, width = image.shape[:2]
image = self.transform(image=image)["image"]
pad_dict = pad_to_size((max(image.shape[:2]), max(image.shape[:2])),
image)
return {
"torched_image": tensor_from_rgb_image(pad_dict["image"]),
"image_path": str(image_path),
"pads": pad_dict["pads"],
"original_width": width,
"original_height": height,
}
def __getitem__(self, idx: int) -> Optional[Dict[str, Any]]:
image_path = self.file_paths[idx]
image = np.array(Image.open(image_path))
image_height, image_width = image.shape[:2]
image = self.resize(image=image)["image"]
paded = pad_to_size(target_size=(self.max_size, self.max_size), image=image)
image = paded["image"]
pads = paded["pads"]
image = self.transform(image=image)["image"]
return {
"torched_image": tensor_from_rgb_image(image),
"image_path": str(image_path),
"pads": np.array(pads),
"image_height": image_height,
"image_width": image_width,
}
def predict_jsons(
self, image: np.array, confidence_threshold: float = 0.7, nms_threshold: float = 0.4
) -> List[Dict[str, Union[List, float]]]:
with torch.no_grad():
original_height, original_width = image.shape[:2]
scale_landmarks = torch.from_numpy(np.tile([self.max_size, self.max_size], 5)).to(self.device)
scale_bboxes = torch.from_numpy(np.tile([self.max_size, self.max_size], 2)).to(self.device)
transformed_image = self.transform(image=image)["image"]
paded = pad_to_size(target_size=(self.max_size, self.max_size), image=transformed_image)
pads = paded["pads"]
torched_image = tensor_from_rgb_image(paded["image"]).to(self.device)
loc, conf, land = self.model(torched_image.unsqueeze(0))
conf = F.softmax(conf, dim=-1)
annotations: List[Dict[str, Union[List, float]]] = []
boxes = decode(loc.data[0], self.prior_box, self.variance)
boxes *= scale_bboxes
scores = conf[0][:, 1]
landmarks = decode_landm(land.data[0], self.prior_box, self.variance)
landmarks *= scale_landmarks
# ignore low scores
valid_index = torch.where(scores > confidence_threshold)[0]
boxes = boxes[valid_index]
landmarks = landmarks[valid_index]
scores = scores[valid_index]
# Sort from high to low
order = scores.argsort(descending=True)
boxes = boxes[order]
landmarks = landmarks[order]
scores = scores[order]
# do NMS
keep = nms(boxes, scores, nms_threshold)
boxes = boxes[keep, :].int()
if boxes.shape[0] == 0:
return [{"bbox": [], "score": -1, "landmarks": []}]
landmarks = landmarks[keep]
scores = scores[keep].cpu().numpy().astype(np.float64)
boxes = boxes.cpu().numpy()
landmarks = landmarks.cpu().numpy()
landmarks = landmarks.reshape([-1, 2])
unpadded = unpad_from_size(pads, bboxes=boxes, keypoints=landmarks)
resize_coeff = max(original_height, original_width) / self.max_size
boxes = (unpadded["bboxes"] * resize_coeff).astype(int)
landmarks = (unpadded["keypoints"].reshape(-1, 10) * resize_coeff).astype(int)
for box_id, bbox in enumerate(boxes):
x_min, y_min, x_max, y_max = bbox
x_min = np.clip(x_min, 0, original_width - 1)
x_max = np.clip(x_max, x_min + 1, original_width - 1)
if x_min >= x_max:
continue
y_min = np.clip(y_min, 0, original_height - 1)
y_max = np.clip(y_max, y_min + 1, original_height - 1)
if y_min >= y_max:
continue
annotations += [
{
"bbox": bbox.tolist(),
"score": scores[box_id],
"landmarks": landmarks[box_id].reshape(-1, 2).tolist(),
}
]
return annotations