def main(base_model_name, weights_file, image_source, predictions_file, img_format='jpg'): # load samples if os.path.isfile(image_source): image_dir, samples = image_file_to_json(image_source) else: image_dir = image_source samples = image_dir_to_json(image_dir, img_type='jpg') # build model and load weights nima = Nima(base_model_name, weights=None) nima.build() nima.nima_model.load_weights(weights_file) # initialize data generator data_generator = TestDataGenerator(samples, image_dir, 64, 10, nima.preprocessing_function(), img_format=img_format) # get predictions predictions = predict(nima.nima_model, data_generator) # calc mean scores and add to samples for i, sample in enumerate(samples): sample['mean_score_prediction'] = calc_mean_score(predictions[i]) print(json.dumps(samples, indent=2)) if predictions_file is not None: save_json(samples, predictions_file)
def predict(self, image_source): img_format = "jpg" if os.path.isfile(image_source): image_dir, samples = image_file_to_json(image_source) else: image_dir = image_source samples = image_dir_to_json(image_dir, img_type='jpg') data_generator = TestDataGenerator(samples, image_dir, 64, 10, self.a_nima.preprocessing_function(), img_format=img_format) a_predictions = predict(self.a_nima.nima_model, data_generator) for i, sample in enumerate(samples): sample['a_mean_score_prediction'] = calc_mean_score(a_predictions[i]) data_generator = TestDataGenerator(samples, image_dir, 64, 10, self.t_nima.preprocessing_function(), img_format=img_format) t_predictions = predict(self.t_nima.nima_model, data_generator) for i, sample in enumerate(samples): sample['t_mean_score_prediction'] = calc_mean_score(t_predictions[i]) print(json.dumps(samples, indent=2)) return samples
def start_Button_clicked(self): print('Start Button Clicked') SelfDestructingBox.showWithTimeout(10, "Please wait while your images are analyzed.\nAnother message will appear when the process is complete.", "In progress") QApplication.processEvents() # build model and load weights nima = Nima("MobileNet", weights=None) nima.build() nima.nima_model.load_weights("weights_mobilenet_aesthetic_0.07.hdf5") totalSamples = [] j = 0 for (root, directories, files) in os.walk(self.selected_directory, topdown=True): print("[INFO] ROOT: " + root) self.samples = [] # Only grab the jpg images img_paths = glob.glob(os.path.join(root, '*.jpg')) for img_path in img_paths: print("[INFO] IMG PATH IN ROOT: " + img_path) totalSamples.append({'image_path': img_path}) self.importedFiles_id = os.path.basename(img_path).split('.')[0] self.samples.append({'image_id': self.importedFiles_id}) # initialize data generator data_generator = TestDataGenerator(self.samples, root, 64, 10, nima.preprocessing_function(), img_format='jpg') # TODO: Circumvent multiprocessing error on Windows predictions = nima.nima_model.predict_generator(data_generator, workers=1, use_multiprocessing=False, verbose=1) # print(predictions) # calc mean scores and add to samples for i, sample in enumerate(self.samples): sample['mean_score_prediction'] = calc_mean_score(predictions[i]) totalSamples[j]['mean_score_prediction'] = calc_mean_score(predictions[i]) j += 1 print("[INFO] Directory specific samples:") print(json.dumps(self.samples, indent=2)) # sort totalSamples by score print("[INFO] Sorted (descending) Total samples:") totalSamples.sort(key=lambda i: i['mean_score_prediction'], reverse=True) print(json.dumps(totalSamples, indent=2)) modelPath = os.path.sep.join(["personalModel", "model.cpickle"]) bUsePM = False paths_with_scores = [] if os.path.exists(modelPath): print("[INFO] MODEL PATH EXISTS: {}".format(modelPath)) bUsePM = True # Initialize model to use its feature extraction method personal_model = PersonalModel(modelPath) # Grab top 50% of General Model's output top_half = totalSamples[:len(totalSamples)//2] feat_vecs = [] img_paths = [] for datum in top_half: feat_vecs.append(personal_model.getFeatureVector(datum['image_path'])) img_paths.append(datum['image_path']) feat_vecs_np = np.array(feat_vecs) feat_vecs_np = np.squeeze(feat_vecs_np) predictions = personal_model.model.predict(feat_vecs_np) print("[INFO] Number of samples for PM: {}".format(len(top_half))) print("[INFO] Predictions: {}".format(predictions)) paths_with_scores = list(zip(img_paths, predictions)) # sort (descending) by the second tuple value (the score) paths_with_scores.sort(key=lambda x: x[1], reverse=True) print("[INFO] ***paths with scores***") for tup in paths_with_scores: print("[INFO] {}, {}".format(tup[0], tup[1])) # Creates new directory name. Checks if previous directory exists. self.new_dir_title = self.create_directory_name() # create new directory to hold top rated pics self.new_dir = os.path.sep.join([self.selected_directory, self.new_dir_title]) if os.path.exists(self.new_dir): shutil.rmtree(self.new_dir) os.mkdir(self.new_dir) # determine number of pics to copy into new dir # 20% for now num_to_pick = int(len(totalSamples) * 0.2) # so before copying pics into new dir, see if we need to use the PM # if so, pass 40% of PM's picks into if bUsePM: print("[INFO] PM Used") for i, sample in enumerate(paths_with_scores): if i == num_to_pick: break shutil.copy(sample[0], self.new_dir) else: # since totalSamples is now sorted with the highest rated # pics first, the loop will only copy the first 'num_to_pick' images print("[INFO] GENERAL used") for i, sample in enumerate(totalSamples): if i == num_to_pick: break shutil.copy(sample['image_path'], self.new_dir) self.done_msg = QtGui.QMessageBox() self.done_msg.setWindowTitle("Processing Complete!") self.done_msg.setText("The best photos were copied into a new directory.\nYou can find this new directory inside of the directory you initially chose.") self.done_msg.setIcon(QMessageBox.Information) self.done_msg.exec_() print("[INFO] Done copying into new dir.")