<reponame>khm159/MFC_ImageClassification_libtorch_vs22
﻿
// ImageClassificationDlg.cpp: 구현 파일
//
#include "pch.h"
#include "framework.h"
#include "ImageClassification.h"
#include "ImageClassificationDlg.h"
#include "afxdialogex.h"
#include <iostream>
#pragma comment(linker, "/entry:WinMainCRTStartup /subsystem:console")

#define IMAGE_SIZE 224
#define CHANNELS 3


#ifdef _DEBUG
#define new DEBUG_NEW
#endif


// 응용 프로그램 정보에 사용되는 CAboutDlg 대화 상자입니다.

class CAboutDlg : public CDialogEx
{
public:
	CAboutDlg();

// 대화 상자 데이터입니다.
#ifdef AFX_DESIGN_TIME
	enum { IDD = IDD_ABOUTBOX };
#endif

	protected:
	virtual void DoDataExchange(CDataExchange* pDX);    // DDX/DDV 지원입니다.

// 구현입니다.
protected:
	DECLARE_MESSAGE_MAP()
};

CAboutDlg::CAboutDlg() : CDialogEx(IDD_ABOUTBOX)
{
}

void CAboutDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialogEx::DoDataExchange(pDX);
}

BEGIN_MESSAGE_MAP(CAboutDlg, CDialogEx)
END_MESSAGE_MAP()


// CImageClassificationDlg 대화 상자

bool loadimagenetlabel(std::string file_name, std::vector<std::string>& labels)
{
	std::ifstream ifs(file_name);
	if (!ifs)
	{
		return false;
	}
	std::string line;
	while (std::getline(ifs, line))
	{
		labels.push_back(line);
	}
	return true;
}

CImageClassificationDlg::CImageClassificationDlg(CWnd* pParent /*=nullptr*/)
	: CDialogEx(IDD_IMAGECLASSIFICATION_DIALOG, pParent)
{
	m_hIcon = AfxGetApp()->LoadIcon(IDR_MAINFRAME);
}

void CImageClassificationDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialogEx::DoDataExchange(pDX);
	DDX_Control(pDX, IDC_PICTUREBOX, m_picture_box);
}

BEGIN_MESSAGE_MAP(CImageClassificationDlg, CDialogEx)
	ON_WM_SYSCOMMAND()
	ON_WM_PAINT()
	ON_WM_QUERYDRAGICON()
	ON_BN_CLICKED(IDC_BUTTON_IMAGELOAD, &CImageClassificationDlg::OnBnClickedButtonImageload)
	ON_MESSAGE(WM_USER_MESSAGE1, OnImageClassification)
END_MESSAGE_MAP()


// CImageClassificationDlg 메시지 처리기

BOOL CImageClassificationDlg::OnInitDialog()
{
	CDialogEx::OnInitDialog();

	// 시스템 메뉴에 "정보..." 메뉴 항목을 추가합니다.

	// IDM_ABOUTBOX는 시스템 명령 범위에 있어야 합니다.
	ASSERT((IDM_ABOUTBOX & 0xFFF0) == IDM_ABOUTBOX);
	ASSERT(IDM_ABOUTBOX < 0xF000);

	CMenu* pSysMenu = GetSystemMenu(FALSE);
	if (pSysMenu != nullptr)
	{
		BOOL bNameValid;
		CString strAboutMenu;
		bNameValid = strAboutMenu.LoadString(IDS_ABOUTBOX);
		ASSERT(bNameValid);
		if (!strAboutMenu.IsEmpty())
		{
			pSysMenu->AppendMenu(MF_SEPARATOR);
			pSysMenu->AppendMenu(MF_STRING, IDM_ABOUTBOX, strAboutMenu);
		}
	}

	// 이 대화 상자의 아이콘을 설정합니다.  응용 프로그램의 주 창이 대화 상자가 아닐 경우에는
	//  프레임워크가 이 작업을 자동으로 수행합니다.
	SetIcon(m_hIcon, TRUE);			// 큰 아이콘을 설정합니다.
	SetIcon(m_hIcon, FALSE);		// 작은 아이콘을 설정합니다.

	// TODO: 여기에 추가 초기화 작업을 추가합니다.
	
	// 1. 이미지 분류 모델 로드 하기 
	char * model_path = "./resnet18.pt";
	char* label_path = "./imagenet_label.txt";
	this->Model = torch::jit::load(model_path);

	// 리소스 아이디와 변수 연결 
	m_edit_class = (CEdit*)GetDlgItem(IDC_EDIT_CLASS);
	m_edit_probability = (CEdit*)GetDlgItem(IDC_EDIT_PROBABILITY);

	// 2. 이미지넷 레이블 로드하기 
	loadimagenetlabel(label_path, this->labels);

	return TRUE;  // 포커스를 컨트롤에 설정하지 않으면 TRUE를 반환합니다.
}

void CImageClassificationDlg::OnSysCommand(UINT nID, LPARAM lParam)
{
	if ((nID & 0xFFF0) == IDM_ABOUTBOX)
	{
		CAboutDlg dlgAbout;
		dlgAbout.DoModal();
	}
	else
	{
		CDialogEx::OnSysCommand(nID, lParam);
	}
}

