You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
class ExampleDataset:
def __init__(self, root, policy_container=None, is_train=True):
self.root = root
self.data_type = "train" if is_train else "val"
self.policy_container = policy_container
# main
# self.imgs = list(sorted(glob.glob(f'{root}/images/{self.data_type}/*.jpg')))
# self.boxes = list(sorted(glob.glob(f'{root}/labels/{self.data_type}/*.txt')))
# test
self.imgs = list(sorted(glob.glob(f'{root}/test_data/*.jpg')))
self.boxes = list(sorted(glob.glob(f'{root}/test_data/*txt')))
self.out_dir = f'{root}/test_data_out/'
if not os.path.exists(self.out_dir):
os.mkdir(self.out_dir)
def __len__(self):
return len(self.imgs)
# def __getitem__(self, idx):
# img = np.array(Image.open(self.imgs[idx]))
# boxes_path = self.boxes[idx]
# height, width, _ = img.shape
# # For convenience I’ve hard coded the label and co-ordinates as label, x_min, y_min, x_max, y_max
# # for each bounding box in the image. For your own model you will need to load
# # in the coordinates and do the appropriate transformations.
# boxes = []
# labels = []
# with open(boxes_path, 'r') as in_box:
# for line in in_box:
# if line:
# line = line.split()
# xywh = list(map(int, map(float, line[1:])))
# xyxy = self.convert_xyxy(xywh, width=width, height=height)
# boxes.append(xyxy)
# labels.append(int(line[0]))
# if self.policy_container:
# # Select a random sub-policy from the policy list
# random_policy = self.policy_container.select_random_policy()
# print(random_policy)
# # Apply this augmentation to the image, returns the augmented image and bounding boxes
# # The boxes must be at a pixel level. e.g. x_min, y_min, x_max, y_max with pixel values
# img_aug, bbs_aug = self.policy_container.apply_augmentation(
# random_policy,
# img,
# boxes,
# labels,
# )
# labels = np.array(labels)
# boxes = np.hstack((np.vstack(labels), np.array(boxes))) # Add the labels to the boxes
# bbs_aug= np.array(bbs_aug)
# # Only return the augmented image and bounded boxes if there are
# # boxes present after the image augmentation
# if bbs_aug.size > 0:
# return img, boxes, img_aug, bbs_aug
# else:
# return img, boxes, [], np.array([])
# return img, boxes
def run(self, num_random=10):
for idx in range(len(self.imgs)):
img = np.array(Image.open(self.imgs[idx]))
boxes_path = self.boxes[idx]
height, width, _ = img.shape
# For convenience I’ve hard coded the label and co-ordinates as label, x_min, y_min, x_max, y_max
# for each bounding box in the image. For your own model you will need to load
# in the coordinates and do the appropriate transformations.
boxes = []
labels = []
with open(boxes_path, 'r') as in_box:
for line in in_box:
if line:
line = line.split()
xywh = list(map(int, map(float, line[1:])))
xyxy = self.convert_xyxy(xywh, width=width, height=height)
boxes.append(xyxy)
labels.append(int(line[0]))
if self.policy_container:
# run $num_random times
print("Processing: " + self.imgs[idx])
i = 0
for i in range(num_random):
# Select a random sub-policy from the policy list
random_policy = self.policy_container.select_random_policy()
# Apply this augmentation to the image, returns the augmented image and bounding boxes
# The boxes must be at a pixel level. e.g. x_min, y_min, x_max, y_max with pixel values
img_aug, bbs_aug = self.policy_container.apply_augmentation(
random_policy,
img,
boxes,
labels
)
labels = np.array(labels)
boxes = np.hstack((np.vstack(labels), np.array(boxes))) # Add the labels to the boxes
bbs_aug= np.array(bbs_aug)
# Only return the augmented image and bounded boxes if there are
# boxes present after the image augmentation
if bbs_aug.size > 0:
print("Step: " + str(i))
print(random_policy)
# img, boxes, img_aug, bbs_aug
# to write
cv2.imwrite(str(i)+"_bbaug_"+self.imgs[idx], img_aug)
with open(str(i)+"_bbaug_"+self.boxes[idx], "w") as fw:
fw.writelines(bbs_aug)
i += 1
def convert_xyxy(self, xywh, width, height):
x, w = xywh[0] * width, xywh[2] * width
y, h = xywh[1] * height, xywh[3] * height
x1 = x - w / 2
x2 = x + w / 2
y1 = y - h / 2
y2 = y + h / 2
return list(map(int, [x1, x2, y1, y2]))
The text was updated successfully, but these errors were encountered:
Code reproduced:
The text was updated successfully, but these errors were encountered: