init

app.py

@@ -3,18 +3,49 @@ import torch
 from PIL import Image
 import numpy as np
 import spaces  # Import spaces for ZeroGPU compatibility
-from einops import rearrange
-
 import math
-import torch
-import argparse
-from PIL import Image
+import re
 from einops import rearrange
 from mmengine.config import Config
 from xtuner.registry import BUILDER
 from xtuner.model.utils import guess_load_checkpoint
 
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
 from scripts.camera.cam_dataset import Cam_Generator
+from scripts.camera.visualization.visualize_batch import make_perspective_figures
+
+
+NUM = r"[+-]?(?:\d+(?:\.\d+)?|\.\d+)(?:[eE][+-]?\d+)?"
+CAM_PATTERN = re.compile(r"(?:camera parameters.*?:|roll.*?:)\s*("+NUM+r")\s*,\s*("+NUM+r")\s*,\s*("+NUM+r")", re.IGNORECASE|re.DOTALL)
+
+def center_crop(image):
+    w, h = image.size
+    s = min(w, h)
+    l = (w - s) // 2
+    t = (h - s) // 2
+    return image.crop((l, t, l + s, t + s))
+
+def preprocess_image(image_path, image_size):
+    image = Image.open(image_path).convert('RGB')
+    image = center_crop(image)
+    image = image.resize((image_size, image_size))
+    x = torch.from_numpy(np.array(image)).float()
+    x = x / 255.0
+    x = 2 * x - 1
+    x = rearrange(x, 'h w c -> c h w')
+    return x
+
+def load_square_rgb(image_path, image_size):
+    image = Image.open(image_path).convert('RGB')
+    image = center_crop(image)
+    image = image.resize((image_size, image_size))
+    arr = np.array(image)[:, :, ::-1]
+    return arr.copy()
+
 
 ##### load model
 config = "configs/pipelines/stage_2_base.py"
@@ -29,19 +60,6 @@ if torch.cuda.is_available():
 else:
     model = model.to(torch.float32)
 
-def expand2square(pil_img, background_color):
-    width, height = pil_img.size
-    if width == height:
-        return pil_img
-    elif width > height:
-        result = Image.new(pil_img.mode, (width, width), background_color)
-        result.paste(pil_img, (0, (width - height) // 2))
-        return result
-    else:
-        result = Image.new(pil_img.mode, (height, height), background_color)
-        result.paste(pil_img, ((height - width) // 2, 0))
-        return result
-
 
 @torch.inference_mode()
 @spaces.GPU(duration=120) 

requirements.txt

@@ -22,4 +22,6 @@ transformers-stream-generator==0.0.5
 triton==3.3.0
 xtuner==0.1.23
 yarl==1.20.0
+deepspeed
+git+https://github.com/KangLiao929/Puffin