#include "stdafx.h"
#include "cv.h"
#include "highgui.h"

using std::cout;
using std::endl;
using namespace cv;

void FindBot( IplImage& img );
void FindColor( IplImage& img, int ); // поиск по цвету, int - код цвета
                            // (0 - оранжевый, 1 - красный, 2 - синий, 3 - зеленый, 4 - желтый)
void Counter( IplImage* img ); // выделение контуров
void Metka( IplImage& img, int& x, int& y, int& r, int area, int color ); // распознавание метки
void ConversionThreshold( IplImage* img ); // пороговое преобразование
void Color( IplImage& img, int& x, int& y, int& r, int area, int color );

IplImage *playing_field = NULL; // картинка для хранения кадра с камеры
IplImage* RGBThreshold = NULL; //cvCreateImage( cvGetSize(playing_field), 8, 3 ); // кадр для хранения порогового преобразования

struct obj {public: char name; CvPoint center; int radius; int area;}; // наименование, центр (x, y) и радиус описанной окружности, площадь фигуры
obj bot, shape;

struct bot {public: int ID; CvPoint center; int radius; int vector;};

int main( int argc, char** argv ) {
    CvCapture *capture = cvCreateCameraCapture( 0 ); // инициализируем камеру
     
    if(capture == NULL) // если камер не обнаружено - выходим
       return 0;

    //cvNamedWindow( "camera", CV_WINDOW_AUTOSIZE ); // Окошко

    while (1) {
       playing_field = cvQueryFrame( capture ); // получаем фрейм из камеры

       if(!playing_field) // если фрейм не получен
          break;

       // Создаем копию кадра
       IplImage* Image = cvCreateImage( cvGetSize(playing_field), 8, 3 ); //кадр для передачи в функции для последовательной обработки
       RGBThreshold = cvCreateImage( cvGetSize(playing_field), 8, 3 ); // кадр для хранения порогового преобразования
        
       cvCopy( playing_field, Image ); // копируем кадр
        
       ConversionThreshold( Image ); // пороговое преобразование
       cvCopy( Image, RGBThreshold ); // копируем кадр порогового перобразования
        
       cvShowImage( "RGB cvThreshold 2", RGBThreshold );
        
       FindBot( *Image ); // поиск ботов в кадре

       cvShowImage( "camera", playing_field );
        
       if(cvWaitKey(1) == 27)
          break; // Если Esc - выходим

       // Освободить ресурсы
       cvReleaseImage( &Image );
       cvReleaseImage( &RGBThreshold );
    }

    // Освободить ресурсы
    cvReleaseCapture( &capture );

    cvDestroyAllWindows(); // закрываем все окна

    return 0;
}

// Поиск ботов
void FindBot( IplImage& Img ) {
    FindColor( Img, 0 ); // выделить цвет
    Counter( &Img ); // выделить контуры
}

// Пороговое преобразование
void ConversionThreshold( IplImage* Image )
{
    // Кадры для хранения отдельных слоёв BGR-изображений
    IplImage* r = cvCreateImage( cvGetSize(Image), IPL_DEPTH_8U, 1 );
    IplImage* g = cvCreateImage( cvGetSize(Image), IPL_DEPTH_8U, 1 );
    IplImage* b = cvCreateImage( cvGetSize(Image), IPL_DEPTH_8U, 1 );

    cvSplit( Image, b, g, r, 0 ); // разбиваем на слои

/*   cvAdaptiveThreshold( r, r, 255, CV_ADAPTIVE_THRESH_GAUSSIAN_C, CV_THRESH_BINARY, 7, 1 );
    cvAdaptiveThreshold( g, g, 255, CV_ADAPTIVE_THRESH_GAUSSIAN_C, CV_THRESH_BINARY, 7, 1 );
    cvAdaptiveThreshold( b, b, 255, CV_ADAPTIVE_THRESH_GAUSSIAN_C, CV_THRESH_BINARY, 7, 1 );
*/
    // Выполняем пороговое преобразование
    cvThreshold( r, r, 120, 250, CV_THRESH_BINARY ); // 150 250
    cvThreshold( g, g, 120, 250, CV_THRESH_BINARY );
    cvThreshold( b, b, 120, 250, CV_THRESH_BINARY );

    // Показать результат порового преобразования
    cvShowImage( "R", r );
    cvShowImage( "G", g );
    cvShowImage( "B", b );

    cvMerge(b, g, r, 0, Image); // суммируем результаты

    // Освобождаем ресурсы
    cvReleaseImage( &r );
    cvReleaseImage( &g );
    cvReleaseImage( &b );
}

