def train(args):
""" Fit a model's parameters given the parameters specified in args.
"""
X = load_image(args.train_data)
# Initialize appropriate algorithm
if args.algorithm == 'mrf':
try:
model = pickle.load(open('./models/image_2_model', 'rb'))
except (EOFError, FileNotFoundError, IOError) as e:
model = models.MRF(J=args.edge_weight,
K=args.num_states,
n_em_iter=args.n_em_iterations,
n_vi_iter=args.n_vi_iterations)
# Train the model
model.fit(X=X)
# Save the model
pickle.dump(model, open(args.model_file, 'wb'))
else:
raise Exception("Algorithm argument not recognized")
# predict most likely latent states for each of the pixels
preds = model.predict(X)
# output model predictions
np.savetxt(args.predictions_file, preds, fmt='%s', delimiter='\t')
# Visualize clustering
visualize_clustering(args.visualize_predictions_file, preds, X,
args.num_states)
import models
from data import load_image
import numpy as np
import pickle
import matplotlib
matplotlib.use('PS')
import matplotlib.pyplot as plt
X = load_image('./data/image_1.jpg')
np.random.seed(1)
model = models.MRF(J=1.2, K=4, n_em_iter=3, n_vi_iter=3) #TODO revert to k=4, both iters 3
model.fit(X=X)
pickle.dump(model, open('./models/testing', 'wb'))
predictions = model.predict(X=X)
print(predictions)
out_file = 'image_1.jpg'
fig = plt.figure()
ax = fig.add_subplot(1, 2, 1)
plt.imshow(X, cmap='gray', vmin=0, vmax=255)
plt.title('Raw data')
plt.axis('off')
ax = fig.add_subplot(1, 2, 2)
plt.imshow(predictions)
plt.title('Segmentation')
plt.axis('off')
plt.savefig('./out/' + out_file)