Update load image in sample (#10223)

This commit is contained in:
Indira Salyahova
2022-02-15 10:18:27 +03:00
committed by GitHub
parent 13c024b7a3
commit acf6185bf3
2 changed files with 3 additions and 3 deletions

View File

@@ -71,7 +71,7 @@ class COCOLoader(DataLoader):
def _read_and_preprocess_image(image_path):
image = cv2.imread(image_path, cv2.IMREAD_COLOR)
image = cv2.resize(image, (640, 640))
return image.transpose(2, 0, 1)
return image
@staticmethod
def prepare_bbox(x, y, weight, height):

View File

@@ -71,10 +71,10 @@ class VOCSegmentationLoader(DataLoader):
# Pad image to destination size
image = central_padding(image, self._image_size, self._image_size)
return image.transpose(2, 0, 1) # Change data layout from HWC to CHW
return image
def _read_and_preprocess_mask(self, mask_path):
mask = self._read_and_preprocess_image(mask_path).transpose(1, 2, 0)
mask = self._read_and_preprocess_image(mask_path)
encoded_mask = np.zeros((mask.shape[0], mask.shape[1]), dtype=np.int16)
for label, color in enumerate(_SEGMENTATION_COLORS):
encoded_mask[np.where(np.all(mask == color, axis=-1))[:2]] = label