def main(): # path = "/home/jitesh/3d/data/coco_data/fp2_400_2020_06_05_14_46_57_coco-data" # path = "/home/jitesh/3d/data/coco_data/fp2_40_2020_06_05_10_37_48_coco-data" # img_path = "/home/jitesh/3d/data/UE_training_results/fp2_40" # path = "/home/jitesh/3d/data/coco_data/bolt_real4" # path = "/home/jitesh/3d/data/coco_data/hc1_1000_2020_06_30_18_43_56_coco-data" path = "/home/jitesh/3d/data/coco_data/bolt/b2_coco-data" # path = "/home/jitesh/3d/data/coco_data/hr1_300_coco-data" # path = "/home/jitesh/3d/data/coco_data/bolt_real1_training_result1" # img_path = "/home/jitesh/3d/data/UE_training_results/fp2_40" img_path = f'{path}/img' # model = "COCO-Keypoints/keypoint_rcnn_R_50_FPN_1x.yaml" # model = "COCO-Keypoints/keypoint_rcnn_R_101_FPN_3x.yaml" # model = "COCO-Detection/faster_rcnn_R_50_FPN_1x.yaml" model = "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml" model_name = model.split('/')[0].split('-')[1] + '_'\ + model.split('/')[1].split('_')[2] + '_'\ + model.split('/')[1].split('_')[3] + '_'\ + model.split('/')[1].split('_')[5].split('.')[0] make_dir_if_not_exists(f'{path}/weights') _output_dir_path = f'{path}/weights/{model_name}' output_dir_path = f"{_output_dir_path}_1" resume = True # resume=False # if not resume: i = 1 while os.path.exists(f"{_output_dir_path}_{i}"): i = i + 1 if resume: output_dir_path = f"{_output_dir_path}_{i-1}" else: output_dir_path = f"{_output_dir_path}_{i}" # coco_ann_path = os.path.join(path, "json/bolt.json") # coco_ann_path = os.path.join(path, "json/bbox_resized.json") coco_ann_path = os.path.join(path, "json/bolt.json") # train(path, coco_ann_path, img_path, output_dir_path, resume=resume, model=model) d2 = D2Trainer( coco_ann_path=coco_ann_path, img_path=img_path, output_dir_path=output_dir_path, resume=resume, model=model, aug_on=False, # num_workers=2, # images_per_batch=2, # base_lr=0.002, # max_iter=10000, # checkpoint_period=100, # batch_size_per_image=512, # num_classes=1, # max_size_train=1024, # min_size_train=1024, # aug_on=True, detectron2_dir_path="/home/jitesh/prj/detectron2") d2.train()
def run(self): if not path_exists(self.ndds_dir): printj.red(f"Input ndds_dir path does not exist./n{self.ndds_dir}") make_dir_if_not_exists(dir_path=self.output_dir) coco_data_dir = os.path.join( self.output_dir, f"{get_filename_from_path(self.ndds_dir)}_coco_data") delete_dir_if_exists(dir_path=coco_data_dir) make_dir(dir_path=coco_data_dir) ndds_dataset = NDDS_Dataset.load_from_dir(json_dir=self.ndds_dir)
def replace_bg_wrt_seg_ann( coco_data_dir: str, json_filename: str, bg_dirs: list, img_dir_name: str = "img", output_img_dir_name: str = "img_", aug_on: bool = False, aug_json: str = None, show_preview: bool = False, ): coco_dataset = COCO_Dataset.load_from_path( json_path=f"{coco_data_dir}/json/{json_filename}.json", check_paths=False) # image_path_list = folder_list(folder1) image_path_list = [] for bg_dir in bg_dirs: image_path_list += dir_contents_path_list_with_extension( dirpath=bg_dir, extension=['.jpg', '.jpeg', '.png']) bg_gen = image_sequence(image_path_list) pbar = tqdm(coco_dataset.images, colour='#44aa44') for image in pbar: pbar.set_description("Changing background") pbar.set_postfix({'file_name': image.file_name}) image_path_split = image.coco_url.split("/") image_path_split[-2] = img_dir_name image_path = "/".join(image_path_split) for ann in coco_dataset.annotations: if ann.image_id == image.id: seg = ann.segmentation background = next(bg_gen) if aug_on: aug = aug_flip_and_rotate(aug_json) background = aug(image=np.array(background))['image'] orig_image = cv2.imread(image_path) assert orig_image.shape[1] == image.width assert orig_image.shape[0] == image.height mask = np.zeros((image.width, image.height), np.uint8) contours = seg.to_contour() cv2.drawContours(mask, contours, -1, (255, 255, 255), -1) final = replace_bg_wrt_mask(orig_image, background, mask) if show_preview: show_image(final) else: output = os.path.join(coco_data_dir, output_img_dir_name) make_dir_if_not_exists(coco_data_dir) make_dir_if_not_exists(output) output_path = os.path.join(output, image.file_name) cv2.imwrite(output_path, final)
def convert_4ch_to_3ch(dataset_path, split_ratio=[0.8, 0.1, 0.1]): b_type_list = [ dir for dir in os.listdir(dataset_path) if dir not in DATA_TYPES ] img_path_list = dict() for data_type in DATA_TYPES: img_path_list[data_type] = [] import random for b_type in b_type_list: data = glob.glob(f'{dataset_path}/{b_type}/*.png') # train_data.append(data[:int(len(data)*0.8)]) random.shuffle(data) s1 = split_ratio[0] s2 = split_ratio[0] + split_ratio[1] assert 1 == split_ratio[0] + split_ratio[1] + split_ratio[2] img_path_list["train"] += data[:int(len(data) * s1)] img_path_list["val"] += data[int(len(data) * s1):int(len(data) * s2)] img_path_list["test"] += data[int(len(data) * s2):] print(f'len(train_data): {len(img_path_list["train"])}') print(f'len(val_data): {len(img_path_list["val"])}') print(f'len(test_data): {len(img_path_list["test"])}') # In[ ]: import pyjeasy.file_utils as f import cv2 filename_list = dict() for data_type in DATA_TYPES: dirname_new = os.path.join(dataset_path, data_type) f.make_dir_if_not_exists(dirname_new) for file_path in tqdm(img_path_list[data_type]): file_path_split = file_path.split("/") dirname_old = file_path_split[-2].split("_")[0] filename_old = file_path_split[-1] # filename_new = dirname_old.replace("b", "") + "_" + filename_old filename_new = dirname_old + "_" + filename_old output_img_path = os.path.join(dirname_new, filename_new) f.delete_file_if_exists(output_img_path) # Converting 4 channel to 3 channel and then writing in different folder img = cv2.imread(file_path) img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR) cv2.imwrite(output_img_path, img) # f.copy_file(src_path=file_path, dst_path=output_img_path, verbose=False) filename_list[data_type] = os.listdir(dirname_new)
def chop( input_img_path, output_path, c_sizes, c_strides, padding_color=[255, 255, 255], input_json_path=None, theshold: float=1.0 ): lbm = None if input_json_path: lbm = LabelmeLoader(input_json_path) orig = cv2.imread(input_img_path) gray = cv2.cvtColor(orig,cv2.COLOR_RGB2GRAY) img = cv2.cvtColor(gray,cv2.COLOR_GRAY2RGB) img_h, img_w = gray.shape img = cv2.copyMakeBorder( img, top=0, bottom=max(c_sizes), left=0, right=max(c_sizes), borderType=cv2.BORDER_CONSTANT, value=padding_color) make_dir_if_not_exists(os.path.abspath(f'{output_path}/..')) make_dir_if_not_exists(output_path) for c_size, c_stride in zip(c_sizes, c_strides): i = 0 h = 0 while h < img_h: sys.stdout.write(f"\rCreating image {i} if size {c_size}") sys.stdout.flush() w = 0 while w < img_w: c_img = img[h:h+c_size, w:w+c_size, :] # if show_image(c_img, f"size: {c_size}, horizntal: {w}-{w+c_size}, vertical: {h}-{h+c_size}", 900): # return output = os.path.join(output_path, f't_{c_size}_{i}.jpg') useful_img = lbm.create_cropped_labelme( c_point1=[w, h], c_point2=[w+c_size, h+c_size], output_img=f'{output}', theshold=theshold) if useful_img: cv2.imwrite(f'{output}', c_img) i += 1 w += c_size - c_stride h += c_size - c_stride sys.stdout.write(f"\rCreated images of sizes {c_sizes}. {' '*50}\n")
def save_ckp(state, checkpoint_dir, best_model_dir, filename_pth='checkpoint.pth'): """ checkpoint = { 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict() } save_ckp(checkpoint, is_best, checkpoint_dir, model_dir) """ make_dir_if_not_exists(checkpoint_dir) make_dir_if_not_exists(best_model_dir) best_fpath = os.path.join(best_model_dir, 'best_model.pth') rename_f_path = os.path.join(best_model_dir, filename_pth) f_path = os.path.join(checkpoint_dir, filename_pth) if file_exists(best_fpath): os.rename(best_fpath, rename_f_path) move_file(rename_f_path, f_path) torch.save(state, best_fpath)
def run(path): printj.yellow(f'Processing {path}') ndds_dataset = NDDS_Dataset.load_from_dir(json_dir=path, show_pbar=True) for frame in tqdm(ndds_dataset.frames): # printj.red(frame) # Fix Naming Convention for ann_obj in frame.ndds_ann.objects: if ann_obj.class_name.startswith('bolt'): bolt_bbox = ann_obj.bounding_box roi_bbox = bolt_bbox for ann_obj in frame.ndds_ann.objects: if ann_obj.class_name.startswith('i-'): imark_bbox = ann_obj.bounding_box roi_bbox = roi_bbox + imark_bbox elif ann_obj.class_name.startswith('o-'): omark_bbox = ann_obj.bounding_box # printj.purple(roi_bbox) roi_bbox = roi_bbox + omark_bbox # printj.cyan(roi_bbox) # printj.cyan(omark_bbox) # printj.green('ghcfxnn') roi_bbox = roi_bbox.to_int() # printj.red(roi_bbox) # printj.red(f"{roi_bbox.xmin}:{roi_bbox.xmax}, {roi_bbox.ymin}:{roi_bbox.ymax}") input_path = frame.img_path output_path = os.path.abspath(f"{input_path}/../../../bolt_cropped")\ + "/"\ + os.path.abspath(f"{input_path}").split('/')[-2]\ + "/"\ + os.path.abspath(f"{input_path}").split('/')[-1] make_dir_if_not_exists(os.path.abspath(f"{output_path}/../../..")) make_dir_if_not_exists(os.path.abspath(f"{output_path}/../..")) make_dir_if_not_exists(os.path.abspath(f"{output_path}/..")) img = cv2.imread(input_path) cv2.imwrite( output_path, img[roi_bbox.ymin:roi_bbox.ymax, roi_bbox.xmin:roi_bbox.xmax])
def infer(self, input_type: Union[str, int], output_type: Union[str, int], input_path: Union[str, List[str]], output_path: Union[str, List[str]], show_max_score_only: bool = False, show_class_label: bool = True, show_class_label_score_only: bool = False, show_keypoint_label: bool = True, show_legends: bool = False, show_bbox: bool = True, show_keypoints: bool = True, show_segmentation: bool = True, color_bbox: list = None, transparent_mask: bool = True, transparency_alpha: float = 0.3, ignore_keypoint_idx=None, gt_path: Union[str, List[str]] = None, result_json_path: str = None): """ Valid options for, === input_type: --- ["image", "image_list", "image_directory", "image_directories_list", "video", "video_list", "video_directory" ] output_type: --- ["return_summary", "show_image", "write_image", "write_video" ] Returns --- The inference result of all data formats. """ self.counter = infinite_sequence() check_value(input_type, check_from=["image", "image_list", "image_directory", "image_directories_list", "video", "video_list", "video_directory"]) check_value(value=output_type, check_from=[ "show_image", "write_image", "write_video"]) self.gt_path = gt_path if self.gt_path: check_file_exists(gt_path) with open(gt_path) as json_file: self.gt_data = json.load(json_file) if result_json_path is None: if dir_exists(output_path): self.result_json_path = f'{output_path}/result.json' else: _p = output_path.split('.') _output_path = '.'.join(_p[:-1]) self.result_json_path = f'{_output_path}.json' predict_image = partial(self._infer_image, show_max_score_only=show_max_score_only, show_class_label=show_class_label, show_class_label_score_only=show_class_label_score_only, show_keypoint_label=show_keypoint_label, show_bbox=show_bbox, show_keypoints=show_keypoints, show_segmentation=show_segmentation, color_bbox=color_bbox, transparent_mask=transparent_mask, transparency_alpha=transparency_alpha, ignore_keypoint_idx=ignore_keypoint_idx, show_legends=show_legends # gt_path=gt_path, ) if input_type == "image": if file_exists(input_path): output = predict_image(input_path) else: raise Error if output_type == "return_summary": return self.pred_dataset elif output_type == "show_image": show_image(output) elif output_type == "write_image": cv2.imwrite(f'{output_path}', output) elif output_type == "write_video": raise NotImplementedError else: raise Exception elif input_type == "image_list": for image_path in tqdm(input_path, colour='#66cc66'): output = predict_image(image_path) if output_type == "show_image": if show_image(output): break elif output_type == "write_image": make_dir_if_not_exists(output_path) filename = f.get_filename(image_path) cv2.imwrite(f'{output_path}/{filename}', output) elif output_type == "write_video": raise NotImplementedError else: raise Exception elif input_type == "image_directory": image_path_list = f.dir_contents_path_list_with_extension( dirpath=input_path, extension=[".png", ".jpg", ".jpeg"]) for image_path in tqdm(image_path_list, colour='#66cc66'): output = predict_image(image_path) # output = self.draw_gt(gt_path, gt_data, image_path, output) if output_type == "show_image": if show_image(output): break elif output_type == "write_image": make_dir_if_not_exists(output_path) filename = f.get_filename(image_path) cv2.imwrite(f'{output_path}/{filename}', output) elif output_type == "write_video": raise NotImplementedError else: raise Exception elif input_type == "image_directories_list": for image_directory in tqdm(input_path): image_path_list = f.dir_contents_path_list_with_extension( dirpath=input_path, extension=[".png", ".jpg", ".jpeg"]) for image_path in tqdm(image_path_list, colour='#66cc66'): output = predict_image(image_path) if output_type == "show_image": if show_image(output): break elif output_type == "write_image": filename = f.get_filename(image_path) directory_name = f.get_directory_name(image_path) if f.dir_exists(f'{output_path}/{directory_name}'): f.delete_all_files_in_dir( f'{output_path}/{directory_name}') else: f.make_dir(f'{output_path}/{directory_name}') cv2.imwrite( f'{output_path}/{directory_name}/{filename}', output) elif output_type == "write_video": raise NotImplementedError else: raise Exception elif input_type == "video": raise NotImplementedError elif input_type == "video_list": raise NotImplementedError elif input_type == "video_directory": raise NotImplementedError else: raise Exception # printj.cyan(self.pred_dataset) self.write_predictions_json() if self.gt_path: # print(self.df) # pip install openpyxl self.df.to_excel(os.path.abspath( f'{result_json_path}/../test_data.xlsx'))
# # palette = [int(c*255) for color in palette for c in color] # palette = [[int(c*255) for c in color] for color in palette ] # printj.red(palette) # printj.red(np.array(color_palette(None, 3))*255) # x() inferer_seg = D2Inferer( weights_path='/home/jitesh/3d/data/coco_data/ty1w_coco-data/weights/InstanceSegmentation_R_50_1x_aug_val_1/model_0019999.pth', confidence_threshold=0.01, class_names=['tropicana'], model='mask_rcnn_R_50_FPN_1x', size_min=1000, size_max=1000, ) now = datetime.now() dt_string3 = now.strftime("%m_%d_%H") # infer_dump_dir = f'/home/jitesh/3d/blender/Learning/tropicana/infer_5k_{dt_string3}' infer_dump_dir = f'/home/jitesh/3d/blender/Learning/tropicana_yellow/infer_5k_{dt_string3}' delete_dir_if_exists(infer_dump_dir) make_dir_if_not_exists(infer_dump_dir) dir_path = "/home/jitesh/Downloads/Photos (3)" # dir_path = "/home/jitesh/3d/blender/Learning/tropicana_yellow/img" # dir_path = "/home/jitesh/3d/data/coco_data/tropi1_coco-data/img" inferer_seg.infer(input_type="image_directory", output_type="show_image", input_path=dir_path, output_path=infer_dump_dir, show_class_label=True, show_max_score_only=True, transparency_alpha=.9, )
import os from os import path from os.path import join import sys import json import printj import pyjeasy.file_utils as f import cv2 from tqdm import tqdm PATH = "/home/jitesh/sekisui/bolt/hexagon_bolts" OUTPUT_PATH = "/home/jitesh/sekisui/bolt/cropped_hexagon_bolts" f.make_dir_if_not_exists(OUTPUT_PATH) json_list = f.get_all_filenames_of_extension(dirpath=PATH, extension="json") # printj.blue(json_list) for json_file in tqdm(json_list): filename = json_file.split(".")[0] output_image_path = os.path.join(OUTPUT_PATH, f"{filename}.jpg") image_path = os.path.join(PATH, f"{filename}.jpg") json_path = os.path.join(PATH, json_file) if f.path_exists(image_path): img = cv2.imread(image_path) with open(json_path) as json_data: data = json.load(json_data) data = data["shapes"] i = 0 for d in data: if d["label"] == "bolt-roi":
def predict(self, model=None, model_path=None, test_dir_path=None, no_of_samples: int = 24, test_label: int = 1, save_csv_path: str = None, write_images=None): if write_images: f.make_dir_if_not_exists(write_images) f.delete_all_files_in_dir(write_images) f.make_dir_if_not_exists(f'{write_images}/{self.b_type_list[0]}') f.make_dir_if_not_exists(f'{write_images}/{self.b_type_list[1]}') if model is None: model = self.model try: model.load_state_dict(torch.load(model_path)['state_dict']) except KeyError: model.load_state_dict(torch.load(model_path)) # "/home/jitesh/prj/classification/test/bolt/ckpt_densenet121_mark_correct_128_s3_5.pth")) data_transform = transforms.Compose([ transforms.Resize((self.img_size, self.img_size)), # transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), transforms.ToTensor(), ]) test_list = [ file for file in os.listdir(test_dir_path) if os.path.isfile(os.path.join(test_dir_path, file)) ] # and "b1" in file] test_list = sorted(test_list) # print(test_list) test_data = BoltDataset(test_list, test_dir_path, mode="test", transform=data_transform, test_label=test_label) testloader = DataLoader(test_data, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers) model.eval() fn_list = [] pred_list = [] for x, fn in testloader: with torch.no_grad(): x = x.to(self.device) output = model(x) pred = torch.argmax(output, dim=1) fn_list += [n[:-4] for n in fn] pred_list += [p.item() for p in pred] if write_images: for filename, prediction in zip(fn, pred): if prediction == 0: f.copy_file( src_path=f'{test_dir_path}/{filename}', dst_path= f'{write_images}/{self.b_type_list[0]}/{filename}' ) elif prediction == 1: f.copy_file( src_path=f'{test_dir_path}/{filename}', dst_path= f'{write_images}/{self.b_type_list[1]}/{filename}' ) submission = pd.DataFrame({"id": fn_list, "label": pred_list}) if save_csv_path is None: save_csv_path = f'preds_densenet121_dir_{self.img_size}_test_.csv' elif '.csv' not in save_csv_path: save_csv_path = save_csv_path + '/1.csv' submission.to_csv(save_csv_path, # index=False )
def move_images( self, dst_img_dir: str, preserve_filenames: bool=False, overwrite_duplicates: bool=False, update_img_paths: bool=True, overwrite: bool=False, show_pbar: bool=True ): """ Combines all image directories specified in the coco_url of each coco image in self.images to a single image directory. dst_img_dir: The directory where you would like to save the combined image set. preserve_filenames: If False, unique filenames will be generated so as to not create a filename conflict. overwrite_duplicates: Only applicable when preserve_filenames=True. In the event that two images with the same filename are moved to dst_img_dir from two different source folders, an error will be raised if overwrite_duplicates=False. If overwrite_duplicates=True, the second copy will be overwrite the first copy. update_img_paths: If True, all coco_url paths specified in self.images will be updated to reflect the new combined image directory. overwrite: If True, all files in dst_img_dir will be deleted before copying images into the folder. """ used_img_dir_list = [] for coco_image in self.images: used_img_dir = get_dirpath_from_filepath(coco_image.coco_url) if used_img_dir not in used_img_dir_list: check_dir_exists(used_img_dir) used_img_dir_list.append(used_img_dir) if len(used_img_dir_list) == 0: logger.error(f"Couldn't parse used_img_dir_list.") logger.error(f"Are the coco_url paths in your dataset's image dictionary correct?") raise Exception make_dir_if_not_exists(dst_img_dir) if get_dir_contents_len(dst_img_dir) > 0: if overwrite: delete_all_files_in_dir(dst_img_dir, ask_permission=False) else: logger.error(f'dst_img_dir={dst_img_dir} is not empty.') logger.error('Please use overwrite=True if you would like to delete the contents before proceeding.') raise Exception pbar = tqdm(total=len(self.images), unit='image(s)') if show_pbar else None pbar.set_description(f'Moving Images...') for coco_image in self.images: if not preserve_filenames: img_extension = get_extension_from_path(coco_image.coco_url) dst_img_path = get_next_dump_path( dump_dir=dst_img_dir, file_extension=img_extension ) dst_img_path = rel_to_abs_path(dst_img_path) else: img_filename = get_filename(coco_image.coco_url) dst_img_path = f'{dst_img_dir}/{img_filename}' if file_exists(dst_img_path) and not overwrite_duplicates: logger.error(f'Failed to copy {coco_image.coco_url} to {dst_img_dir}') logger.error(f'{img_filename} already exists in destination directory.') logger.error(f'Hint: In order to use preserve_filenames=True, all filenames in the dataset must be unique.') logger.error( f'Suggestion: Either update the filenames to be unique or use preserve_filenames=False' + \ f' in order to automatically assign the destination filename.' ) raise Exception copy_file(src_path=coco_image.coco_url, dest_path=dst_img_path, silent=True) if update_img_paths: coco_image.coco_url = dst_img_path coco_image.file_name = get_filename(dst_img_path) if pbar is not None: pbar.update(1) if pbar is not None: pbar.close()
def replace_bg_wrt_isimg( ndds_data_dir: str, coco_data_dir: str, bg_dirs: list, json_filename: str = None, bg_iscolor: list = None, output_img_dir_name: str = "img_", aug_on: bool = False, aug_json: str = None, show_preview: bool = False, verbose: bool = False, ): make_dir_if_not_exists( os.path.abspath(os.path.join(coco_data_dir, '../..'))) make_dir_if_not_exists(os.path.abspath(os.path.join(coco_data_dir, '..'))) make_dir_if_not_exists(coco_data_dir) # Load NDDS Dataset # ndds_dataset = NDDS_Dataset.load_from_dir( # json_dir=ndds_data_dir, # show_pbar=True # ) coco_dataset = COCO_Dataset.load_from_path( json_path=f"{coco_data_dir}/json/{json_filename}.json", check_paths=False) image_path_list = [] for bg_dir in bg_dirs: image_path_list += dir_contents_path_list_with_extension( dirpath=bg_dir, extension=['.jpg', '.jpeg', '.png']) bg_gen = image_sequence(image_path_list) # pbar = tqdm(ndds_dataset.frames, colour='#44aa44') pbar = tqdm(coco_dataset.images, colour='#44aa44', total=len(coco_dataset.images)) bg_gen = image_sequence(image_path_list) for image in pbar: pbar.set_description("Changing background") # pbar.set_postfix({'file_name': image.file_name}) is_path = ndds_data_dir + '/' + \ image.file_name.split('.')[0]+'.is.'+image.file_name.split('.')[-1] # img_path = ndds_data_dir + if verbose: printj.green(image.coco_url) printj.green(is_path) img = cv2.imread(image.coco_url) is_img = cv2.imread(is_path) is_img2 = is_img.copy() # from PIL import Image # img0 = Image.open(is_path) # colors = img0.convert('RGB').getcolors() # printj.red(colors) if bg_iscolor: mask = create_mask(img=is_img2, color=list(reversed(bg_iscolor)), difference=2, min_limit=0, max_limit=255) else: # img.convert('RGB').getcolors() colors = get_all_colors(img_path=is_path) colors = tuple(sorted(colors, key=itemgetter(0), reverse=True)) _bg_iscolor = list(colors[0][1]) if verbose: printj.cyan( f"\nAll {len(colors)} colors in the image: {colors}") printj.yellow(f'Background color is {_bg_iscolor}') mask = create_mask(img=is_img2, color=list(reversed(_bg_iscolor)), difference=2, min_limit=0, max_limit=255) background = next(bg_gen) if aug_on: aug = aug_flip_and_rotate(aug_json) background = aug(image=np.array(background))['image'] background = resize_img(src=background, size=(img.shape[1], img.shape[0])) # while (img.shape[1] > background.shape[1]) and (img.shape[0] > background.shape[0]): # background1 = cv2.hconcat([background, next(bg_gen)]) # background2 = cv2.hconcat([next(bg_gen), next(bg_gen)]) # background = cv2.vconcat([background1, background2]) # background = background[:img.shape[1], :img.shape[0]] bg = cv2.bitwise_or(background, background, mask=mask) mask = cv2.bitwise_not(mask) fg = cv2.bitwise_or(img, img, mask=mask) final = cv2.bitwise_or(fg, bg) output = os.path.join(coco_data_dir, output_img_dir_name) make_dir_if_not_exists(coco_data_dir) make_dir_if_not_exists(output) collaged_output = os.path.join(output, image.file_name) if show_preview: quit = show_image(final) if quit: break else: cv2.imwrite(collaged_output, final)
def __init__( self, coco_ann_path: str, img_path: str, val_coco_ann_path: str, val_img_path: str, output_dir_path: str, resume: bool = True, class_names: List[str] = None, num_classes: int = None, keypoint_names: List[str] = None, num_keypoints: int = None, model: str = "mask_rcnn_R_50_FPN_1x", instance_train: str = "training_instance1", min_size_train: int = None, max_size_train: int = None, min_size_test: int = None, max_size_test: int = None, max_iter: int = 10000, batch_size_per_image: int = 512, checkpoint_period: int = None, score_thresh: int = None, key_seg_together: bool = False, aug_on: bool = True, train_val: bool = False, aug_settings_file_path: str = None, aug_vis_save_path: str = 'aug_vis.png', show_aug_seg: bool = False, aug_n_rows: int = 3, aug_n_cols: int = 5, aug_save_dims: Tuple[int] = (3 * 500, 5 * 500), device: str = 'cuda', num_workers: int = 2, images_per_batch: int = 2, base_lr: float = 0.003, decrease_lr_by_ratio: float = 0.1, lr_steps: tuple = (30000, ), detectron2_dir_path: str = None, val_on: bool = False, instance_test: str = "test_instance1", val_eval_period: int = 100, vis_period: int = 0, train_type: str = None, ): """ D2Trainer ========= Parameters: ------ output_dir_path: str class_names: List[str] = None, num_classes: int = None, keypoint_names: List[str] = None, num_keypoints: int = None, model: str = "mask_rcnn_R_50_FPN_1x", confidence_threshold: float = 0.5, min_size_train: int = None, max_size_train: int = None, key_seg_together: bool = False, detectron2_dir_path: str = "/home/jitesh/detectron/detectron2" """ self.key_seg_together = key_seg_together self.coco_ann_path = coco_ann_path self.img_path = img_path self.val_coco_ann_path = val_coco_ann_path self.val_img_path = val_img_path self.output_dir_path = output_dir_path self.instance_train = instance_train self.resume = resume self.device = device self.num_workers = num_workers self.images_per_batch = images_per_batch self.batch_size_per_image = batch_size_per_image self.checkpoint_period = checkpoint_period self.score_thresh = score_thresh self.base_lr = base_lr self.decrease_lr_by_ratio = decrease_lr_by_ratio self.lr_steps = lr_steps self.max_iter = max_iter self.val_on = val_on self.instance_test = instance_test self.val_eval_period = val_eval_period self.vis_period = vis_period """ Load annotations json """ with open(self.coco_ann_path) as json_file: self.coco_ann_data = json.load(json_file) self.categories = self.coco_ann_data["categories"] if class_names is None: # self.class_names = [''] self.class_names = [ category["name"] for category in self.categories ] else: self.class_names = class_names if num_classes is None: self.num_classes = len(self.class_names) else: printj.red(f'num_classes: {num_classes}') printj.red(f'len(self.class_names): {len(self.class_names)}') assert num_classes == len(self.class_names) self.num_classes = num_classes if keypoint_names is None: self.keypoint_names = [''] else: self.keypoint_names = keypoint_names if num_keypoints is None: if keypoint_names == ['']: self.num_keypoints = 0 else: self.num_keypoints = len(self.keypoint_names) else: assert num_keypoints == len(self.keypoint_names) self.num_keypoints = num_keypoints self.model = model if "COCO-Detection" in self.model: self.model = self.model train_type = 'bbox' elif "COCO-Keypoints" in self.model: self.model = self.model train_type = 'kpt' elif "COCO-InstanceSegmentation" in self.model: self.model = self.model train_type = 'seg' elif "COCO-PanopticSegmentation" in self.model: self.model = self.model train_type = 'seg' elif "LVIS-InstanceSegmentation" in self.model: self.model = self.model train_type = 'seg' elif "Misc" in model: self.model = model train_type = 'seg' elif "rpn" in model or "fast" in model: self.model = "COCO-Detection/" + model train_type = 'bbox' elif "keypoint" in model: self.model = "COCO-Keypoints/" + model train_type = 'kpt' elif "mask" in model: self.model = "COCO-InstanceSegmentation/" + model train_type = 'seg' elif train_type: self.model = model train_type = train_type else: printj.red.bold_on_black(f'{model} is not in the dictionary.\ Choose the correct model.') raise Exception if ".yaml" in self.model: self.model = self.model else: self.model = self.model + ".yaml" if detectron2_dir_path: model_conf_path = f"{detectron2_dir_path}/configs/{self.model}" else: model_conf_path = model_zoo.get_config_file(self.model) if not file_exists(model_conf_path): printj.red(f"Invalid model: {model}\nOr") printj.red(f"File not found: {model_conf_path}") raise Exception """ register """ register_coco_instances(name=self.instance_train, metadata={}, json_file=self.coco_ann_path, image_root=self.img_path) MetadataCatalog.get( self.instance_train).thing_classes = self.class_names # sys.exit(self.class_names) if val_on: register_coco_instances(name=self.instance_test, metadata={}, json_file=self.val_coco_ann_path, image_root=self.val_img_path) MetadataCatalog.get( self.instance_test).thing_classes = self.class_names """ cfg """ self.cfg = get_cfg() self.cfg.merge_from_file(model_conf_path) self.cfg.DATASETS.TRAIN = tuple([self.instance_train]) self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(self.model) self.cfg.MODEL.ROI_HEADS.NUM_CLASSES = self.num_classes self.cfg.MODEL.ROI_KEYPOINT_HEAD.NUM_KEYPOINTS = self.num_keypoints self.cfg.DATALOADER.NUM_WORKERS = self.num_workers self.cfg.SOLVER.IMS_PER_BATCH = self.images_per_batch self.cfg.SOLVER.BASE_LR = self.base_lr self.cfg.MODEL.DEVICE = self.device self.cfg.OUTPUT_DIR = self.output_dir_path if self.lr_steps: self.cfg.SOLVER.GAMMA = self.decrease_lr_by_ratio self.cfg.SOLVER.STEPS = self.lr_steps if self.max_iter: self.cfg.SOLVER.MAX_ITER = self.max_iter if self.batch_size_per_image: self.cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = self.batch_size_per_image if self.checkpoint_period: self.cfg.SOLVER.CHECKPOINT_PERIOD = self.checkpoint_period if self.vis_period: self.cfg.VIS_PERIOD = self.vis_period if score_thresh: self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = score_thresh if self.val_on: self.cfg.DATASETS.TEST = tuple([self.instance_test]) self.cfg.TEST.EVAL_PERIOD = self.val_eval_period make_dir_if_not_exists(self.cfg.OUTPUT_DIR) if not self.resume: delete_dir_if_exists(self.cfg.OUTPUT_DIR) make_dir_if_not_exists(self.cfg.OUTPUT_DIR) if "mask" in self.model.lower() or "segmentation" in self.model.lower( ): self.cfg.MODEL.MASK_ON = True else: self.cfg.MODEL.MASK_ON = False # self.cfg.MODEL.SEM_SEG_HEAD.LOSS_WEIGHT=0.5 # Train Size Parameters if min_size_train is not None: self.cfg.INPUT.MIN_SIZE_TRAIN = min_size_train if max_size_train is not None: self.cfg.INPUT.MAX_SIZE_TRAIN = max_size_train # Test Size Parameters if min_size_test is not None: self.cfg.INPUT.MIN_SIZE_TEST = min_size_test elif min_size_train is not None: self.cfg.INPUT.MIN_SIZE_TEST = min_size_train if max_size_test is not None: self.cfg.INPUT.MAX_SIZE_TEST = max_size_test elif max_size_train is not None: self.cfg.INPUT.MAX_SIZE_TEST = max_size_train self.cfg.INPUT.MIN_SIZE_TEST = min_size_train """ def train() """ self.aug_settings_file_path = aug_settings_file_path self.aug_on = aug_on self.train_val = train_val self.train_type = train_type self.aug_vis_save_path = aug_vis_save_path self.show_aug_seg = show_aug_seg self.aug_n_rows = aug_n_rows self.aug_n_cols = aug_n_cols self.aug_save_dims = aug_save_dims