diff --git a/random/src/main.cpp b/random/src/main.cpp
index 88aeada..d75a9a1 100644
--- a/random/src/main.cpp
+++ b/random/src/main.cpp
@@ -38,7 +38,7 @@ enum SegmentType {
 	ST_OBJECTNESS,
 };
 
-cv::Mat get_bboxes_(const cv::Mat & image, const cv::Mat & seg, const cv::Mat & edge, uint8_t flags, const cv::Mat & weights, uint32_t n, std::string selection_prior, std::string segment_type, float threshold) {
+cv::Mat get_bboxes_(const cv::Mat & image, const cv::Mat & seg, const cv::Mat & edge, uint8_t flags, const cv::Mat & weights, uint32_t n, std::string selection_prior, std::string segment_type, uint32_t max_size) {
 	double max_id_;
 	cv::minMaxIdx(seg, nullptr, &max_id_);
 	int max_id = max_id_;
@@ -150,10 +150,14 @@ cv::Mat get_bboxes_(const cv::Mat & image, const cv::Mat & seg, const cv::Mat &
 
 	for (uint32_t i = 0; i < n; i++) {
 		std::shared_ptr<Segment> s = segments[prior.poll()];
-		RandomStoppingCriterion stop(threshold);
+		//RandomStoppingCriterion stop(threshold);
+		uint32_t seg_size = rand() % max_size;
 		cv::Rect r(s->min_p, s->max_p);
 
-		while (s->neighbours.size() && stop.stop(image, r) == false) {
+		for (int j = 0; j < seg_size; j++) {
+		//while (s->neighbours.size() && stop.stop(image, r) == false) {
+			if (s->neighbours.size() == 0)
+				break;
 #ifdef DEBUG
 			cv::Mat red = cv::Mat::zeros(seg.size(), CV_8UC1);//edge * 0.5;
 			cv::Mat green;
@@ -237,13 +241,13 @@ cv::Mat get_bboxes_(const cv::Mat & image, const cv::Mat & seg, const cv::Mat &
 	return bboxes;
 }
 
-PyObject * get_bboxes(PyObject * image_, PyObject * seg_, PyObject * edge_, uint8_t flags, PyObject * weights_, uint32_t n, std::string selection_prior, std::string segment_type, float threshold) {
+PyObject * get_bboxes(PyObject * image_, PyObject * seg_, PyObject * edge_, uint8_t flags, PyObject * weights_, uint32_t n, std::string selection_prior, std::string segment_type, uint32_t max_size) {
 	NDArrayConverter cvt;
 	cv::Mat image   = cvt.toMat(image_);
 	cv::Mat seg     = cvt.toMat(seg_);
 	cv::Mat edge    = cvt.toMat(edge_);
 	cv::Mat weights = cvt.toMat(weights_);
-	return cvt.toNDArray(get_bboxes_(image, seg, edge, flags, weights, n, selection_prior, segment_type, threshold));
+	return cvt.toNDArray(get_bboxes_(image, seg, edge, flags, weights, n, selection_prior, segment_type, max_size));
 }
 
 static void init_ar() {
@@ -285,7 +289,7 @@ int main(int argc, char * argv[]) {
 // 	cv::namedWindow("Image", cv::WINDOW_NORMAL);
 //	while (cv::waitKey() != 'q') {
 		std::clock_t begin = std::clock();
-		cv::Mat bboxes = get_bboxes_(image, seg, edge, COLOR_SIMILARITY /*| TEXTURE_SIMILARITY*/  | SIZE_SIMILARITY | BBOX_SIMILARITY, weights, iterations, selection_prior, segment_type, 0.85f); 
+		cv::Mat bboxes = get_bboxes_(image, seg, edge, COLOR_SIMILARITY /*| TEXTURE_SIMILARITY*/  | SIZE_SIMILARITY | BBOX_SIMILARITY, weights, iterations, selection_prior, segment_type, 18);
 		std::clock_t end = std::clock();
 		double elapsed_secs = double(end - begin) / CLOCKS_PER_SEC;
 		std::cout << "Times passed in seconds: " << elapsed_secs << std::endl;
diff --git a/random/src/stochastic.cpp b/random/src/stochastic.cpp
index 659dace..e36938a 100644
--- a/random/src/stochastic.cpp
+++ b/random/src/stochastic.cpp
@@ -8,7 +8,7 @@
 
 #include "connection.h"
 #include "adjacency.h"
-// #include "location_prior.h"
+#include "location_prior.h"
 
 #include "random_stopping_criterion.h"
 
@@ -19,7 +19,7 @@ std::random_device rd_;
 std::mt19937 gen_(rd_());
 
 cv::Mat get_bboxes_(const cv::Mat & seg, uint32_t max_size, float sigma, uint32_t n) {
-	//LocationPrior lp(sigma);
+	LocationPrior lp(sigma);
 
 	double max_id_;
 	cv::minMaxIdx(seg, nullptr, &max_id_);
@@ -64,6 +64,8 @@ cv::Mat get_bboxes_(const cv::Mat & seg, uint32_t max_size, float sigma, uint32_
 	cv::Mat bboxes;
 	float similarity_sum = 0.f;
 	for (auto & s: segments) {
+		if (s->size == 0)
+			continue;
 		cv::Mat bbox = cv::Mat(1, 4, CV_32SC1);
 		bbox.at<int>(0) = s->min_p.x;
 		bbox.at<int>(1) = s->min_p.y;
@@ -75,21 +77,21 @@ cv::Mat get_bboxes_(const cv::Mat & seg, uint32_t max_size, float sigma, uint32_
 			cv::vconcat(bboxes, bbox, bboxes);
 	}
 
-// 	SelectionPriorMap prior;
-// 
-// 	prior = lp.computeSelectionPrior(image, segments);
-// #ifdef DEBUG
-// 	prior.visualize(seg);
-// #endif
+	SelectionPriorMap prior;
+
+	prior = lp.computeSelectionPrior(seg, segments);
+#ifdef DEBUG
+	prior.visualize(seg);
+#endif
 
 #ifdef DEBUG
 	cv::namedWindow("SelectionLikelihood", cv::WINDOW_AUTOSIZE);
 #endif
 
 	for (uint32_t i = 0; i < n; i++) {
-// 		std::shared_ptr<Segment> s = segments[prior.poll()];
+		std::shared_ptr<Segment> s = segments[prior.poll()];
+// 		std::shared_ptr<Segment> s = segments[rand() % segments.size()];
 		uint32_t seg_size = rand() % max_size;
-		std::shared_ptr<Segment> s = segments[rand() % segments.size()];
 		//RandomStoppingCriterion stop(threshold);
 		cv::Rect r(s->min_p, s->max_p);