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