// 대화 상자에 최소화 단추를 추가할 경우 아이콘을 그리려면
//  아래 코드가 필요합니다.  문서/뷰 모델을 사용하는 MFC 애플리케이션의 경우에는
//  프레임워크에서 이 작업을 자동으로 수행합니다.

void CImageClassificationDlg::OnPaint()
{
	if (IsIconic())
	{
		CPaintDC dc(this); // 그리기를 위한 디바이스 컨텍스트입니다.

		SendMessage(WM_ICONERASEBKGND, reinterpret_cast<WPARAM>(dc.GetSafeHdc()), 0);

		// 클라이언트 사각형에서 아이콘을 가운데에 맞춥니다.
		int cxIcon = GetSystemMetrics(SM_CXICON);
		int cyIcon = GetSystemMetrics(SM_CYICON);
		CRect rect;
		GetClientRect(&rect);
		int x = (rect.Width() - cxIcon + 1) / 2;
		int y = (rect.Height() - cyIcon + 1) / 2;

		// 아이콘을 그립니다.
		dc.DrawIcon(x, y, m_hIcon);
	}
	else
	{
		CDialogEx::OnPaint();
		
		
	}
}

// 사용자가 최소화된 창을 끄는 동안에 커서가 표시되도록 시스템에서
//  이 함수를 호출합니다.
HCURSOR CImageClassificationDlg::OnQueryDragIcon()
{
	return static_cast<HCURSOR>(m_hIcon);
}


void CImageClassificationDlg::OnBnClickedButtonImageload()
{
	static TCHAR BASED_CODE szFilter[] = _T("이미지 파일(*.BMP, *.GIF, *.JPG) | *.BMP;*.GIF;*.JPG;*.bmp;*.jpg;*.gif |모든파일(*.*)|*.*||");

	// FileDialog 생성
	CFileDialog dlg(TRUE, _T("*.jpg"), _T("image"), OFN_HIDEREADONLY, szFilter);

	if (IDOK == dlg.DoModal())
	{
		CString pathName = dlg.GetPathName(); // 선택된 파일 경로 확인
		MessageBox(pathName);                 // 로드 경로 확인
		this->m_image_path = pathName;        // 이미지 경로 저장 
	}

	// 제대로 저장 되었는지 확인하기
	std::wcout << (const wchar_t*)this->m_image_path;

	// 픽처박스 dc pointer 불러 오기 
	CDC* pdc = m_picture_box.GetWindowDC();
	// 이미지 로드 
	this->m_image.Load(this->m_image_path);
	m_picture_box.GetWindowRect(m_image_rect);
	m_image.StretchBlt(pdc->m_hDC, 0, 0, m_image_rect.Width(), m_image_rect.Height(), SRCCOPY);
	ReleaseDC(pdc);

	// Image Classification 메세지 Post 
	PostMessage(WM_USER_MESSAGE1, 100, 200);
}

bool loadimage(CString file_name, cv::Mat& image)
{
	// CString to std::string for cv::imread 

	std::string tr_file_name;
	tr_file_name = std::string(CT2CA(file_name.operator LPCWSTR()));

	image = cv::imread(tr_file_name);

	cv::cvtColor(image, image, cv::COLOR_BGR2RGB);

	// scale image to fit
	cv::Size scale(IMAGE_SIZE, IMAGE_SIZE);
	cv::resize(image, image, scale);
	// convert [unsigned int] to [float]
	image.convertTo(image, CV_32FC3, 1.0f / 255.0f);

	return true;
}


LRESULT CImageClassificationDlg::OnImageClassification(WPARAM wParam, LPARAM lParam)
{
	std::cout << "\nOnImageClassification() 핸들러 실행!";

	//이미지 로드 하기
	if (loadimage(this->m_image_path, this->m_inference_img))
	{
		//cv2 이미지 -> torch tensor 
		auto input_tensor = torch::from_blob(this->m_inference_img.data, { 1, IMAGE_SIZE, IMAGE_SIZE, CHANNELS });
		input_tensor = input_tensor.permute({ 0, 3, 1, 2 });
		
		// 모델 인퍼런스 
		torch::Tensor out_tensor = this->Model.forward({ input_tensor }).toTensor();

		// softmax and get argmax 구하기 
		auto results = out_tensor.sort(-1, true);
		auto softmaxs = std::get<0>(results)[0].softmax(0);
		auto indexs = std::get<1>(results)[0];
		auto idx = indexs[0].item<int>();

		// 결과 display 
		std::cout << "label:  " << labels[idx] << std::endl;
		std::cout << "probability:  " << softmaxs[0].item<float>() * 100.0f << "%" << std::endl;
		
		CString output_text = CString::CStringT(CA2CT(labels[idx].c_str()));
		float probability = softmaxs[0].item<float>() * 100.0f;

		m_edit_class->SetWindowTextW(output_text);
		m_edit_probability->SetWindowTextW(CString::CStringT(CA2CT(std::to_string(probability).c_str())));
	}
	else
	{
		std::cout << "이미지 로드에 실패하였습니다!" << std::endl;
	}




	return 0;
}


