www.gusucode.com > Deep Learning for Defect Detection on Raspberry Pi 源码程序matlab > Deep Learning for Defect Detection on Raspberry Pi/RaspPi/main_nutsDet.cpp

    /* Copyright 2018 The MathWorks, Inc. */

#if defined(_WIN32) || (_MSC_VER)
#define VC_MODE
#endif

/* Headers */
#include "cnn_exec.hpp"
#include <cuda.h>
#include <cuda_runtime_api.h>

#include <stdio.h>
#include <string.h>
#include <math.h>
#include <iostream>

/* Opencv includes*/
#include "opencv2/opencv.hpp"

#ifdef VC_MODE
#include <sys/types.h>
#include <sys/timeb.h>
#else
#include <sys/time.h>
#endif

using namespace cv;

/* pre/post processing source headers*/
#include "myNDNet_Preprocess.h"
#include "myNDNet_Postprocess.h"

#ifdef VC_MODE
double cpuSecond()
{
    _timeb tp;
    _ftime(&tp);
    return ((double)tp.time + (double)tp.millitm / 1000.0);
}
#else
double cpuSecond()
{
    struct timeval tp;
    gettimeofday(&tp, NULL);
    return ((double)tp.tv_sec + (double)tp.tv_usec * 1.e-6);
}
#endif

void readData(float *input, unsigned char nutsroi[], int num)
{ 
    for(int j=0;j<128*128;j++)
    {       
        //BGR to RGB
        input[j]=(float)(nutsroi[num*128*128+j]);
    }
}

int main(int argc, char* argv[]) {
    
    unsigned int wi = 320;
    unsigned int he = 240;
    unsigned int ch = 3;
    unsigned int size = wi*he*ch;
    
    int sw;
    sw = atoi(argv[1]);
    
    /* Input arguments */
    VideoCapture cap;
    if ( sw < 3 ) {
        if ( sw < 2 ) {
            //VideoCapture cap(atoi(argv[2])); //use device number for camera.
            cap.open(atoi(argv[2])); //use device number for camera.
            if (!cap.isOpened()) {
                printf("Could not open the video capture device.\n");
                return -1;
            }
            cap.set(CAP_PROP_FRAME_WIDTH, wi);
            cap.set(CAP_PROP_FRAME_HEIGHT, he);
            //cap.set(CAP_PROP_AUTOFOCUS, 0);
            //cap.set(28, 70);
        } else {
            //VideoCapture cap(argv[2]);
            cap.open(argv[2]);
            if (!cap.isOpened()) {
                printf("Could not open the video capture device.\n");
                return -1;
            }
        }
    }
    
    namedWindow("Nuts Defect Detection Demo",CV_WINDOW_NORMAL);

    double fps=0;	
    unsigned char output[size];
    unsigned char output2[size];
    unsigned char nuts[65536];
	double numnuts;
    double bbox[16];
    double scores[8];
    double iStart, iElaps;
    int i;
    int frame=0;
    Mat inImg;
    Mat im;
    
    float *inputBuffer = (float*)calloc(sizeof(float),128*128);
    float *outputBuffer = (float*)calloc(sizeof(float),2);

    CnnMain* net = new CnnMain;

    net->batchSize = 1;
    net->setup();
    
    /* Start reading frames */
    for(;;)
    {
		
        if ( sw < 3 ) {
            cap >> inImg;
        } else {
            inImg = imread(argv[2], 1);
        }
            
		if (inImg.empty()) {
            if ( sw == 2 ) {
                cap.set(CAP_PROP_POS_FRAMES, 0);
                cap >> inImg;
            } else {
                break;
            }
        }

		Size size(wi,he);
		resize(inImg,im,size);

        iStart = cpuSecond();
        myNDNet_Preprocess(im.data, output, nuts, &numnuts, bbox);
        
        if (numnuts > 0) { 
            for (i = 0; i < numnuts; i++) {
                readData(inputBuffer, nuts, i);
                cudaMemcpy(net->inputData, inputBuffer, sizeof(float)*128*128, cudaMemcpyHostToDevice );
                net->predict();
                cudaMemcpy(outputBuffer, net->outputData, sizeof(float)*2, cudaMemcpyDeviceToHost );
                scores[i*2]   = *outputBuffer;
                scores[i*2+1] = *(outputBuffer+1);
            }
        }
        
        myNDNet_Postprocess(output, numnuts, bbox, scores, output2);
        iElaps = cpuSecond() - iStart;
        
        Size sizeo(wi,he);
        Mat outImg(sizeo,im.type(),output2);
        
        fps = fps*.9+1.0/iElaps*.1;
        
        char strbuf[50];
        sprintf (strbuf, "%.2f FPS", fps);
        putText(outImg, strbuf, cvPoint(100,30), CV_FONT_HERSHEY_DUPLEX, 0.7, CV_RGB(255,241,15), 1.5);
               
        imshow("Nuts Defect Detection Demo", outImg);
        
        
        if( waitKey(50)%256 == 27 ) break; // stop capturing by pressing ESC
        
        if( sw == 3 ) {
            if( frame > 100) {
                imwrite("outImg.png", outImg);
                break;
            } else {
                frame = frame + 1;
            }
        }
    }
    
    destroyWindow("Nuts Defect Detection Demo");
    
    return 0;
}