diff --git a/lmi_utils/label_utils/augment_hf.py b/lmi_utils/label_utils/augment_hf.py new file mode 100644 index 00000000..c955117c --- /dev/null +++ b/lmi_utils/label_utils/augment_hf.py @@ -0,0 +1,98 @@ + +import numpy as np +from label_utils.csv_utils import load_csv +from label_utils.shapes import Mask, Rect, Keypoint, Brush +from label_utils.csv_utils import write_to_csv +import os +import cv2 +import collections +import logging + +logging.basicConfig() +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) + + +def main(): + import argparse + ap = argparse.ArgumentParser() + + ap.add_argument('--path_imgs', '-i', required=True, help='the path of a image folder') + ap.add_argument('--path_csv', default='labels.csv', help='[optinal] the path of a csv file that corresponds to path_imgs, default="labels.csv" in path_imgs') + ap.add_argument('--path_out', '-o', required=True, help='the output path') + ap.add_argument('--symc',required=True, help='the comma separated symetrical classes. Every two classes are symetrical to each other.') + args = vars(ap.parse_args()) + path_imgs = args['path_imgs'] + symc = args['symc'].split(',') + if len(symc) % 2 != 0: + raise ValueError('symc should have even number of elements') + symc_dict = {} + for i in range(0, len(symc), 2): + symc_dict[symc[i]] = symc[i+1] + symc_dict[symc[i+1]] = symc[i] + path_csv = args['path_csv'] if args['path_csv']!='labels.csv' else os.path.join(path_imgs, args['path_csv']) + fname_to_shapes,class_to_id = load_csv(path_csv, path_imgs, zero_index=True) + if not os.path.exists(args['path_out']): + os.makedirs(args['path_out']) + + + annots = fname_to_shapes.copy() + for fname in fname_to_shapes: + logger.info(f'processing {fname}') + ext = os.path.basename(fname).split('.')[-1] + updated_fname = os.path.basename(fname).replace(f'.{ext}', f'_augmented_hf.{ext}') + img = cv2.imread(os.path.join(args['path_imgs'], fname)) + height, width = img.shape[:2] + flipped_image = img[:, ::-1].copy() + logger.info(f'flipped image shape: {flipped_image.shape}') + + for shape in fname_to_shapes[fname]: + new_symc = shape.category + if shape.category in symc_dict: + new_symc = symc_dict[shape.category] + + if isinstance(shape, Rect): + x1,y1 = shape.up_left + x2,y2 = shape.bottom_right + bbox = [x1,y1,x2,y2] + + flipped_bbox = [width - x2, y1, width - x1, y2] + annots[updated_fname].append(Rect( + up_left=[flipped_bbox[0], flipped_bbox[1]], + bottom_right=[flipped_bbox[2], flipped_bbox[3]], + angle=shape.angle, + confidence=shape.confidence, + category=shape.category, + im_name=updated_fname, + fullpath=os.path.join(args['path_out'], updated_fname) + )) + elif isinstance(shape, Keypoint): + + # flip the keypoint + x = shape.x + y = shape.y + + flipped_x = width - x + flipped_y = y + annots[updated_fname].append(Keypoint( + x=flipped_x, + y=flipped_y, + confidence=shape.confidence, + category=new_symc, + im_name=updated_fname, + fullpath=os.path.join(args['path_out'], updated_fname) + )) + cv2.imwrite(os.path.join(args['path_out'], fname), img) + cv2.imwrite(os.path.join(args['path_out'], updated_fname), flipped_image) + + # write the updated shapes to a csv file + + write_to_csv(annots, os.path.join(args['path_out'], 'labels.csv')) + logger.info('done augmenting hf') + + + +if __name__ == '__main__': + main() + + diff --git a/lmi_utils/label_utils/crop_by_foreground.py b/lmi_utils/label_utils/crop_by_foreground.py new file mode 100644 index 00000000..1070fd57 --- /dev/null +++ b/lmi_utils/label_utils/crop_by_foreground.py @@ -0,0 +1,243 @@ + +import numpy as np +from label_utils.csv_utils import load_csv +from label_utils.shapes import Mask, Rect, Keypoint, Brush +from label_utils.csv_utils import write_to_csv +import os +import cv2 +import collections +import logging + +logging.basicConfig() +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) + +def crop_by_percent(image, crop_percent, crop_from = 'top'): + height, width = image.shape[:2] + x1, y1 = 0, 0 + x2, y2 = width, height + if crop_from == 'top': + x1, y1 = 0, 0 + x2, y2 = width, int(height * crop_percent) + elif crop_from == 'bottom': + x1, y1 = 0, int(height * (1 - crop_percent)) + x2, y2 = width, height + return x1, y1, x2, y2 + +def crop_kp(bbox, shape): + x1,y1,x2,y2 = bbox + w,h = x2-x1, y2-y1 + x,y = shape.x, shape.y + x -= x1 + y -= y1 + + valid = True + if x<0 or x>=w or y<0 or y>=h: + valid = False + logger.warning(f'in {shape.im_name}, keypoint {x:.4f},{y:.4f} is out of the foreground bbox: {x1:.4f},{y1:.4f},{x2:.4f},{y2:.4f}. skip') + + return x,y,valid + + +def crop_bbox(bbox1, bbox2): + crop_x1, crop_y1, crop_x2, crop_y2 = bbox1 + target_x1, target_y1, target_x2, target_y2 = bbox2 + adjusted_x1 = target_x1 - crop_x1 + adjusted_y1 = target_y1 - crop_y1 + adjusted_x2 = target_x2 - crop_x1 + adjusted_y2 = target_y2 - crop_y1 + return adjusted_x1, adjusted_y1, adjusted_x2, adjusted_y2 + +def crop_mask(bbox, mask=None, polygon_mask=None, bbox_format="xywh"): + # Interpret the bounding box coordinates + if bbox_format == "xywh": + x, y, w, h = bbox + xmin, ymin, xmax, ymax = x, y, x + w, y + h + elif bbox_format == "xyxy": + xmin, ymin, xmax, ymax = bbox + w, h = xmax - xmin, ymax - ymin + else: + raise ValueError("bbox_format must be either 'xywh' or 'xyxy'") + cropped_mask = None + if mask is not None: + cropped_mask = mask[ymin:ymax, xmin:xmax] + + # If a polygon mask is provided, adjust its coordinates relative to the crop. + cropped_polygon = None + if polygon_mask is not None: + # Ensure the input is a numpy array of shape (N, 2) + polygon_mask = np.asarray(polygon_mask) + if polygon_mask.ndim != 2 or polygon_mask.shape[1] != 2: + raise ValueError("polygon_mask must be a 2D array with shape (N_points, 2)") + + # Shift the polygon by the top-left corner of the bounding box. + cropped_polygon = polygon_mask - np.array([xmin, ymin]) + + cropped_polygon[:, 0] = np.clip(cropped_polygon[:, 0], 0, w) + cropped_polygon[:, 1] = np.clip(cropped_polygon[:, 1], 0, h) + + return cropped_mask, cropped_polygon + + +def main(): + import argparse + ap = argparse.ArgumentParser() + + ap.add_argument('--path_imgs', '-i', required=True, help='the path of a image folder') + ap.add_argument('--path_csv', default='labels.csv', help='[optinal] the path of a csv file that corresponds to path_imgs, default="labels.csv" in path_imgs') + ap.add_argument('--path_out', '-o', required=True, help='the output path') + ap.add_argument('--target_classes',required=True, help='the comma separated target classes to crop') + ap.add_argument('--crop_by_percent', type=float, default=0.0, help='the percentage of the image to crop', required=False) + ap.add_argument('--crop_from', default='top', help='the direction to crop the image', required=False) + args = vars(ap.parse_args()) + path_imgs = args['path_imgs'] + path_csv = args['path_csv'] if args['path_csv']!='labels.csv' else os.path.join(path_imgs, args['path_csv']) + target_classes = args['target_classes'].split(',') + fname_to_shapes,class_to_id = load_csv(path_csv, path_imgs, zero_index=True) + if not os.path.exists(args['path_out']): + os.makedirs(args['path_out']) + + foreground_shapes = {} + for fname in fname_to_shapes: + if fname not in foreground_shapes: + foreground_shapes[fname] = { + 'foreground': [] + } + for shape in fname_to_shapes[fname]: + #get class ID + if shape.category not in target_classes: + continue + + + if isinstance(shape, Rect): + x0,y0 = shape.up_left + x2,y2 = shape.bottom_right + foreground_shapes[fname]['foreground'] = list(map(int, [x0,y0,x2,y2])) + if len(foreground_shapes[fname]['foreground'])==0: + logger.warning(f'no foreground found in {fname}') + + + annots = collections.defaultdict(list) + for fname in fname_to_shapes: + + ext = os.path.basename(fname).split('.')[-1] + updated_fname = os.path.basename(fname).replace(f'.{ext}', f'_cropped.{ext}') + if fname not in foreground_shapes: + continue + if len(foreground_shapes[fname]['foreground'])==0: + continue + logger.info(f'processing {fname}') + + image = cv2.imread(os.path.join(path_imgs, fname)) + if image is None: + logger.warning(f'failed to read {fname}, skip') + continue + + if args['crop_by_percent']>0: + x1,y1,x2,y2 = foreground_shapes[fname]['foreground'] + sx,sy,ex,ey = crop_by_percent(image[y1:y2, x1:x2], args['crop_by_percent'], args['crop_from']) + logger.info(f'cropping {fname} by {args["crop_by_percent"]*100:.2f}% from {args["crop_from"]}') + + + + H,W = image.shape[:2] + for shape in fname_to_shapes[fname]: + + if shape.category in target_classes: + continue + + + if isinstance(shape, Rect): + x1,y1 = shape.up_left + x2,y2 = shape.bottom_right + bbox = [x1,y1,x2,y2] + cropped_bbox = crop_bbox(foreground_shapes[fname]['foreground'], bbox) + if args['crop_by_percent']>0: + cropped_bbox = crop_bbox([sx,sy,ex,ey], cropped_bbox) + annots[updated_fname].append( + Rect( + up_left=[cropped_bbox[0], cropped_bbox[1]], + bottom_right=[cropped_bbox[2], cropped_bbox[3]], + angle=shape.angle, + confidence=shape.confidence, + category=shape.category, + im_name=updated_fname, + fullpath=os.path.join(args['path_out'], updated_fname) + ) + ) + + if isinstance(shape, Brush): + mask = shape.to_mask((H,W)) + mask = mask.astype(np.uint8)*255 + # mask = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR) + cropped_mask, _ = crop_mask(foreground_shapes[fname]['foreground'], mask, bbox_format="xyxy") + if args['crop_by_percent']>0: + cropped_mask, _ = crop_mask([sx,sy,ex,ey], cropped_mask, bbox_format="xyxy") + # add to brush labels + annots[updated_fname].append( + Brush( + mask=cropped_mask>128, + confidence=shape.confidence, + category=shape.category, + im_name=updated_fname, + fullpath=os.path.join(args['path_out'], updated_fname) + ) + ) + + if isinstance(shape, Keypoint): + x,y,is_valid = crop_kp(foreground_shapes[fname]['foreground'], shape) + if not is_valid: + continue + if args['crop_by_percent']>0: + x,y,is_valid = crop_kp([sx,sy,ex,ey], Keypoint(x=x, y=y, confidence=shape.confidence, category=shape.category, im_name=updated_fname, fullpath=os.path.join(args['path_out'], updated_fname))) + if not is_valid: + continue + annots[updated_fname].append( + Keypoint( + x=x, + y=y, + confidence=shape.confidence, + category=shape.category, + im_name=updated_fname, + fullpath=os.path.join(args['path_out'], updated_fname) + ) + ) + + # suporting polygon mask + if isinstance(shape, Mask): + polygon_x = shape.X + polygon_y = shape.Y + polygon_mask = np.stack([polygon_x, polygon_y], axis=1) + _, cropped_polygon = crop_mask(foreground_shapes[fname]['foreground'], mask=None, polygon_mask=polygon_mask, bbox_format="xyxy") + if args['crop_by_percent']>0: + _, cropped_polygon = crop_mask([sx,sy,ex,ey], mask=None, polygon_mask=cropped_polygon, bbox_format="xyxy") + # add to mask labels + annots[updated_fname].append( + Mask( + x_vals=cropped_polygon[:, 0].tolist(), + y_vals=cropped_polygon[:, 1].tolist(), + confidence=shape.confidence, + category=shape.category, + im_name=updated_fname, + fullpath=os.path.join(args['path_out'], updated_fname) + ) + ) + + # save the cropped image + x1,y1,x2,y2 = foreground_shapes[fname]['foreground'] + cropped_image = image[y1:y2, x1:x2] + if args['crop_by_percent']>0: + cropped_image = cropped_image[sy:ey, sx:ex] + + cv2.imwrite(os.path.join(args['path_out'], updated_fname), cropped_image) + + + # save the updated shapes + write_to_csv(annots, os.path.join(args['path_out'], 'labels.csv')) + + + +if __name__ == '__main__': + main() + + diff --git a/lmi_utils/label_utils/csv_to_yolo.py b/lmi_utils/label_utils/csv_to_yolo.py index 8262695f..a9005515 100644 --- a/lmi_utils/label_utils/csv_to_yolo.py +++ b/lmi_utils/label_utils/csv_to_yolo.py @@ -178,20 +178,22 @@ def convert_to_txt(fname_to_shapes, target_classes, is_seg=False, is_convert=Fal rows.append([class_name]+xywh) elif isinstance(shape, Keypoint): x,y = shape.x, shape.y - row = [x/W, y/H] + row = [shape.category, x/W, y/H] ignore_cls.add(shape.category) kps.append(row) # assign keypoint to bbox new_rows = {} + kps.sort(key=lambda x: x[0]) # sort by class name for kp in kps: - x,y = kp + c,x,y = kp + x,y = x*W, y*H hit = 0 for row in rows: if len(row)!=5: logging.warning(f'key point can only be assign to a bbox, but got {len(row)} values in a row. Skip it.') continue - xc,yc,w,h = row[1:] + xc,yc,w,h = np.array(row[1:]*np.array([W,H,W,H])) x1,y1 = xc-w/2, yc-h/2 x2,y2 = xc+w/2, yc+h/2 if x1<=x<=x2 and y1<=y<=y2: @@ -199,7 +201,7 @@ def convert_to_txt(fname_to_shapes, target_classes, is_seg=False, is_convert=Fal key = ','.join(str(v) for v in row) if key not in new_rows: new_rows[key] = row[:] - new_rows[key].extend(kp) + new_rows[key].extend(kp[1:]) if not hit: raise Exception(f'key point ({x},{y}) is not in any bbox. Fix it.') @@ -229,7 +231,7 @@ def assign_class_id(fname_to_rows, class_to_id): row[0] = class_to_id[row[0]] -def write_txts(fname_to_rows, path_txts): +def write_txts(fname_to_rows, path_txts, names=None): """ write to the yolo format txts Arugments: @@ -237,7 +239,10 @@ def write_txts(fname_to_rows, path_txts): path_txts: the output folder contains txt files """ os.makedirs(path_txts, exist_ok=True) + num_files = 0 for fname in fname_to_rows: + if names is not None and fname not in names: + continue txt_file = os.path.join(path_txts, fname) with open(txt_file, 'w') as f: for shape in fname_to_rows[fname]: @@ -248,7 +253,8 @@ def write_txts(fname_to_rows, path_txts): row2 += f'{pt:.4f} ' row2 += '\n' f.write(row2) - logger.info(f' wrote {len(fname_to_rows)} txt files to {path_txts}') + num_files += 1 + logger.info(f' wrote {num_files} txt files to {path_txts}') def copy_images_in_folder(path_img, path_out, fnames=None): @@ -347,13 +353,21 @@ def copy_images_in_folder(path_img, path_out, fnames=None): #write labels/annotations path_txts = os.path.join(args['path_out'], 'labels') - write_txts(fname_to_rows, path_txts) #write images path_img_out = os.path.join(args['path_out'], 'images') if args['bg']: logger.info('save background images') copy_images_in_folder(path_imgs, path_img_out) + write_txts(fname_to_rows, path_txts) else: logger.info('skip background images') - copy_images_in_folder(path_imgs, path_img_out, fname_to_shapes.keys()) + names = [] + txt_files = [] + for name in fname_to_shapes: + if len(fname_to_rows[name.replace('.png', '.txt').replace('.jpg', '.txt')]) != 0: + names.append(name) + txt_files.append(name.replace('.png', '.txt').replace('.jpg', '.txt')) + write_txts(fname_to_rows, path_txts, txt_files) + copy_images_in_folder(path_imgs, path_img_out, fnames=names) + \ No newline at end of file diff --git a/lmi_utils/label_utils/lst_to_csv.py b/lmi_utils/label_utils/lst_to_csv.py index 23aff3de..14f19131 100644 --- a/lmi_utils/label_utils/lst_to_csv.py +++ b/lmi_utils/label_utils/lst_to_csv.py @@ -5,6 +5,7 @@ import numpy as np import collections import glob +import pathlib from label_studio_sdk.converter.brush import decode_rle from label_utils.csv_utils import write_to_csv @@ -19,6 +20,38 @@ PRED_NAME = 'preds.csv' +class Node: + def __init__(self, name): + self.name = name + self.children = [] + + def append(self, node): + self.children.append(node) + + +def common_node(paths): + """find the deepest node that is the common parent of all paths + """ + root = Node('') + for path in paths: + path = pathlib.Path(path).as_posix() + cur = root + for part in path.split('/'): + found = False + for child in cur.children: + if child.name == part: + cur = child + found = True + break + if not found: + new_node = Node(part) + cur.append(new_node) + cur = new_node + cur = root + while len(cur.children) == 1: + cur = cur.children[0] + return cur + def lst_to_shape(result:dict, fname:str, load_confidence=False): """parse the result from label studio result dict, return a Shape object @@ -76,24 +109,47 @@ def get_annotations_from_json(path_json): json_files=[path_json] else: json_files=glob.glob(os.path.join(path_json,'*.json')) - - annots = collections.defaultdict(list) - preds = collections.defaultdict(list) + + # get common parent node + for path_json in json_files: + with open(path_json) as f: + l = json.load(f) + + img_lists = [] + for dt in l: + if 'data' not in dt: + raise Exception('missing "data" in json file. Ensure that the label studio export format is not JSON-MIN.') + f = dt['data']['image'] + img_lists.append(f) + root = common_node(img_lists) + prefixes = set(n.name for n in root.children) + + annots = {} + preds = {} for path_json in json_files: logger.info(f'Extracting labels from: {path_json}') with open(path_json) as f: l = json.load(f) - + cnt_anno = 0 cnt_image = 0 cnt_pred = 0 cnt_wrong = 0 for dt in l: - # load file name - if 'data' not in dt: - raise Exception('missing "data" in json file. Ensure that the label studio export format is not JSON-MIN.') - f = dt['data']['image'] # image web path - fname = os.path.basename(f) + f = pathlib.Path(dt['data']['image']) + l2 = f.as_posix().split('/') + i = next((l2.index(p) for p in prefixes if p in l2), -1) + if i==-1: + raise Exception('Cannot find common parent node') + key = '_'.join(l2[i:-1]) + + if key not in annots: + annots[key] = collections.defaultdict(list) + preds[key] = collections.defaultdict(list) + + fname = f.name + if fname in annots[key]: + raise Exception('Found duplicate name') if 'annotations' in dt: cnt = 0 @@ -104,14 +160,14 @@ def get_annotations_from_json(path_json): for result in annot['result']: shape = lst_to_shape(result,fname) if shape is not None: - annots[fname].append(shape) + annots[key][fname].append(shape) cnt_anno += 1 if 'prediction' in annot and 'result' in annot['prediction']: for result in annot['prediction']['result']: shape = lst_to_shape(result,fname,load_confidence=True) if shape is not None: - preds[fname].append(shape) + preds[key][fname].append(shape) cnt_pred += 1 if cnt>0: cnt_image += 1 @@ -147,7 +203,9 @@ def get_annotations_from_json(path_json): if os.path.isfile(args.path_out): raise Exception('The output path should be a directory') - if not os.path.isdir(args.path_out): - os.makedirs(args.path_out) - write_to_csv(annots, os.path.join(args.path_out, LABEL_NAME)) - write_to_csv(preds, os.path.join(args.path_out, PRED_NAME)) + for k in annots: + path_out = os.path.join(args.path_out, k) + os.makedirs(path_out,exist_ok=1) + + write_to_csv(annots[k], os.path.join(path_out, LABEL_NAME)) + write_to_csv(preds[k], os.path.join(path_out, PRED_NAME)) \ No newline at end of file