Load jpeg from memory

dlib.go

@@ -2,5 +2,5 @@ package dlib
 
 // #cgo pkg-config: dlib-1
 // #cgo CFLAGS: -Wall -O3 -march=native
-// #cgo LDFLAGS: -lgif
+// #cgo LDFLAGS: -ljpeg
 import "C"

error.go

@@ -16,12 +16,6 @@ func (e SerializationError) Error() string {
 	return string(e)
 }
 
-type RecognizeError string
-
-func (e RecognizeError) Error() string {
-	return string(e)
-}
-
 type UnknownError string
 
 func (e UnknownError) Error() string {
@@ -35,8 +29,6 @@ func makeError(s string, code int) error {
 		return ImageLoadError(s)
 	case C.SERIALIZATION_ERROR:
 		return SerializationError(s)
-	case C.RECOGNIZE_ERROR:
-		return RecognizeError(s)
 	default:
 		return UnknownError(s)
 	}

facerec.cc

@@ -1,8 +1,8 @@
 #include <dlib/dnn.h>
-#include <dlib/image_io.h>
+#include <dlib/image_loader/image_loader.h>
 #include <dlib/image_processing/frontal_face_detector.h>
-
 #include "facerec.h"
+#include "jpeg_mem_loader.h"
 
 using namespace dlib;
 
@@ -36,6 +36,7 @@ using anet_type = loss_metric<fc_no_bias<128,avg_pool_everything<
 
 typedef matrix<float,0,1> descriptor;
 typedef std::pair<std::vector<rectangle>, std::vector<descriptor>> faces;
+
 static const size_t RECT_SIZE = 4 * sizeof(long);
 static const size_t DESCR_SIZE = 128 * sizeof(float);
 
@@ -53,10 +54,7 @@ public:
 	}
 
 	// TODO(Kagami): Jittering?
