Helper function which creates masks and croops Args: source_image_dir: image directory containing input images source_annotation_dir: annotation directory containing csv annotations dest_images_dir: destination directory for storing images dest_masks_sir: des
(source_image_dir,
source_annotation_dir,
dest_images_dir,
dest_masks_sir,
crops_dest_dir,
model)
| 53 | } |
| 54 | |
| 55 | def prepare_data(source_image_dir, |
| 56 | source_annotation_dir, |
| 57 | dest_images_dir, |
| 58 | dest_masks_sir, |
| 59 | crops_dest_dir, |
| 60 | model): |
| 61 | """ |
| 62 | Helper function which creates masks and croops |
| 63 | Args: |
| 64 | source_image_dir: image directory containing input images |
| 65 | source_annotation_dir: annotation directory containing csv annotations |
| 66 | dest_images_dir: destination directory for storing images |
| 67 | dest_masks_sir: destination directory for storing masks |
| 68 | dest_crops_dir: destinations directory for storing crops |
| 69 | model: super resolution model |
| 70 | """ |
| 71 | annotations_list = glob(os.path.join(source_annotation_dir, "*.csv")) |
| 72 | count = 0 |
| 73 | cnt = 0 |
| 74 | transform = A.Compose([ |
| 75 | A.augmentations.transforms.CLAHE(clip_limit=4.0, |
| 76 | tile_grid_size=(8, 8), |
| 77 | always_apply=False, |
| 78 | p=1.0) |
| 79 | ]) |
| 80 | with tqdm(total=len(annotations_list)) as pbar: |
| 81 | for annotation in annotations_list: |
| 82 | df = pd.read_csv(annotation) |
| 83 | # checking if atleast 1 designation is present in annotation |
| 84 | if df["Designator"].isna().sum() != df.shape[0]: |
| 85 | image_name = list(df["Image File"].unique()) |
| 86 | if os.path.exists(os.path.join(source_image_dir, image_name[0])): |
| 87 | img = cv2.imread(os.path.join(source_image_dir, image_name[0])) |
| 88 | img1 = cv2.cvtColor(img, cv2.COLOR_BGR2HLS) |
| 89 | mask = np.zeros(shape=img.shape, dtype=np.uint8) |
| 90 | transformed = transform(image=img1, mask=mask) |
| 91 | img1 = transformed['image'] |
| 92 | vertices_list = list(df["Vertices"]) |
| 93 | designator_list = list(df["Designator"]) |
| 94 | for (anote, cat) in zip(vertices_list, designator_list): |
| 95 | if cat in color_values: |
| 96 | color_code = color_values[cat] |
| 97 | else: |
| 98 | continue |
| 99 | try: |
| 100 | pts = np.array(ast.literal_eval(anote))[0].reshape((-1, 1, 2)) |
| 101 | except: |
| 102 | continue |
| 103 | mask = cv2.polylines(mask, [pts], True, color_code, 2) |
| 104 | mask = cv2.fillPoly(mask, [pts], color=color_code) |
| 105 | # create crops |
| 106 | mask_copy = np.zeros(shape=img.shape, dtype=np.uint8) |
| 107 | mask_copy = cv2.polylines(mask_copy, [pts], True, color_code, 2) |
| 108 | mask_copy = cv2.fillPoly(mask_copy, [pts], color=color_code) |
| 109 | mask_copy = cv2.cvtColor(mask_copy, cv2.COLOR_BGR2GRAY) |
| 110 | contours, _ = cv2.findContours(mask_copy, cv2.RETR_EXTERNAL, |
| 111 | cv2.CHAIN_APPROX_NONE) |
| 112 | x,y,w,h = cv2.boundingRect(contours[0]) |