// Поиск меток цвета
void FindColor( IplImage& Img, int color ) {
    IplImage* Image = cvCreateImage( cvGetSize(&Img), 8, 3 );
    cvCopy(&Img, Image);

    // Теперь необходимо получить доступ ко всем пикселям.
    uchar* ptr1;
    ptr1 = (uchar*) (Image->imageData );

    // Выделить все пиксели удовлетворяющие искомому цвету (BGR)
    for (int i = 0; i < Image->height; i++)
       for (int j = 0; j < Image->width; j++) {
          switch (color) {
          case 0: // черный цвет
             if (ptr1[j * 3 + i * Image->widthStep] == 0 && // B
             ptr1[j * 3 + 1 + i * Image->widthStep] == 0 && // G
             ptr1[j * 3 + 2 + i * Image->widthStep] == 0) { // R
                ptr1[j * 3 + i * Image->widthStep] = 255;
                ptr1[j * 3 + 1 + i * Image->widthStep] = 255;
                ptr1[j * 3 + 2 + i * Image->widthStep] = 255;
             }
             else {
                ptr1[j * 3 + i * Image->widthStep] = 0;
                ptr1[j * 3 + 1 + i * Image->widthStep] = 0;
                ptr1[j * 3 + 2 + i * Image->widthStep] = 0;
             }
             break;
          case 1: // белый цвет
             if (ptr1[j * 3 + i * Image->widthStep] > 200 && // B
             ptr1[j * 3 + 1 + i * Image->widthStep] > 200 && // G
             ptr1[j * 3 + 2 + i * Image->widthStep] > 200) { // R
                ptr1[j * 3 + i * Image->widthStep] = 255;
                ptr1[j * 3 + 1 + i * Image->widthStep] = 255;
                ptr1[j * 3 + 2 + i * Image->widthStep] = 255;
             }
             else {
                ptr1[j * 3 + i * Image->widthStep] = 0;
                ptr1[j * 3 + 1 + i * Image->widthStep] = 0;
                ptr1[j * 3 + 2 + i * Image->widthStep] = 0;
             }
             break;
          case 2: // синий цвет
             if (ptr1[j * 3 + i * Image->widthStep] > 200 && // B
             ptr1[j * 3 + 1 + i * Image->widthStep] == 0 && // G
             ptr1[j * 3 + 2 + i * Image->widthStep] == 0) { // R
                ptr1[j * 3 + i * Image->widthStep] = 255;
                ptr1[j * 3 + 1 + i * Image->widthStep] = 255;
                ptr1[j * 3 + 2 + i * Image->widthStep] = 255;
             }
             else {
                ptr1[j * 3 + i * Image->widthStep] = 0;
                ptr1[j * 3 + 1 + i * Image->widthStep] = 0;
                ptr1[j * 3 + 2 + i * Image->widthStep] = 0;
             }
             break;
          case 3: // красный цвет
             if (ptr1[j * 3 + i * Image->widthStep] == 0 && // B
                ptr1[j * 3 + 1 + i * Image->widthStep] == 0 && // G
                ptr1[j * 3 + 2 + i * Image->widthStep] > 200) { // R
                ptr1[j * 3 + i * Image->widthStep] = 255;
                ptr1[j * 3 + 1 + i * Image->widthStep] = 255;
                ptr1[j * 3 + 2 + i * Image->widthStep] = 255;
             }
             else {
                ptr1[j * 3 + i * Image->widthStep] = 0;
                ptr1[j * 3 + 1 + i * Image->widthStep] = 0;
                ptr1[j * 3 + 2 + i * Image->widthStep] = 0;
             }
             break;
          case 4: // желтый цвет
             if (ptr1[j * 3 + i * Image->widthStep] == 0 && // B
             ptr1[j * 3 + 1 + i * Image->widthStep] > 200 && // G
             ptr1[j * 3 + 2 + i * Image->widthStep] > 200) { // R
                ptr1[j * 3 + i * Image->widthStep] = 255;
                ptr1[j * 3 + 1 + i * Image->widthStep] = 255;
                ptr1[j * 3 + 2 + i * Image->widthStep] = 255;
             }
             else {
                ptr1[j * 3 + i * Image->widthStep] = 0;
                ptr1[j * 3 + 1 + i * Image->widthStep] = 0;
                ptr1[j * 3 + 2 + i * Image->widthStep] = 0;
             }
             break;
          }
       }

       cvCopy( Image, &Img );

    // Освободить ресурсы
    cvReleaseImage( &Image );
}

