Multires baker: fix wrong normalization if baking happens to multiple images

Previously normalization will happen per image buffer individually, which
was wrong in cases different faces of the sane mesh uses different images.

Also solved possible threading issues when calculating min.max displacement.

3 files changed, 90 insertions, 111 deletions

diff --git a/source/blender/render/extern/include/RE_shader_ext.h b/source/blender/render/extern/include/RE_shader_ext.h
index 8818aa67c32..d686de21517 100644
--- a/source/blender/render/extern/include/RE_shader_ext.h
+++ b/source/blender/render/extern/include/RE_shader_ext.h
@@ -186,7 +186,6 @@ typedef struct ShadeInput {
 
 typedef struct BakeImBufuserData {
 	float *displacement_buffer;
-	float displacement_min, displacement_max;
 	char *mask_buffer;
 } BakeImBufuserData;
 
@@ -212,7 +211,7 @@ struct Object;
 int RE_bake_shade_all_selected(struct Render *re, int type, struct Object *actob, short *do_update, float *progress);
 struct Image *RE_bake_shade_get_image(void);
 void RE_bake_ibuf_filter(struct ImBuf *ibuf, char *mask, const int filter);
-void RE_bake_ibuf_normalize_displacement(struct ImBuf *ibuf, float *displacement, char *mask, float global_displacement_min, float global_displacement_max);
+void RE_bake_ibuf_normalize_displacement(struct ImBuf *ibuf, float *displacement, char *mask, float displacement_min, float displacement_max);
 
 #define BAKE_RESULT_OK			0
 #define BAKE_RESULT_NO_OBJECTS		1
diff --git a/source/blender/render/intern/source/multires_bake.c b/source/blender/render/intern/source/multires_bake.c
index a86c78ce259..13c9e6e88bd 100644
--- a/source/blender/render/intern/source/multires_bake.c
+++ b/source/blender/render/intern/source/multires_bake.c
@@ -60,14 +60,17 @@
 #include "rayobject.h"
 #include "rendercore.h"
 
-typedef void (*MPassKnownData)(DerivedMesh *lores_dm, DerivedMesh *hires_dm, const void *bake_data,
-                               ImBuf *ibuf, const int face_index, const int lvl, const float st[2],
-                               float tangmat[3][3], const int x, const int y);
+typedef void (*MPassKnownData)(DerivedMesh *lores_dm, DerivedMesh *hires_dm, void *thread_data,
+                               void *bake_data, ImBuf *ibuf, const int face_index, const int lvl,
+                               const float st[2], float tangmat[3][3], const int x, const int y);
 
 typedef void * (*MInitBakeData)(MultiresBakeRender *bkr, Image *ima);
-typedef void   (*MApplyBakeData)(void *bake_data);
 typedef void   (*MFreeBakeData)(void *bake_data);
 
+typedef struct MultiresBakeResult {
+	float height_min, height_max;
+} MultiresBakeResult;
+
 typedef struct {
 	MVert *mvert;
 	MFace *mface;
@@ -79,6 +82,7 @@ typedef struct {
 	int i0, i1, i2;
 	DerivedMesh *lores_dm, *hires_dm;
 	int lvl;
+	void *thread_data;
 	void *bake_data;
 	ImBuf *ibuf;
 	MPassKnownData pass_data;
@@ -95,7 +99,6 @@ typedef struct {
 
 typedef struct {
 	float *heights;
-	float height_min, height_max;
 	Image *ima;
 	DerivedMesh *ssdm;
 	const int *orig_index_mf_to_mpoly;
@@ -161,9 +164,11 @@ static void multiresbake_get_normal(const MResolvePixelData *data, float norm[],
 
 static void init_bake_rast(MBakeRast *bake_rast, const ImBuf *ibuf, const MResolvePixelData *data, MFlushPixel flush_pixel)
 {
+	BakeImBufuserData *userdata = (BakeImBufuserData *) ibuf->userdata;
+
 	memset(bake_rast, 0, sizeof(MBakeRast));
 
-	bake_rast->texels = ibuf->userdata;
+	bake_rast->texels = userdata->mask_buffer;
 	bake_rast->w = ibuf->x;
 	bake_rast->h = ibuf->y;
 	bake_rast->data = data;
@@ -222,7 +227,7 @@ static void flush_pixel(const MResolvePixelData *data, const int x, const int y)
 		zero_m3(to_tang);
 	}
 
-	data->pass_data(data->lores_dm, data->hires_dm, data->bake_data,
+	data->pass_data(data->lores_dm, data->hires_dm, data->thread_data, data->bake_data,
 	                data->ibuf, data->face_index, data->lvl, st, to_tang, x, y);
 }
 
@@ -348,6 +353,9 @@ typedef struct MultiresBakeThread {
 	/* thread-specific data */
 	MBakeRast bake_rast;
 	MResolvePixelData data;
+
+	/* displacement-specific data */
+	float height_min, height_max;
 } MultiresBakeThread;
 
 static int multires_bake_queue_next_face(MultiresBakeQueue *queue)
@@ -450,7 +458,7 @@ static void init_ccgdm_arrays(DerivedMesh *dm)
 }
 
 static void do_multires_bake(MultiresBakeRender *bkr, Image *ima, int require_tangent, MPassKnownData passKnownData,
-                             MInitBakeData initBakeData, MApplyBakeData applyBakeData, MFreeBakeData freeBakeData)
+                             MInitBakeData initBakeData, MFreeBakeData freeBakeData, MultiresBakeResult *result)
 {
 	DerivedMesh *dm = bkr->lores_dm;
 	const int lvl = bkr->lvl;
@@ -514,9 +522,13 @@ static void do_multires_bake(MultiresBakeRender *bkr, Image *ima, int require_ta
 			handle->data.hires_dm = bkr->hires_dm;
 			handle->data.lvl = lvl;
 			handle->data.pass_data = passKnownData;
+			handle->data.thread_data = handle;
 			handle->data.bake_data = bake_data;
 			handle->data.ibuf = ibuf;
 
+			handle->height_min = FLT_MAX;
+			handle->height_max = -FLT_MAX;
+
 			init_bake_rast(&handle->bake_rast, ibuf, &handle->data, flush_pixel);
 
 			if (tot_thread > 1)
@@ -529,12 +541,18 @@ static void do_multires_bake(MultiresBakeRender *bkr, Image *ima, int require_ta
 		else
 			do_multires_bake_thread(&handles[0]);
 
+		/* construct bake result */
+		result->height_min = handles[0].height_min;
+		result->height_max = handles[0].height_max;
+
+		for (i = 1; i < tot_thread; i++) {
+			result->height_min = min_ff(result->height_min, handles[i].height_min);
+			result->height_max = max_ff(result->height_max, handles[i].height_max);
+		}
+
 		BLI_spin_end(&queue.spin);
 
 		/* finalize baking */
-		if (applyBakeData)
-			applyBakeData(bake_data);
-
 		if (freeBakeData)
 			freeBakeData(bake_data);
 
@@ -676,13 +694,15 @@ static void *init_heights_data(MultiresBakeRender *bkr, Image *ima)
 	MHeightBakeData *height_data;
 	ImBuf *ibuf = BKE_image_acquire_ibuf(ima, NULL, NULL);
 	DerivedMesh *lodm = bkr->lores_dm;
+	BakeImBufuserData *userdata = ibuf->userdata;
+
+	if (userdata->displacement_buffer == NULL)
+		userdata->displacement_buffer = MEM_callocN(sizeof(float) * ibuf->x * ibuf->y, "MultiresBake heights");
 
 	height_data = MEM_callocN(sizeof(MHeightBakeData), "MultiresBake heightData");
 
 	height_data->ima = ima;
-	height_data->heights = MEM_callocN(sizeof(float) * ibuf->x * ibuf->y, "MultiresBake heights");
-	height_data->height_max = -FLT_MAX;
-	height_data->height_min = FLT_MAX;
+	height_data->heights = userdata->displacement_buffer;
 
 	if (!bkr->use_lores_mesh) {
 		SubsurfModifierData smd = {{NULL}};
@@ -710,48 +730,6 @@ static void *init_heights_data(MultiresBakeRender *bkr, Image *ima)
 	return (void *)height_data;
 }
 
-static void apply_heights_data(void *bake_data)
-{
-	MHeightBakeData *height_data = (MHeightBakeData *)bake_data;
-	ImBuf *ibuf = BKE_image_acquire_ibuf(height_data->ima, NULL, NULL);
-	int x, y, i;
-	float height, *heights = height_data->heights;
-	float min = height_data->height_min, max = height_data->height_max;
-
-	for (x = 0; x < ibuf->x; x++) {
-		for (y = 0; y < ibuf->y; y++) {
-			i = ibuf->x * y + x;
-
-			if (((char *)ibuf->userdata)[i] != FILTER_MASK_USED)
-				continue;
-
-			if (ibuf->rect_float) {
-				float *rrgbf = ibuf->rect_float + i * 4;
-
-				if (max - min > 1e-5f) height = (heights[i] - min) / (max - min);
-				else height = 0;
-
-				rrgbf[0] = rrgbf[1] = rrgbf[2] = height;
-			}
-			else {
-				char *rrgb = (char *)ibuf->rect + i * 4;
-
-				if (max - min > 1e-5f) height = (heights[i] - min) / (max - min);
-				else height = 0;
-
-				rrgb[0] = rrgb[1] = rrgb[2] = FTOCHAR(height);
-			}
-		}
-	}
-
-	if (ibuf->rect_float)
-		ibuf->userflags |= IB_RECT_INVALID;
-
-	ibuf->userflags |= IB_DISPLAY_BUFFER_INVALID;
-
-	BKE_image_release_ibuf(height_data->ima, ibuf, NULL);
-}
-
 static void free_heights_data(void *bake_data)
 {
 	MHeightBakeData *height_data = (MHeightBakeData *)bake_data;
@@ -759,7 +737,6 @@ static void free_heights_data(void *bake_data)
 	if (height_data->ssdm)
 		height_data->ssdm->release(height_data->ssdm);
 
-	MEM_freeN(height_data->heights);
 	MEM_freeN(height_data);
 }
 
@@ -769,13 +746,14 @@ static void free_heights_data(void *bake_data)
  *   - find coord of point and normal with specified UV in lo-res mesh (or subdivided lo-res
  *     mesh to make texture smoother) let's call this point p0 and n.
  *   - height wound be dot(n, p1-p0) */
-static void apply_heights_callback(DerivedMesh *lores_dm, DerivedMesh *hires_dm, const void *bake_data,
+static void apply_heights_callback(DerivedMesh *lores_dm, DerivedMesh *hires_dm, void *thread_data_v, void *bake_data,
                                    ImBuf *ibuf, const int face_index, const int lvl, const float st[2],
                                    float UNUSED(tangmat[3][3]), const int x, const int y)
 {
 	MTFace *mtface = CustomData_get_layer(&lores_dm->faceData, CD_MTFACE);
 	MFace mface;
 	MHeightBakeData *height_data = (MHeightBakeData *)bake_data;
+	MultiresBakeThread *thread_data = (MultiresBakeThread *) thread_data_v;
 	float uv[2], *st0, *st1, *st2, *st3;
 	int pixel = ibuf->x * y + x;
 	float vec[3], p0[3], p1[3], n[3], len;
@@ -822,15 +800,18 @@ static void apply_heights_callback(DerivedMesh *lores_dm, DerivedMesh *hires_dm,
 	len = dot_v3v3(n, vec);
 
 	height_data->heights[pixel] = len;
-	if (len < height_data->height_min) height_data->height_min = len;
-	if (len > height_data->height_max) height_data->height_max = len;
+
+	thread_data->height_min = min_ff(thread_data->height_min, len);
+	thread_data->height_max = max_ff(thread_data->height_max, len);
 
 	if (ibuf->rect_float) {
 		float *rrgbf = ibuf->rect_float + pixel * 4;
+		rrgbf[0] = rrgbf[1] = rrgbf[2] = len;
 		rrgbf[3] = 1.0f;
 	}
 	else {
 		char *rrgb = (char *)ibuf->rect + pixel * 4;
+		rrgb[0] = rrgb[1] = rrgb[2] = FTOCHAR(len);
 		rrgb[3] = 255;
 	}
 }
@@ -862,9 +843,9 @@ static void free_normal_data(void *bake_data)
  *   - find coord and normal of point with specified UV in hi-res mesh
  *   - multiply it by tangmat
  *   - vector in color space would be norm(vec) /2 + (0.5, 0.5, 0.5) */
-static void apply_tangmat_callback(DerivedMesh *lores_dm, DerivedMesh *hires_dm, const void *bake_data,
-                                   ImBuf *ibuf, const int face_index, const int lvl, const float st[2],
-                                   float tangmat[3][3], const int x, const int y)
+static void apply_tangmat_callback(DerivedMesh *lores_dm, DerivedMesh *hires_dm, void *UNUSED(thread_data),
+                                   void *bake_data, ImBuf *ibuf, const int face_index, const int lvl,
+                                   const float st[2], float tangmat[3][3], const int x, const int y)
 {
 	MTFace *mtface = CustomData_get_layer(&lores_dm->faceData, CD_MTFACE);
 	MFace mface;
@@ -1099,9 +1080,9 @@ static int trace_ao_ray(MAOBakeData *ao_data, float ray_start[3], float ray_dire
 	return RE_rayobject_raycast(ao_data->raytree, &isect);
 }
 
-static void apply_ao_callback(DerivedMesh *lores_dm, DerivedMesh *hires_dm, const void *bake_data,
-                              ImBuf *ibuf, const int face_index, const int lvl, const float st[2],
-                              float UNUSED(tangmat[3][3]), const int x, const int y)
+static void apply_ao_callback(DerivedMesh *lores_dm, DerivedMesh *hires_dm, void *UNUSED(thread_data),
+                              void *bake_data, ImBuf *ibuf, const int face_index, const int lvl,
+                              const float st[2], float UNUSED(tangmat[3][3]), const int x, const int y)
 {
 	MAOBakeData *ao_data = (MAOBakeData *) bake_data;
 	MTFace *mtface = CustomData_get_layer(&lores_dm->faceData, CD_MTFACE);
@@ -1231,7 +1212,7 @@ static void count_images(MultiresBakeRender *bkr)
 		mtface[a].tpage->id.flag &= ~LIB_DOIT;
 }
 
-static void bake_images(MultiresBakeRender *bkr)
+static void bake_images(MultiresBakeRender *bkr, MultiresBakeResult *result)
 {
 	LinkData *link;
 
@@ -1240,18 +1221,19 @@ static void bake_images(MultiresBakeRender *bkr)
 		ImBuf *ibuf = BKE_image_acquire_ibuf(ima, NULL, NULL);
 
 		if (ibuf->x > 0 && ibuf->y > 0) {
-			ibuf->userdata = MEM_callocN(ibuf->y * ibuf->x, "MultiresBake imbuf mask");
+			BakeImBufuserData *userdata = MEM_callocN(sizeof(BakeImBufuserData), "MultiresBake userdata");
+			userdata->mask_buffer = MEM_callocN(ibuf->y * ibuf->x, "MultiresBake imbuf mask");
+			ibuf->userdata = userdata;
 
 			switch (bkr->mode) {
 				case RE_BAKE_NORMALS:
-					do_multires_bake(bkr, ima, TRUE, apply_tangmat_callback, init_normal_data, NULL, free_normal_data);
+					do_multires_bake(bkr, ima, TRUE, apply_tangmat_callback, init_normal_data, free_normal_data, result);
 					break;
 				case RE_BAKE_DISPLACEMENT:
-					do_multires_bake(bkr, ima, FALSE, apply_heights_callback, init_heights_data,
-					                 apply_heights_data, free_heights_data);
+					do_multires_bake(bkr, ima, FALSE, apply_heights_callback, init_heights_data, free_heights_data, result);
 					break;
 				case RE_BAKE_AO:
-					do_multires_bake(bkr, ima, FALSE, apply_ao_callback, init_ao_data, NULL, free_ao_data);
+					do_multires_bake(bkr, ima, FALSE, apply_ao_callback, init_ao_data, free_ao_data, result);
 					break;
 			}
 		}
@@ -1262,18 +1244,25 @@ static void bake_images(MultiresBakeRender *bkr)
 	}
 }
 
-static void finish_images(MultiresBakeRender *bkr)
+static void finish_images(MultiresBakeRender *bkr, MultiresBakeResult *result)
 {
 	LinkData *link;
+	int use_displacement_buffer = bkr->mode == RE_BAKE_DISPLACEMENT;
 
 	for (link = bkr->image.first; link; link = link->next) {
 		Image *ima = (Image *)link->data;
 		ImBuf *ibuf = BKE_image_acquire_ibuf(ima, NULL, NULL);
+		BakeImBufuserData *userdata = (BakeImBufuserData *) ibuf->userdata;
 
 		if (ibuf->x <= 0 || ibuf->y <= 0)
 			continue;
 
-		RE_bake_ibuf_filter(ibuf, (char *)ibuf->userdata, bkr->bake_filter);
+		RE_bake_ibuf_filter(ibuf, userdata->mask_buffer, bkr->bake_filter);
+
+		if (use_displacement_buffer) {
+			RE_bake_ibuf_normalize_displacement(ibuf, userdata->displacement_buffer, userdata->mask_buffer,
+			                                    result->height_min, result->height_max);
+		}
 
 		ibuf->userflags |= IB_BITMAPDIRTY | IB_DISPLAY_BUFFER_INVALID;
 
@@ -1286,7 +1275,11 @@ static void finish_images(MultiresBakeRender *bkr)
 		}
 
 		if (ibuf->userdata) {
-			MEM_freeN(ibuf->userdata);
+			if (userdata->displacement_buffer)
+				MEM_freeN(userdata->displacement_buffer);
+
+			MEM_freeN(userdata->mask_buffer);
+			MEM_freeN(userdata);
 			ibuf->userdata = NULL;
 		}
 
@@ -1296,7 +1289,9 @@ static void finish_images(MultiresBakeRender *bkr)
 
 void RE_multires_bake_images(MultiresBakeRender *bkr)
 {
+	MultiresBakeResult result;
+
 	count_images(bkr);
-	bake_images(bkr);
-	finish_images(bkr);
+	bake_images(bkr, &result);
+	finish_images(bkr, &result);
 }
diff --git a/source/blender/render/intern/source/rendercore.c b/source/blender/render/intern/source/rendercore.c
index e4e6c9f7594..acb4fe062cf 100644
--- a/source/blender/render/intern/source/rendercore.c
+++ b/source/blender/render/intern/source/rendercore.c
@@ -2020,7 +2020,7 @@ typedef struct BakeShade {
 	/* displacement buffer used for normalization with unknown maximal distance */
 	int use_displacement_buffer;
 	float *displacement_buffer;
-	float *displacement_min, *displacement_max;
+	float displacement_min, displacement_max;
 	
 	int use_mask;
 	char *rect_mask; /* bake pixel mask */
@@ -2270,8 +2270,8 @@ static void bake_displacement(void *handle, ShadeInput *UNUSED(shi), float dist,
 	if (bs->displacement_buffer) {
 		float *displacement = bs->displacement_buffer + (bs->rectx * y + x);
 		*displacement = disp;
-		*bs->displacement_min = min_ff(*bs->displacement_min, disp);
-		*bs->displacement_max = max_ff(*bs->displacement_max, disp);
+		bs->displacement_min = min_ff(bs->displacement_min, disp);
+		bs->displacement_max = max_ff(bs->displacement_max, disp);
 	}
 
 	if (bs->rect_float && !bs->vcol) {
@@ -2670,8 +2670,8 @@ static void shade_verts(BakeShade *bs)
 	bs->rect = NULL;
 	bs->rect_float = NULL;
 	bs->displacement_buffer = NULL;
-	bs->displacement_min = NULL;
-	bs->displacement_max = NULL;
+	bs->displacement_min = FLT_MAX;
+	bs->displacement_max = -FLT_MAX;
 
 	bs->quad = 0;
 
@@ -2735,8 +2735,6 @@ static void shade_tface(BakeShade *bs)
 	bs->quad= 0;
 	bs->rect_mask = NULL;
 	bs->displacement_buffer = NULL;
-	bs->displacement_min = NULL;
-	bs->displacement_max = NULL;
 
 	if (bs->use_mask || bs->use_displacement_buffer) {
 		BakeImBufuserData *userdata = bs->ibuf->userdata;
@@ -2749,11 +2747,8 @@ static void shade_tface(BakeShade *bs)
 			if (bs->use_mask)
 				userdata->mask_buffer = MEM_callocN(sizeof(char) * bs->rectx * bs->recty, "BakeMask");
 
-			if (bs->use_displacement_buffer) {
+			if (bs->use_displacement_buffer)
 				userdata->displacement_buffer = MEM_callocN(sizeof(float) * bs->rectx * bs->recty, "BakeDisp");
-				userdata->displacement_min = FLT_MAX;
-				userdata->displacement_max = -FLT_MAX;
-			}
 
 			bs->ibuf->userdata = userdata;
 
@@ -2762,8 +2757,6 @@ static void shade_tface(BakeShade *bs)
 
 		bs->rect_mask = userdata->mask_buffer;
 		bs->displacement_buffer = userdata->displacement_buffer;
-		bs->displacement_min = &userdata->displacement_min;
-		bs->displacement_max = &userdata->displacement_max;
 	}
 	
 	/* get pixel level vertex coordinates */
@@ -2837,14 +2830,14 @@ void RE_bake_ibuf_filter(ImBuf *ibuf, char *mask, const int filter)
 	}
 }
 
-void RE_bake_ibuf_normalize_displacement(ImBuf *ibuf, float *displacement, char *mask, float global_displacement_min, float global_displacement_max)
+void RE_bake_ibuf_normalize_displacement(ImBuf *ibuf, float *displacement, char *mask, float displacement_min, float displacement_max)
 {
 	int i;
 	float *current_displacement = displacement;
 	char *current_mask = mask;
 	float max_distance;
 
-	max_distance = max_ff(fabsf(global_displacement_min), fabsf(global_displacement_max));
+	max_distance = max_ff(fabsf(displacement_min), fabsf(displacement_max));
 
 	for (i = 0; i < ibuf->x * ibuf->y; i++) {
 		if (*current_mask == FILTER_MASK_USED) {
@@ -2853,7 +2846,7 @@ void RE_bake_ibuf_normalize_displacement(ImBuf *ibuf, float *displacement, char
 			if (max_distance > 1e-5f)
 				normalized_displacement = (*current_displacement + max_distance) / (max_distance * 2);
 			else
-				normalized_displacement = 0.0f;
+				normalized_displacement = 0.5f;
 
 			if (ibuf->rect_float) {
 				/* currently baking happens to RGBA only */
@@ -2948,6 +2941,9 @@ int RE_bake_shade_all_selected(Render *re, int type, Object *actob, short *do_up
 
 		handles[a].do_update = do_update; /* use to tell the view to update */
 		
+		handles[a].displacement_min = FLT_MAX;
+		handles[a].displacement_max = -FLT_MAX;
+
 		BLI_insert_thread(&threads, &handles[a]);
 	}
 	
@@ -2970,23 +2966,12 @@ int RE_bake_shade_all_selected(Render *re, int type, Object *actob, short *do_up
 
 	/* filter and refresh images */
 	if ((R.r.bake_flag & R_BAKE_VCOL) == 0) {
-		float global_displacement_min = FLT_MAX, global_displacement_max = -FLT_MAX;
+		float displacement_min = FLT_MAX, displacement_max = -FLT_MAX;
 
 		if (use_displacement_buffer) {
-			for (ima = G.main->image.first; ima; ima = ima->id.next) {
-				if ((ima->id.flag & LIB_DOIT)==0) {
-					ImBuf *ibuf = BKE_image_acquire_ibuf(ima, NULL, NULL);
-					BakeImBufuserData *userdata;
-
-					if (!ibuf)
-						continue;
-
-					userdata = (BakeImBufuserData *) ibuf->userdata;
-					global_displacement_min = min_ff(global_displacement_min, userdata->displacement_min);
-					global_displacement_max = max_ff(global_displacement_max, userdata->displacement_max);
-
-					BKE_image_release_ibuf(ima, ibuf, NULL);
-				}
+			for (a = 0; a < re->r.threads; a++) {
+				displacement_min = min_ff(displacement_min, handles[a].displacement_min);
+				displacement_max = max_ff(displacement_max, handles[a].displacement_max);
 			}
 		}
 
@@ -3006,7 +2991,7 @@ int RE_bake_shade_all_selected(Render *re, int type, Object *actob, short *do_up
 
 				if (use_displacement_buffer) {
 					RE_bake_ibuf_normalize_displacement(ibuf, userdata->displacement_buffer, userdata->mask_buffer,
-					                                    global_displacement_min, global_displacement_max);
+					                                    displacement_min, displacement_max);
 				}
 
 				ibuf->userflags |= IB_BITMAPDIRTY;