patch_kmeans1.diff

use parallel_for labels assignment if TBB is available - Boris Mansencal, 2011-07-29 04:31 pm

Download (2.2 kB)

 
modules/core/src/matrix.cpp (working copy)
2259 2259
    }
2260 2260
}
2261 2261

  
2262
#ifdef HAVE_TBB 
2263

  
2264
class KMeansDistanceComputer
2265
{
2266
  float *distances;
2267
  int *labels;
2268
  const cv::Mat& data;
2269
  const cv::Mat& centers;
2270
  
2271
public:
2272

  
2273
  KMeansDistanceComputer(float *_distances,
2274
			 int *_labels,
2275
			 const cv::Mat& _data,
2276
			 const cv::Mat& _centers)
2277
    : distances(_distances),
2278
      labels(_labels),
2279
      data(_data),
2280
      centers(_centers)
2281
  {
2282
    CV_DbgAssert(centers.cols == data.cols);
2283
  }
2284

  
2285
  void operator()(const cv::BlockedRange& range) const
2286
  {
2287
    const int begin = range.begin();
2288
    const int end = range.end();
2289
    const int K = centers.rows;
2290
    const int dims = centers.cols;
2291

  
2292
    const float *sample;
2293
    for( int i=begin; i<end; ++i) 
2294
    {
2295

  
2296
      sample = data.ptr<float>(i);
2297
      int k_best = 0;
2298
      double min_dist = DBL_MAX;
2299
      
2300
      for( int k = 0; k < K; k++ )
2301
      {
2302
	  const float* center = centers.ptr<float>(k);
2303
	  const double dist = distance(sample, center, dims);
2304
	  
2305
	  if( min_dist > dist )
2306
	  {
2307
	      min_dist = dist;
2308
	      k_best = k;
2309
	  }
2310
      }
2311
      
2312
      distances[i] = min_dist;
2313
      labels[i] = k_best;
2314
    }
2315
    
2316
  }
2317
  
2318
};
2319

  
2320
#endif //HAVE_TBB
2321

  
2262 2322
}
2263 2323
    
2264 2324
double cv::kmeans( InputArray _data, int K,
......
2338 2398
        }
2339 2399
    }
2340 2400

  
2401
#ifdef HAVE_TBB 
2402
    cv::Mat dists(N, 1, CV_32F);
2403
#endif //HAVE_TBB
2404
    
2341 2405
    for( a = 0; a < attempts; a++ )
2342 2406
    {
2343 2407
        double max_center_shift = DBL_MAX;
......
2422 2486
            }
2423 2487

  
2424 2488
            // assign labels
2489
#ifndef HAVE_TBB 
2425 2490
            compactness = 0;
2426 2491
            for( i = 0; i < N; i++ )
2427 2492
            {
......
2444 2509
                compactness += min_dist;
2445 2510
                labels[i] = k_best;
2446 2511
            }
2512
#else
2513
	    float* dist = dists.ptr<float>(0);
2514
	    cv::parallel_for(cv::BlockedRange(0, N),
2515
			     KMeansDistanceComputer(dist, labels, data, centers));
2516
	    compactness = 0;
2517
	    for( i = 0; i < N; i++ )
2518
	    {
2519
	      compactness += dist[i];
2520
	    }
2521

  
2522
#endif //HAVE_TBB
2447 2523
        }
2448 2524

  
2449 2525
        if( compactness < best_compactness )