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