// Выделить все найденные объекты
void Counter( IplImage* img ) {
    IplImage* img_gray = cvCreateImage( cvSize(img->width, img->height), 8, 1 );

    CvSeq* contours = 0;
    CvMemStorage* storage = cvCreateMemStorage(0);
    cvCvtColor( img, img_gray, CV_BGR2GRAY );

    // Поиск контуров
    cvFindContours( img_gray, storage, &contours, sizeof(CvContour),
                CV_RETR_EXTERNAL, CV_CHAIN_APPROX_SIMPLE, cvPoint(0,0) );
     
    CvSeq* h_next = 0;

    CvPoint2D32f center;
    float radius;

    // Поиск объектов
    for(CvSeq* c = contours; c != NULL; c = c->h_next) {
       if (c != contours) {
          cvMinEnclosingCircle( h_next, ¢er, &radius ); // поиск центра и радиуса объекта
           
          shape.radius = radius;

          if (radius >= 20 && radius < 50) { // ограничить минимальный и максимальный радиус
             // Координаты центра окружности
             shape.center.x = center.x;
             shape.center.y = center.y;

             if ((shape.center.x - shape.radius > 0 && shape.center.y - shape.radius > 0) && // только если объект полностью в поле кадра
             (shape.center.x + shape.radius < img->width && shape.center.y + shape.radius < img->height)) {          
                shape.area = cvContourArea( h_next ); // запомнить площадь объекта

                for (int color = 1; color < 5; color++) {
                   Metka( *RGBThreshold, shape.center.x, shape.center.y, shape.radius, shape.area, color ); // распознать метку
                }

                cvDrawContours( img, h_next, CV_RGB(255,0,0), CV_RGB(0,255,0), 0, 1, CV_AA, cvPoint(0,0) ); // нарисовать контур объекта
                cvCircle( playing_field, shape.center, shape.radius, CV_RGB(255,0,0), 1, 8, 0 ); // выделить найденный объект в кадре
             }
          }
       }
       h_next = c;
    }
     
    cvShowImage( "Contours", img );

    // Освободить ресурсы
    cvReleaseMemStorage( &storage );
    cvReleaseImage( &img_gray );
}

// Распознать метку
void Metka( IplImage& img, int& x, int& y, int& r, int area, int color ) {

    IplImage* roii = cvCreateImage( cvGetSize(&img), 8, 3 ); // кадр для ROI
    cvCopy( &img, roii ); // скопировать кадр в ROI

    cvSetImageROI( roii, cvRect(x - r, y - r, r * 2, r * 2) ); // устанавливаем ROI с изображением полученного объекта
    CvPoint  pt1, pt2; pt1.x = x - r; pt1.y = y - r; pt2.x = pt1.x + r * 2; pt2.y = pt1.y + r * 2;
    FindColor( *roii, color ); // выделить цвет

    Color( *roii, x, y, r, area, color ); // пометить найденный цвет на метке

    cvShowImage( "ROI", roii );

    // Освободить ресурсы
    cvResetImageROI( roii ); // сброс ROI
    cvReleaseImage( &roii );
}

void Color( IplImage& img, int& x, int& y, int& r, int area, int color )
{    
    IplImage* img_gray = cvCreateImage( cvGetSize(&img), 8, 1 );
    cvConvertImage( &img, img_gray, CV_BGR2GRAY );
        
    //cvCvtColor( roii, img_gray, CV_BGR2GRAY );
    CvSeq* contours = 0;
    CvMemStorage* storage = cvCreateMemStorage(0);
     
    // Поиск контуров
    cvFindContours( img_gray, storage, &contours, sizeof(CvContour),
                CV_RETR_LIST, CV_CHAIN_APPROX_SIMPLE, cvPoint(0, 0) );

    CvSeq* h_next = 0;
    CvPoint2D32f center;
    float radius;

    // Выбор цвета выделения
    int red, grin, blu;
    switch (color) {
    case 0:
       red = 0; grin = 255; blu = 255;
       break;
    case 1:
       red = 255; grin = 255; blu = 0;
       break;
    case 2:
       red = 0; grin = 0; blu = 255;
       break;
    case 3:
       red = 255; grin = 0; blu = 0;
       break;
    case 4:
       red = 255; grin = 237; blu = 0;
       break;
    }

    // Поиск меток
    for (CvSeq* c = contours; c != NULL; c = c->h_next) {
       if (c != contours) {
          cvMinEnclosingCircle( h_next, ¢er, &radius ); // поиск центра и радиуса объекта
           
          int rad = radius;

          int are = cvContourArea( h_next );
           
          if ((are < area / 2 && are > area / 10) && // площадь цветовой метки должен быть в 3 раза меньше площади объекта
          ((center.x + rad) < r * 2 && (center.y + rad) < r * 2) && // только если объект полностью в поле ROI
          (center.x - rad > 0 && center.y - rad > 0)) {             
             // Координаты центра окружности
             CvPoint p;
             p.x = (x - r) + center.x;
             p.y = (y - r) + center.y;

             cvCircle( playing_field, p, rad, CV_RGB(red, grin, blu), 1, 8, 0 ); // выделить найденный объект в кадре
          }
       }
       h_next = c;
    }
    cvReleaseMemStorage( &storage );
    cvReleaseImage( &img_gray );
}