image post resize and light refine

app.py

@@ -5,8 +5,6 @@ from vision_tower import DINOv2_MLP
 from transformers import AutoImageProcessor
 import torch
 import os
-import matplotlib.pyplot as plt
-import io
 from PIL import Image
 
 import torch.nn.functional as F
@@ -78,14 +76,6 @@ def get_3angle_infer_aug(origin_img, rm_bkg_img):
     angles[3]  = confidence
     return angles
 
-
-def figure_to_img(fig):
-    with io.BytesIO() as buf:
-        fig.savefig(buf, format='JPG', bbox_inches='tight')
-        buf.seek(0)
-        image = Image.open(buf).copy()
-    return image
-
 def infer_func(img, do_rm_bkg, do_infer_aug):
     origin_img = Image.fromarray(img)
     if do_infer_aug:

render/core.py

@@ -195,14 +195,18 @@ def dot_product(a: Vec3d, b: Vec3d):
 def cross_product(a: Vec3d, b: Vec3d):
     return Vec3d(*speedup.cross_product(*a.arr, *b.arr))
 
-BASE_LIGHT = 0.3
+BASE_LIGHT = 0.9
 def get_light_intensity(face) -> float:
-    light0 = Vec3d(-2, 4, -10)
-    
-    light1 = Vec3d(10, 4, -2)
+    # lights = [Vec3d(-2, 4, -10), Vec3d(10, 4, -2), Vec3d(8, 8, -8), Vec3d(0, 0, -8)]
+    lights = [Vec3d(-2, 4, -10)]
+    # lights = []
+
     v1, v2, v3 = face
     up = normalize(cross_product(v2 - v1, v3 - v1))
-    return dot_product(up, normalize(light0))*0.6 + dot_product(up, normalize(light1))*0.6 + BASE_LIGHT
+    intensity = BASE_LIGHT
+    for light in lights:
+        intensity += dot_product(up, normalize(light))*0.2
+    return intensity
 
 
 def look_at(eye: Vec3d, target: Vec3d, up: Vec3d = Vec3d(0, -1, 0)) -> Mat4d:

utils.py

@@ -2,10 +2,11 @@ import rembg
 import random
 import torch
 import numpy as np
-from PIL import Image
+from PIL import Image, ImageOps
 import PIL
 from typing import Any
 import matplotlib.pyplot as plt
+import io
 
 def resize_foreground(
     image: Image,
@@ -232,8 +233,16 @@ def matplotlib_2D_arrow(angles, rm_bkg_img):
     ax.set_xlim(-5, 5)
     ax.set_ylim(-5, 5)
 
+def figure_to_img(fig):
+    with io.BytesIO() as buf:
+        fig.savefig(buf, format='JPG', bbox_inches='tight')
+        buf.seek(0)
+        image = Image.open(buf).copy()
+    return image
+
 from render import render, Model
 import math
+axis_model = Model("./axis.obj", texture_filename="./axis.png")
 def render_3D_axis(phi, theta, gamma):
     radius = 240
     # camera_location = [radius * math.cos(phi), radius * math.sin(phi), radius * math.tan(theta)]
@@ -241,7 +250,7 @@ def render_3D_axis(phi, theta, gamma):
     camera_location = [-1*radius * math.cos(phi), -1*radius * math.tan(theta), radius * math.sin(phi)]
     img = render(
         # Model("res/jinx.obj", texture_filename="res/jinx.tga"),
-        Model("./axis.obj", texture_filename="./axis.png"),
+        axis_model,
         height=512,
         width=512,
         filename="tmp_render.png",
@@ -269,22 +278,27 @@ def overlay_images_with_scaling(center_image: Image.Image, background_image, tar
     
     # 缩放背景图像，确保其适合前景图像的尺寸
     bg_width, bg_height = background_image.size
-    target_width, target_height = target_size
     
     # 按宽度或高度等比例缩放背景
-    scale = max(target_width / bg_width, target_height / bg_height)
-    new_size = (int(bg_width * scale), int(bg_height * scale))
-    resized_background = background_image.resize(new_size)
-    
-    # 裁剪背景图像至目标大小
-    left = (new_size[0] - target_width) // 2
-    top = (new_size[1] - target_height) // 2
-    right = left + target_width
-    bottom = top + target_height
-    cropped_background = resized_background.crop((left, top, right, bottom))
+    scale = target_size[0] / max(bg_width, bg_height)
+    new_width = int(bg_width * scale)
+    new_height = int(bg_height * scale)
+    resized_background = background_image.resize((new_width, new_height))
+    # 计算需要的填充量
+    pad_width = target_size[0] - new_width
+    pad_height = target_size[0] - new_height
+
+    # 计算上下左右的 padding
+    left = pad_width // 2
+    right = pad_width - left
+    top = pad_height // 2
+    bottom = pad_height - top
+
+    # 添加 padding
+    resized_background = ImageOps.expand(resized_background, border=(left, top, right, bottom), fill=(255,255,255,255))
     
     # 将前景图像叠加到背景图像上
-    result = cropped_background.copy()
+    result = resized_background.copy()
     result.paste(center_image, (0, 0), mask=center_image)
     
     return result