-	faces Recognize(const char* img_path, int max_faces) {
-		matrix<rgb_pixel> img;
-		load_image(img, img_path);
-
+	faces Recognize(matrix<rgb_pixel>& img, int max_faces) {
 		std::vector<rectangle> rects = detector_(img);
 		std::vector<descriptor> descrs;
 
@@ -98,11 +96,14 @@ facerec* facerec_init(const char* model_dir) {
 	return rec;
 }
 
-faceret* facerec_recognize(facerec* rec, const char* img_path, int max_faces) {
+faceret* facerec_recognize(facerec* rec, const uint8_t* img_data, int len, int max_faces) {
 	faceret* ret = (faceret*)calloc(1, sizeof(faceret));
 	FaceRec* cls = (FaceRec*)(rec->cls);
+	matrix<rgb_pixel> img;
 	try {
-		faces faces = cls->Recognize(img_path, max_faces);
+		// TODO(Kagami): Support more file types?
+		load_mem_jpeg(img, img_data, len);
+		faces faces = cls->Recognize(img, max_faces);
 		std::vector<rectangle> rects = faces.first;
 		std::vector<descriptor> descrs = faces.second;
 		ret->num_faces = descrs.size();

facerec.go

@@ -6,6 +6,8 @@ package dlib
 import "C"
 import (
 	"image"
+	"io/ioutil"
+	"os"
 	"unsafe"
 )
 
@@ -41,10 +43,15 @@ func NewFaceRec(modelDir string) (rec *FaceRec, err error) {
 	return
 }
 
-func (rec *FaceRec) recognize(imgPath string, maxFaces int) (faces []Face, err error) {
-	cImgPath := C.CString(imgPath)
-	defer C.free(unsafe.Pointer(cImgPath))
-	ret := C.facerec_recognize(rec.p, cImgPath, C.int(maxFaces))
+func (rec *FaceRec) recognize(imgData []byte, maxFaces int) (faces []Face, err error) {
+	if len(imgData) == 0 {
+		err = ImageLoadError("Empty image")
+		return
+	}
+	cImgData := (*C.uchar)(&imgData[0])
+	cLen := C.int(len(imgData))
+	cMaxFaces := C.int(maxFaces)
+	ret := C.facerec_recognize(rec.p, cImgData, cLen, cMaxFaces)
 	defer C.free(unsafe.Pointer(ret))
 
 	if ret.err_str != nil {
@@ -84,20 +91,43 @@ func (rec *FaceRec) recognize(imgPath string, maxFaces int) (faces []Face, err e
 	return
 }
 
-// Recognize all image faces.
-// Empty list is returned if there are no faces, error is returned if
-// there was some error while decoding/processing image.
-func (rec *FaceRec) Recognize(imgPath string) (faces []Face, err error) {
-	return rec.recognize(imgPath, 0)
-}
-
-// Recognize if image has single face or return nil otherwise.
-func (rec *FaceRec) RecognizeSingle(imgPath string) (face *Face, err error) {
-	faces, err := rec.recognize(imgPath, 1)
+// Convenient method.
+func (rec *FaceRec) recognizeFile(imgPath string, maxFaces int) (face []Face, err error) {
+	fd, err := os.Open(imgPath)
 	if err != nil {
 		return
 	}
-	if len(faces) != 1 {
+	imgData, err := ioutil.ReadAll(fd)
+	if err != nil {
+		return
+	}
+	return rec.recognize(imgData, maxFaces)
+}
+
+// Recognize all image faces.
+// Empty list is returned if there are no faces, error is returned if
+// there was some error while decoding/processing image.
+func (rec *FaceRec) Recognize(imgData []byte) (faces []Face, err error) {
+	return rec.recognize(imgData, 0)
+}
+
+// Recognize if image has single face or return nil otherwise.
+func (rec *FaceRec) RecognizeSingle(imgData []byte) (face *Face, err error) {
+	faces, err := rec.recognize(imgData, 1)
+	if err != nil || len(faces) != 1 {
+		return
+	}
+	face = &faces[0]
+	return
+}
+
+func (rec *FaceRec) RecognizeFile(imgPath string) (faces []Face, err error) {
+	return rec.recognizeFile(imgPath, 0)
+}
+
+func (rec *FaceRec) RecognizeSingleFile(imgPath string) (face *Face, err error) {
+	faces, err := rec.recognizeFile(imgPath, 1)
+	if err != nil || len(faces) != 1 {
 		return
 	}
 	face = &faces[0]

facerec.h

@@ -7,7 +7,6 @@ extern "C" {
 typedef enum {
 	IMAGE_LOAD_ERROR,
 	SERIALIZATION_ERROR,
-	RECOGNIZE_ERROR,
 	UNKNOWN_ERROR,
 } err_code;
 
@@ -26,7 +25,7 @@ typedef struct faceret {
 } faceret;
 
 facerec* facerec_init(const char* model_dir);
-faceret* facerec_recognize(facerec* rec, const char* img_path, int max_faces);
+faceret* facerec_recognize(facerec* rec, const uint8_t* img_data, int len, int max_faces);
 void facerec_free(facerec* rec);
 
 #ifdef __cplusplus

facerec_test.go

@@ -10,7 +10,7 @@ func TestNumFaces(t *testing.T) {
 		t.Fatalf("Can't init face recognizer: %v", err)
 	}
 	defer rec.Close()
-	faces, err := rec.Recognize("testdata/pristin.jpg")
+	faces, err := rec.RecognizeFile("testdata/pristin.jpg")
 	if err != nil {
 		t.Fatalf("Can't get faces: %v", err)
 	}

jpeg_mem_loader.cc

@@ -0,0 +1,59 @@
+#include <stdio.h>
+#include <setjmp.h>
+#include <jpeglib.h>
+#include <dlib/image_loader/image_loader.h>
+#include "jpeg_mem_loader.h"
+
+// Based on dlib's JPEG loader and libjpeg example.
+struct jpeg_loader_error_mgr {
+	jpeg_error_mgr pub;  // "public" fields
+	jmp_buf setjmp_buffer;  // For return to caller
+};
+
+static void jpeg_loader_error_exit(j_common_ptr cinfo) {
+	// cinfo->err really points to a jpeg_loader_error_mgr struct, so coerce pointer.
+	jpeg_loader_error_mgr* myerr = (jpeg_loader_error_mgr*)cinfo->err;
+	// Return control to the setjmp point.
+	longjmp(myerr->setjmp_buffer, 1);
+}
+
+static void jpeg_loader_emit_message(j_common_ptr cinfo, int msg_level) {
+	// Just exit on any warning because file is probably corrupted.
+	if (msg_level < 0)
+		jpeg_loader_error_exit(cinfo);
+}
+
+// From memory JPEG loader because dlib lacks one.
+void load_mem_jpeg(dlib::matrix<dlib::rgb_pixel>& img, const uint8_t* img_data, int len) {
+	jpeg_decompress_struct cinfo;
+
+	jpeg_loader_error_mgr jerr;
+	cinfo.err = jpeg_std_error(&jerr.pub);
+	jerr.pub.error_exit = jpeg_loader_error_exit;
+	jerr.pub.emit_message = jpeg_loader_emit_message;
+	if (setjmp(jerr.setjmp_buffer)) {
+		char buffer[JMSG_LENGTH_MAX];
+		(jerr.pub.format_message)((j_common_ptr)&cinfo, buffer);
+		jpeg_destroy_decompress(&cinfo);
+		throw dlib::image_load_error(std::string("jpeg_mem_loader: decode error: ") + buffer);
+	}
+
+	jpeg_create_decompress(&cinfo);
+	jpeg_mem_src(&cinfo, img_data, len);
+	jpeg_read_header(&cinfo, TRUE);
+	jpeg_start_decompress(&cinfo);
+
+	if (cinfo.output_components != 3) {
+		jpeg_destroy_decompress(&cinfo);
+		throw dlib::image_load_error("jpeg_mem_loader: unsupported pixel size");
+	}
+
+	img.set_size(cinfo.output_height, cinfo.output_width);
+	while (cinfo.output_scanline < cinfo.output_height) {
+		uint8_t* buffer_array[1] = { (uint8_t*)&img(cinfo.output_scanline, 0) };
+		jpeg_read_scanlines(&cinfo, buffer_array, 1);
+	}
+
+	jpeg_finish_decompress(&cinfo);
+	jpeg_destroy_decompress(&cinfo);
+}

jpeg_mem_loader.h

@@ -0,0 +1,3 @@
+#pragma once
+
+void load_mem_jpeg(dlib::matrix<dlib::rgb_pixel>& img, const uint8_t* img_data, int len);