def check_extend_dataset(main_window, dataset_dir, prev_fnames,
proj_file_path):
all_image_names = [f for f in os.listdir(dataset_dir) if is_image(f)]
new_image_names = [f for f in all_image_names if f not in prev_fnames]
button_reply = QtWidgets.QMessageBox.question(
main_window, 'Confirm',
f"There are {len(new_image_names)} new images in the dataset."
" Are you sure you want to extend the project to include these new images?",
QtWidgets.QMessageBox.Yes | QtWidgets.QMessageBox.No,
QtWidgets.QMessageBox.No)
if button_reply == QtWidgets.QMessageBox.Yes:
# shuffle the new file names
shuffle(new_image_names)
# load the project json for reading and writing
settings = json.load(open(proj_file_path, 'r'))
# read the file_names
all_file_names = settings['file_names'] + new_image_names
settings['file_names'] = all_file_names
# Add the new_files to the list
# then save the json again
json.dump(settings, open(proj_file_path, 'w'), indent=4)
return True, all_file_names
else:
return False, all_image_names
def validate(self):
self.proj_name = self.name_edit_widget.name
if not self.proj_name:
self.info_label.setText("Name must be specified to create project")
self.create_project_btn.setEnabled(False)
return
if not self.selected_dir:
self.info_label.setText(
"Directory must be specified to create project")
self.create_project_btn.setEnabled(False)
return
if not self.use_random_weights and not self.selected_model:
self.info_label.setText(
"Starting model must be specified to create project")
self.create_project_btn.setEnabled(False)
return
cur_files = os.listdir(self.selected_dir)
cur_files = [is_image(f) for f in cur_files]
if not cur_files:
message = "Folder contains no images."
self.info_label.setText(message)
self.create_project_btn.setEnabled(False)
return
self.project_location = os.path.join('projects', self.proj_name)
if os.path.exists(os.path.join(self.sync_dir, self.project_location)):
self.info_label.setText(
f"Project with name {self.proj_name} already exists")
self.create_project_btn.setEnabled(False)
else:
self.info_label.setText(
f"Project location: {self.project_location}")
self.create_project_btn.setEnabled(True)
def create_project(self):
project_name = self.proj_name
project_location = Path(self.project_location)
dataset_path = os.path.abspath(self.selected_dir)
datasets_dir = str(self.sync_dir / 'datasets')
if not dataset_path.startswith(datasets_dir):
message = (
"When creating a project the selected dataset must be in "
"the datasets folder. The selected dataset is "
f"{dataset_path} and the datasets folder is "
f"{datasets_dir}.")
QtWidgets.QMessageBox.about(self, 'Project Creation Error',
message)
return
os.makedirs(self.sync_dir / project_location)
proj_file_path = (self.sync_dir / project_location /
(project_name + '.seg_proj'))
os.makedirs(self.sync_dir / project_location / 'annotations' / 'train')
os.makedirs(self.sync_dir / project_location / 'annotations' / 'val')
os.makedirs(self.sync_dir / project_location / 'segmentations')
os.makedirs(self.sync_dir / project_location / 'models')
os.makedirs(self.sync_dir / project_location / 'messages')
os.makedirs(self.sync_dir / project_location / 'logs')
if self.use_random_weights:
original_model_file = 'random weights'
else:
model_num = 1
model_name = str(model_num).zfill(6)
model_name += '_' + str(int(round(time.time()))) + '.pkl'
shutil.copyfile(
self.selected_model,
self.sync_dir / project_location / 'models' / model_name)
original_model_file = self.selected_model
dataset = os.path.basename(dataset_path)
# get files in random order for training.
all_fnames = os.listdir(dataset_path)
# images only
all_fnames = [a for a in all_fnames if is_image(a)]
all_fnames = sorted(all_fnames)
random.shuffle(all_fnames)
# create project file.
project_info = {
'name': project_name,
'dataset': dataset,
'original_model_file': original_model_file,
'location': str(PurePosixPath(project_location)),
'file_names': all_fnames
}
# 'classes': self.palette_edit_widget.get_brush_data()
with open(proj_file_path, 'w') as json_file:
json.dump(project_info, json_file, indent=4)
self.created.emit(proj_file_path)
self.close()
def run(self):
while True:
done_fnames = os.listdir(self.segment_dir)
done_fnames = [f for f in done_fnames if is_image(f)]
count = len(done_fnames)
if count >= self.total_images:
self.done.emit()
break
else:
self.progress_change.emit(count, self.total_images)
time.sleep(0.2)
def segment_folder(self):
selected_models = self.selected_models
input_dir = self.input_dir
output_dir = self.output_dir
all_fnames = os.listdir(str(input_dir))
all_fnames = [f for f in all_fnames if is_image(f)]
# need to make sure all train photos are copied now.
content = {
"model_paths": selected_models,
"dataset_dir": input_dir,
"seg_dir": output_dir,
"file_names": all_fnames
}
send_instruction('segment', content, self.instruction_dir, self.sync_dir)
self.progress_widget = SegmentProgressWidget()
self.progress_widget.run(output_dir, len(all_fnames))
self.progress_widget.show()
self.close()