1 files changed, 238 insertions, 0 deletions
diff --git a/shot.cpp b/shot.cpp
new file mode 100644
index 0000000..b0c1fdb
--- /dev/null
+++ b/shot.cpp
@@ -0,0 +1,238 @@
+#include <iostream>
+#include <iomanip>
+#include <cstdio>
+#include <cstdlib>
+#include <cassert>
+#include <set>
+#include <unistd.h>
+#include <tesseract/baseapi.h>
+#include <leptonica/allheaders.h>
+#include "lodepng.h"
+
+using namespace std;
+
+struct __attribute__((packed)) Clr {
+	uint8_t r, g, b;
+
+	uint8_t avg() const {
+		return ((int)r + g + b) / 3;
+	}
+
+	int bright() const {
+		return (int)r + g + b;
+	}
+
+	float blueness() const {
+		return ((int)b + 1) / (((int)r + g) / 2.0f + 1);
+	}
+};
+
+bool operator<(const Clr &a, const Clr &b) {
+	return a.bright() < b.bright();
+}
+
+ostream& operator<<(ostream &os, const Clr &c) {
+	return os << (unsigned)c.r << " " << (unsigned)c.g << " " << (unsigned)c.b;
+}
+
+void writeimg(const Clr *g, int gw, int gh, const string &fname) {
+	assert(lodepng_encode24_file(fname.data(), (uint8_t*)g, gw, gh) == 0);
+}
+
+tesseract::TessBaseAPI *tessapi = nullptr;
+
+char tesseract_recog(const Clr *g, int gw, int gh) {
+	if (!tessapi) {
+		tessapi = new tesseract::TessBaseAPI();
+		// Initialize tesseract-ocr with English, without specifying tessdata path
+		if (tessapi->Init(nullptr, "eng")) {
+			fprintf(stderr, "Could not initialize tesseract.\n");
+			exit(1);
+		}
+	}
+
+	const char *fname = ".tesseract_recog.png";
+	writeimg(g, gw, gh, fname);
+
+	// Open input image with leptonica library
+	Pix *image = pixRead(fname);
+	unlink(fname);
+	pixSetResolution(image, 70, 70);
+	tessapi->SetImage(image);
+	// Get OCR result
+	char *outText = tessapi->GetUTF8Text();
+	char c = outText[0];
+	delete[] outText;
+
+	// Destroy used object and release memory
+	// tessapi->End();
+	pixDestroy(&image);
+
+	return c;
+}
+
+void tesseract_finish() {
+	tessapi->End();
+	delete tessapi;
+}
+
+Clr *img;
+int w, h;
+
+template <bool (*F)(const Clr &c)>
+int flood(int x, int y, bool *mark) {
+	mark[w*y+x] = true;
+	int t = 1;
+	if (x > 0 && !mark[w*y+x-1] && F(img[w*y+x-1])) t += flood<F>(x-1, y, mark);
+	if (y > 0 && !mark[w*(y-1)+x] && F(img[w*(y-1)+x])) t += flood<F>(x, y-1, mark);
+	if (x < w-1 && !mark[w*y+x+1] && F(img[w*y+x+1])) t += flood<F>(x+1, y, mark);
+	if (y < h-1 && !mark[w*(y+1)+x] && F(img[w*(y+1)+x])) t += flood<F>(x, y+1, mark);
+	return t;
+}
+
+void cmd_histo() {
+	int cnt[256] = {};
+
+	for (int y = 0; y < h; y++) {
+		int avg = 0;
+		for (int x = 0; x < w; x++) avg += img[w*y+x].avg();
+		avg /= w;
+		cnt[avg]++;
+	}
+	for(int i = 0; i < 256; i++) {
+		cout << setw(3) << i;
+		for (int j = 0; j < cnt[i]; j++) cout << "*";
+		cout << endl;
+	}
+}
+
+void cmd_clrlist() {
+	set<Clr> s;
+
+	for (int i = 0; i < w * h; i++) {
+		s.insert(img[i]);
+	}
+
+	for (const Clr &c : s) {
+		cout << c << "  " << c.bright() << " " << c.blueness() << endl;
+	}
+}
+
+bool clr_interest(const Clr &c) {
+	return c.bright() <= 300 && c.blueness() <= 1.6;
+}
+
+Clr* subimage(const Clr *orig, int ow, int sx, int sy, int sw, int sh) {
+	Clr *sub = new Clr[sw * sh];
+	for (int y = 0; y < sh; y++) {
+		memcpy(sub + sw * y, orig + ow * (y + sy) + sx, 3 * sw);
+	}
+	return sub;
+}
+
+void cmd_recog() {
+	for (int i = 0; i < w * h; i++) {
+		if (!clr_interest(img[i])) {
+			img[i] = {255, 255, 255};
+		}
+	}
+
+	// writeimg(img, w, h, "out1.png");
+
+	bool *mark = new bool[w*h]();
+	int maxn = -1, maxi = -1;
+
+	for (int y = 0; y < h; y++) {
+		for (int x = 0; x < w; x++) {
+			if (!mark[w*y+x]) {
+				int n = flood<clr_interest>(x, y, mark);
+				if (n > maxn) {
+					maxn = n;
+					maxi = w*y + x;
+				}
+			}
+		}
+	}
+
+	memset(mark, 0, w*h*sizeof(bool));
+	flood<clr_interest>(maxi%w, maxi/w, mark);
+
+	int minx = w, maxx = 0, miny = h, maxy = 0;
+
+	for (int y = 0; y < h; y++) {
+		for (int x = 0; x < w; x++) {
+			if (mark[w*y+x]) {
+				img[w*y+x] = {255, 255, 255};
+				if (x < minx) minx = x;
+				if (y < miny) miny = y;
+				if (x > maxx) maxx = x;
+				if (y > maxy) maxy = y;
+			} else if (img[w*y+x].bright() != 3*255) {
+				img[w*y+x] = {0, 0, 0};
+			}
+		}
+	}
+
+	delete[] mark;
+
+	/* int w2 = maxx - minx + 1;
+	int h2 = maxy - miny + 1;
+
+	Clr *img2 = new Clr[w2*h2];
+	for (int y = 0; y < h2; y++) {
+		for (int x = 0; x < w2; x++) {
+			img2[w2*y+x] = img[w*(y+miny)+x+minx];
+		}
+	}
+
+	writeimg(img2, w2, h2, "out.png");
+
+	delete[] img2; */
+
+	int result[81];
+	for (int yi = 0; yi < 9; yi++) {
+		for (int xi = 0; xi < 9; xi++) {
+			int sx = minx + (maxx - minx) * xi / 9;
+			int sy = miny + (maxy - miny) * yi / 9;
+			int sw = (maxx - minx) / 9;
+			int sh = (maxy - miny) / 9;
+			Clr *sub = subimage(img, w, sx, sy, sw, sh);
+			// writeimg(sub, sw, sh, "sub" + to_string(yi) + "_" + to_string(xi) + ".png");
+			char c = tesseract_recog(sub, sw, sh);
+			if (c == '\0') result[9*yi+xi] = -1;
+			else result[9*yi+xi] = c - '0';
+		}
+	}
+
+	for (int yi = 0; yi < 9; yi++) {
+		if (yi > 0 && yi % 3 == 0) cout << endl;
+		for (int xi = 0; xi < 9; xi++) {
+			if (xi > 0 && xi % 3 == 0) cout << ' ';
+			if (xi > 0) cout << ' ';
+			cout << ".0123456789"[result[9*yi+xi] + 1];
+		}
+		cout << endl;
+	}
+
+	tesseract_finish();
+}
+
+int main(int argc, char **argv){
+	if (argc != 3) {
+		cout << "Usage: " << argv[0] << " <cmd> <file.png>" << endl;
+		return 1;
+	}
+
+	const char *cmd = argv[1];
+	const char *srcfile = argv[2];
+
+	uint8_t *imgdata = nullptr;
+
+	assert(lodepng_decode24_file(&imgdata, (unsigned*)&w, (unsigned*)&h, srcfile) == 0);
+	img = (Clr*)imgdata;
+
+	if (strcmp(cmd, "h") == 0) cmd_histo();
+	else if (strcmp(cmd, "l") == 0) cmd_clrlist();
+	else if (strcmp(cmd, "r") == 0) cmd_recog();
+	else cout << "Unknown command" << endl;
+}