From c4d2d3208758f6827cc81a6345474e7660c446dd Mon Sep 17 00:00:00 2001
From: dongchangxi <458593490@qq.com>
Date: Wed, 12 Feb 2025 15:58:13 +0800
Subject: [PATCH] 1

---
 main_step1.py | 31 ++++++++++++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/main_step1.py b/main_step1.py
index 842053a..af522f9 100644
--- a/main_step1.py
+++ b/main_step1.py
@@ -1,6 +1,6 @@
 import os, sys, time, shlex, subprocess, shutil, requests, cv2, numpy as np
 from PIL import Image
-import platform,socket,atexit
+import platform,socket,atexit,redis
 if platform.system() == 'Windows':
     sys.path.append('e:\\libs\\')
     #sys.path.append('libs')
@@ -206,12 +206,41 @@ def main(pid, experience=False, makeloop=True,task_distributed_id="",isNoColorTe
                 if pid == '':
                     time.sleep(3)
                     continue
+                #判断影棚类型，如果是方形影棚，则推入到另外一个重建队列，使用新影棚的建模方式
+                ps_type = get_ps_type(pid)
+                if ps_type == 2:
+                    r = create_redis_connection()
+                    print(f'{time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())} {pid} 方形影棚数据推入到重建队列用新的建模方式')
+                    r.rpush('ai:ai_repeat_for_align_info',pid)
+                    time.sleep(3)
+                    continue
+
             else:
                 experience = False
             step1(pid, experience, makeloop,task_distributed_id,isNoColorTexture)
     else:
         step1(pid, experience, makeloop,task_distributed_id,isNoColorTexture)
 
+def create_redis_connection():
+    """创建 Redis 连接，若连接失败则重试"""
+    while True:
+        try:
+            r = redis.Redis(host="106.14.158.208", password="kcV2000", port=6379, db=6)
+            # 尝试进行一次操作，检查连接是否有效
+            r.ping()  # ping 操作是一个简单的连接测试
+            print("Redis连接成功！")
+            return r
+        except ConnectionError:
+            print("Redis连接失败，正在重试...")
+            time.sleep(5) 
+
+
+def get_ps_type(pid):
+    # return 1:圆形影棚 2:方形影棚
+    res = requests.get("https://mp.api.suwa3d.com/api/takephotoOrder/photoStudioInfo", params={'pid': pid})
+    return res.json()['data']['type']
+
+
 if __name__ == '__main__':    
     # 取云端redis任务，完成第一步的数据预处理后，将数据放入共享存储目录，将第二步任务塞入本地mysql队列
     # 默认循环值守，可传参数运行单一任务，以方便调试