de

app.py

@@ -9,6 +9,10 @@ import cv2
 import gradio as gr
 from torchvision import transforms 
 from controlnet_aux import OpenposeDetector
+import random
+import open3d as o3d
+from collections import Counter
+import trimesh
 
 ratios_map =  {
     0.5:{"width":704,"height":1408},
@@ -85,45 +89,170 @@ def resize_image_old(image):
 
 
 @spaces.GPU
-def generate_(prompt, negative_prompt, pose_image, input_image, num_steps, controlnet_conditioning_scale, seed):
-    generator = torch.Generator("cuda").manual_seed(seed)    
+def generate_(prompt, negative_prompt, pose_image, input_image, controlnet_conditioning_scale):
+    generator = torch.Generator()
+    generator.manual_seed(random.randint(0, 2147483647))
     images = pipe(
-    prompt, negative_prompt=negative_prompt, image=pose_image, num_inference_steps=num_steps, controlnet_conditioning_scale=float(controlnet_conditioning_scale),
+    prompt, negative_prompt=negative_prompt, image=pose_image, num_inference_steps=20, controlnet_conditioning_scale=float(controlnet_conditioning_scale),
     generator=generator, height=input_image.size[1], width=input_image.size[0],
     ).images 
     return images
 
 @spaces.GPU
-def process(input_image, prompt, negative_prompt, num_steps, controlnet_conditioning_scale, seed):
+def process(input_image, prompt, negative_prompt, controlnet_conditioning_scale):
     
     # resize input_image to 1024x1024
     input_image = resize_image(input_image)
     
     pose_image = openpose(input_image, include_body=True, include_hand=True, include_face=True)
   
-    images = generate_(prompt, negative_prompt, pose_image, input_image, num_steps, controlnet_conditioning_scale, seed)
+    images = generate_(prompt, negative_prompt, pose_image, input_image, controlnet_conditioning_scale)
 
     return [pose_image,images[0]]
+
+@spaces.GPU
+def predict_image(cond_image, prompt, negative_prompt, controlnet_conditioning_scale):
+  print("predict position map")
+  global pipe
+  generator = torch.Generator()
+  generator.manual_seed(random.randint(0, 2147483647))
+  image = pipe(
+      prompt,
+      negative_prompt=negative_prompt,
+      image = cond_image,
+      width=1024,
+      height=1024,
+      guidance_scale=8,
+      num_inference_steps=20,
+      generator=generator,
+      guess_mode = True,
+      controlnet_conditioning_scale = controlnet_conditioning_scale
+  ).images[0]
+  
+  return image
+
+def convert_pil_to_opencv(pil_image):
+  return np.array(pil_image)
+
+def inv_func(y,
+  c = -712.380100,
+  a = 137.375240,
+  b = 192.435866):
+  return (np.exp((y - c) / a) - np.exp(-c/a)) / 964.8468371292845
+
+def create_point_cloud(img1, img2):
+  if img1.shape != img2.shape:
+    raise ValueError("Both images must have the same dimensions.")
+
+  h, w, _ = img1.shape
+  points = []
+  colors = []
+  for y in range(h):
+    for x in range(w):
+      # ピクセル位置 (x, y) のRGBをXYZとして取得
+      r, g, b = img1[y, x]
+      r = inv_func(r) * 0.9
+      g = inv_func(g) / 1.7 * 0.6
+      b = inv_func(b)
+      r *= 150
+      g *= 150
+      b *= 150
+      points.append([g, b, r])  # X, Y, Z
+      # 対応するピクセル位置の画像2の色を取得
+      colors.append(img2[y, x] / 255.0)  # 色は0〜1にスケール
+
+  return np.array(points), np.array(colors)
+
+def point_cloud_to_glb(points, colors):
+  # Open3Dでポイントクラウドを作成
+  pc = o3d.geometry.PointCloud()
+  pc.points = o3d.utility.Vector3dVector(points)
+  pc.colors = o3d.utility.Vector3dVector(colors)
+  
+  # 一時的にPLY形式で保存
+  temp_ply_file = "temp_output.ply"
+  o3d.io.write_point_cloud(temp_ply_file, pc)
+  
+  # PLYをGLBに変換
+  mesh = trimesh.load(temp_ply_file)
+  glb_file = "output.glb"
+  mesh.export(glb_file)
+
+  return glb_file
+
+def visualize_3d(image1, image2):
+  print("Processing...")
+  # PIL画像をOpenCV形式に変換
+  img1 = convert_pil_to_opencv(image1)
+  img2 = convert_pil_to_opencv(image2)
+
+  # ポイントクラウド生成
+  points, colors = create_point_cloud(img1, img2)
+
+  # GLB形式に変換
+  glb_file = point_cloud_to_glb(points, colors)
+
+  return glb_file
+
+def scale_image(original_image):
+  aspect_ratio = original_image.width / original_image.height
+
+  if original_image.width > original_image.height:
+    new_width = 1024
+    new_height = round(new_width / aspect_ratio)
+  else:
+    new_height = 1024
+    new_width = round(new_height * aspect_ratio)
+
+  resized_original = original_image.resize((new_width, new_height), Image.LANCZOS)
+
+  return resized_original
+
+def get_edge_mode_color(img, edge_width=10):
+  # 外周の10ピクセル領域を取得
+  left = img.crop((0, 0, edge_width, img.height))  # 左端
+  right = img.crop((img.width - edge_width, 0, img.width, img.height))  # 右端
+  top = img.crop((0, 0, img.width, edge_width))  # 上端
+  bottom = img.crop((0, img.height - edge_width, img.width, img.height))  # 下端
+
+  # 各領域のピクセルデータを取得して結合
+  colors = list(left.getdata()) + list(right.getdata()) + list(top.getdata()) + list(bottom.getdata())
+
+  # 最頻値（mode）を計算
+  mode_color = Counter(colors).most_common(1)[0][0]  # 最も頻繁に出現する色を取得
+
+  return mode_color
+
+def paste_image(resized_img):
+  # 外周10pxの最頻値を背景色に設定
+  mode_color = get_edge_mode_color(resized_img, edge_width=10)
+  mode_background = Image.new("RGBA", (1024, 1024), mode_color)
+  mode_background = mode_background.convert('RGB')
+
+  x = (1024 - resized_img.width) // 2
+  y = (1024 - resized_img.height) // 2
+  mode_background.paste(resized_img, (x, y))
+
+  return mode_background
+
+def outpaint_image(image):
+  if type(image) == type(None):
+    return None
+  resized_img = scale_image(image)
+  image = paste_image(resized_img)
+  
+  return image
     
 block = gr.Blocks().queue()
 
 with block:
     gr.Markdown("## BRIA 2.3 ControlNet Pose")
-    gr.HTML('''
-      <p style="margin-bottom: 10px; font-size: 94%">
-        This is a demo for ControlNet Pose that using
-        <a href="https://huggingface.co/briaai/BRIA-2.3" target="_blank">BRIA 2.3 text-to-image model</a> as backbone. 
-        Trained on licensed data, BRIA 2.3 provide full legal liability coverage for copyright and privacy infringement.
-      </p>
-    ''')
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(sources=None, type="pil") # None for upload, ctrl+v and webcam
             prompt = gr.Textbox(label="Prompt")
             negative_prompt = gr.Textbox(label="Negative prompt", value="Logo,Watermark,Text,Ugly,Morbid,Extra fingers,Poorly drawn hands,Mutation,Blurry,Extra limbs,Gross proportions,Missing arms,Mutated hands,Long neck,Duplicate,Mutilated,Mutilated hands,Poorly drawn face,Deformed,Bad anatomy,Cloned face,Malformed limbs,Missing legs,Too many fingers")
-            num_steps = gr.Slider(label="Number of steps", minimum=25, maximum=100, value=50, step=1)
             controlnet_conditioning_scale = gr.Slider(label="ControlNet conditioning scale", minimum=0.1, maximum=2.0, value=1.0, step=0.05)
-            seed = gr.Slider(label="Seed", minimum=0, maximum=2147483647, step=1, randomize=True,)
             run_button = gr.Button(value="Run")
                         
         with gr.Column():
@@ -131,8 +260,7 @@ with block:
                 pose_image_output = gr.Image(label="Pose Image", type="pil", interactive=False)
                 generated_image_output = gr.Image(label="Generated Image", type="pil", interactive=False)
 
-    ips = [input_image, prompt, negative_prompt, num_steps, controlnet_conditioning_scale, seed]
-    run_button.click(fn=process, inputs=ips, outputs=[pose_image_output, generated_image_output])
+    run_button.click(fn=process, inputs=[input_image, prompt, negative_prompt, controlnet_conditioning_scale], outputs=[pose_image_output, generated_image_output])
 
 
 block.launch(debug = True)