[first]

__init__.py

@@ -0,0 +1,3 @@
+from custom.clip_ebc import ClipEBC
+
+__all__ = ["ClipEBC"]

app.py

@@ -0,0 +1,60 @@
+import gradio as gr
+from custom.clip_ebc_onnx import ClipEBCOnnx
+import numpy as np
+import matplotlib.pyplot as plt
+
+# ONNX 모델 초기화
+model = ClipEBCOnnx()
+
+def predict_crowd(image):
+    """
+    이미지를 받아서 군중 수를 예측하고 시각화 결과를 반환합니다.
+    
+    Args:
+        image: Gradio에서 받은 이미지 (numpy array)
+        
+    Returns:
+        tuple: (예측된 군중 수, 밀도 맵 시각화, 점 시각화)
+    """
+    count = model.predict(image)
+    
+    # 밀도 맵 시각화
+    fig_density, density_map = model.visualize_density_map()
+    plt.close(fig_density)  # 메모리 누수 방지
+    # 점 시각화
+    canvas, dot_map = model.visualize_dots()
+    plt.close(canvas.figure)  
+    
+    return (
+        f"예측된 군중 수: {count:.1f}명",
+        density_map,
+        dot_map
+    )
+
+with gr.Blocks(title="CLIP-EBC Crowd Counter") as app:
+    gr.Markdown("# CLIP-EBC Crowd Counter")
+    gr.Markdown("이미지를 업로드하여 군중 수를 예측하고 시각화합니다.")
+    
+    with gr.Row():
+        input_image = gr.Image(type="numpy", label="입력 이미지")
+    
+    with gr.Row():
+        predict_btn = gr.Button("예측", variant="primary")
+    
+    with gr.Row():
+        count_text = gr.Textbox(label="예측 결과")
+    
+    with gr.Row():
+        with gr.Column():
+            density_output = gr.Image(label="밀도 맵")
+        with gr.Column():
+            dots_output = gr.Image(label="점 시각화")
+    
+    predict_btn.click(
+        fn=predict_crowd,
+        inputs=input_image,
+        outputs=[count_text, density_output, dots_output]
+    )
+
+if __name__ == "__main__":
+    app.launch(share=False)

assets/__init__.py

@@ -0,0 +1,35 @@
+from huggingface_hub import hf_hub_download
+import os
+
+def download_required_files():
+    """Initialize required files from Hugging Face Hub"""
+    try:
+        cache_dir = "assets/"
+        if not os.path.exists(os.path.join(cache_dir, "CLIP_EBC_nwpu_rmse_onnx.onnx")):
+            hf_hub_download(
+                repo_id="PIA-SPACE-LAB/CLIP_EBC_nwpu_rmse_onnx",
+                filename="CLIP_EBC_nwpu_rmse_onnx.onnx",
+                # cache_dir=cache_dir,
+                local_dir=cache_dir
+            )
+        print("Required files downloaded successfully")
+    except Exception as e:
+        print(f"Error downloading required files: {e}")
+
+def download_required_files2():
+    """Initialize required files from Hugging Face Hub"""
+    try:
+        cache_dir = "assets/"
+        if not os.path.exists(os.path.join(cache_dir, "CLIP_EBC_nwpu_rmse.pth")):
+            hf_hub_download(
+                repo_id="PIA-SPACE-LAB/CLIP_EBC_nwpu_rmse",
+                filename="CLIP_EBC_nwpu_rmse.pth",
+                # cache_dir=cache_dir,
+                local_dir=cache_dir
+            )
+        print("Required files downloaded successfully")
+    except Exception as e:
+        print(f"Error downloading required files: {e}")
+
+download_required_files()
+download_required_files2()

configs/reduction_16.json

@@ -0,0 +1,33 @@
+{
+    "8":{
+        "qnrf": {
+            "bins": {
+                "fine":[
+                    [0, 0], [1, 1], [2, 2], [3, 3], [4, 4],
+                    [5, 5], [6, 6], [7, 7], [8, "inf"]
+                ],
+                "dynamic": [
+                    [0, 0], [1, 1], [2, 2], [3, 3],
+                    [4, 5], [6, 7], [8, "inf"]
+                ],
+                "coarse": [
+                    [0, 0], [1, 2], [3, 4], [5, 6], [7, "inf"]
+                ]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2, 3, 4, 5, 6, 7, 8],
+                    "average": [0, 1, 2, 3, 4, 5, 6, 7, 9.23349]
+                },
+                "dynamic": {
+                    "middle": [0, 1, 2, 3, 4.5, 6.5, 8],
+                    "average": [0, 1, 2, 3, 4.29278, 6.31441, 9.23349]
+                },
+                "coarse": {
+                    "middle": [0, 1.5, 3.5, 5.5, 7],
+                    "average": [0, 1.14978, 3.27641, 5.30609, 8.11466]
+                }
+            }
+        }
+    }
+}

configs/reduction_32.json

@@ -0,0 +1,56 @@
+{
+    "19": {
+        "qnrf": {
+            "bins": {
+                "fine": [
+                    [0, 0], [1, 1], [2, 2], [3, 3], [4, 4],
+                    [5, 5], [6, 6], [7, 7], [8, 8], [9, 9],
+                    [10, 10], [11, 11], [12, 12], [13, 13], [14, 14],
+                    [15, 15], [16, 16], [17, 17], [18, 18], [19, "inf"]
+                ],
+                "dynamic": [
+                    [0, 0], [1, 1], [2, 2], [3, 3], [4, 4],
+                    [5, 5], [6, 6], [7, 7], [8, 8], [9, 9],
+                    [10, 11], [12, 13], [14, 15], [16, 17], [18, "inf"]
+                ],
+                "coarse": [
+                    [0, 0], [1, 2], [3, 4], [5, 6], [7, 8],
+                    [9, 10], [11, 12], [13, 14], [15, 16], [17, 18],
+                    [19, "inf"]
+                ]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [
+                        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+                        11, 12, 13, 14, 15, 16, 17, 18, 19
+                    ],
+                    "average": [
+                        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+                        11, 12, 13, 14, 15, 16, 17, 18, 23.01897
+                    ]
+                },
+                "dynamic": {
+                    "middle": [
+                        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10.5,
+                        12.5, 14.5, 16.5, 18
+                    ],
+                    "average": [
+                        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10.42903,
+                        12.43320, 14.43341, 16.43521, 21.93548
+                    ]
+                },
+                "coarse": {
+                    "middle": [
+                        0, 1.5, 3.5, 5.5, 7.5, 9.5,
+                        11.5, 13.5, 15.5, 17.5, 19
+                    ],
+                    "average": [
+                        0, 1.23498, 3.36108, 5.40298, 7.41406, 9.42356,
+                        11.43094, 13.43244, 15.43697, 17.43759, 23.01897
+                    ]
+                }
+            }
+        }
+    }
+}

configs/reduction_8.json

@@ -0,0 +1,129 @@
+{   
+    "2": {
+        "sha": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2],
+                    "average": [0, 1, 2.24479]
+                }
+            }
+        },
+        "shb": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2],
+                    "average": [0, 1, 2.15171]
+                }
+            }
+        },
+        "nwpu": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2],
+                    "average": [0, 1, 2.10737]
+                }
+            }
+        },
+        "qnrf": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2],
+                    "average": [0, 1, 2.09296]
+                }
+            }
+        },
+        "jhu": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2],
+                    "average": [0, 1, 2.18589]
+                }
+            }
+        }
+    },
+    "4": {
+        "sha": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, 2], [3, 3], [4, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2, 3, 4],
+                    "average": [0, 1, 2, 3, 4.29992]
+                }
+            }
+        },
+        "shb": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, 2], [3, 3], [4, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2, 3, 4],
+                    "average": [0, 1, 2, 3, 4.41009]
+                }
+            }
+        },
+        "nwpu": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, 2], [3, 3], [4, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2, 3, 4],
+                    "average": [0, 1, 2, 3, 4.21931]
+                }
+            }
+        },
+        "qnrf": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, 2], [3, 3], [4, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2, 3, 4],
+                    "average": [0, 1, 2, 3, 4.21937]
+                }
+            }
+        },
+        "jhu": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, 2], [3, 3], [4, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2, 3, 4],
+                    "average": [0, 1, 2, 3, 4.24058]
+                }
+            }
+        }
+    },
+    "11": {
+        "qnrf": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7], [8, 8], [9, 9], [10, 10], [11, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
+                    "average": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
+                }
+            }
+        }
+    }
+}

custom/airport_color.py

@@ -0,0 +1,54 @@
+from PIL import Image, ImageDraw
+from custom.json2seg import get_segmentation_by_id
+import random
+INCHEON = "/home/jungseoik/data/PR/CLIP-EBC/assets/incheon.jpg"
+COLOR_PAIR = {1: '빨간색', 2: '주황색', 3: '노란색', 4: '초록색', 5: '빨간색', 6: '초록색'}
+
+def generate_random_color_pair():
+    colors = ['빨간색', '주황색', '노란색', '초록색']
+    return {i: random.choice(colors) for i in range(1, 7)}
+
+def create_mask(segmentation, img_size, color):
+    mask = Image.new('RGBA', img_size, (0, 0, 0, 0))
+    draw = ImageDraw.Draw(mask)
+
+    polygon = segmentation[0] 
+    points = [(polygon[i], polygon[i+1]) for i in range(0, len(polygon), 2)]
+    
+    color_map = {
+        '빨간색': (255, 0, 0, 128),    
+        '주황색': (255, 165, 0, 128),  
+        '노란색': (255, 255, 0, 128),  
+        '초록색': (0, 255, 0, 128),
+        '파란색': (0, 0, 255, 128),
+        '보라색': (128, 0, 128, 128)     
+    }
+    
+    draw.polygon(points, fill=color_map[color])
+    return mask
+
+def create_all_masks(img_size, region_color_pairs):
+    """
+    Parameters:
+    - img_size: 이미지 크기
+    - region_color_pairs: Dictionary 형태로 {region_id: color} 매핑
+        예: {1: '빨간색', 2: '초록색', 3: '노란색', ...}
+    """
+    # 최종 마스크 생성
+    final_mask = Image.new('RGBA', img_size, (0, 0, 0, 0))
+    
+    # 입력받은 region_color_pairs에 따라 마스크 생성 및 합성
+    for region_id, color in region_color_pairs.items():
+        segmentation = get_segmentation_by_id(target_id=region_id)
+        region_mask = create_mask(segmentation, img_size, color)
+        final_mask = Image.alpha_composite(final_mask, region_mask)
+    
+    return final_mask
+
+def airport_map_color(color_pairs = COLOR_PAIR):
+    # region_color_pairs = COLOR_PAIR
+    region_color_pairs = generate_random_color_pair()
+    image = Image.open(INCHEON)
+    all_masks = create_all_masks(image.size, region_color_pairs)
+    result = Image.alpha_composite(image.convert('RGBA'), all_masks)
+    return result

custom/clip_ebc.py

@@ -0,0 +1,346 @@
+import os
+import sys
+import torch
+from torchvision.transforms.functional import normalize, to_pil_image
+from torchvision.transforms import ToTensor, Normalize
+import matplotlib.pyplot as plt
+import json
+from models import get_model
+from utils import resize_density_map, sliding_window_predict
+from PIL import Image
+import numpy as np
+from scipy.ndimage import gaussian_filter
+from sklearn.cluster import KMeans
+import datetime
+from typing import Optional
+from typing import Union
+
+project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.append(project_root)
+os.environ['CUDA_LAUNCH_BLOCKING'] = "1"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+class ClipEBC:
+    """
+    CLIP-EBC (Efficient Boundary Counting) 이미지 처리 클래스입니다.
+    
+    CLIP 모델을 사용하여 이미지를 처리하며, 슬라이딩 윈도우 예측 기능을 포함한
+    다양한 설정 옵션을 제공합니다.
+    
+    Attributes:
+        truncation (int): 잘라내기 매개변수. 기본값 4.
+        reduction (int): 축소 비율. 기본값 8.
+        granularity (str): 세분화 수준. 기본값 "fine".
+        anchor_points (str): 앵커 포인트 방법. 기본값 "average".
+        model_name (str): CLIP 모델 이름. 기본값 "clip_vit_b_16".
+        input_size (int): 입력 이미지 크기. 기본값 224.
+        window_size (int): 슬라이딩 윈도우 크기. 기본값 224.
+        stride (int): 슬라이딩 윈도우 이동 간격. 기본값 224.
+        prompt_type (str): 프롬프트 유형. 기본값 "word".
+        dataset_name (str): 데이터셋 이름. 기본값 "qnrf".
+        num_vpt (int): 비주얼 프롬프트 토큰 수. 기본값 32.
+        vpt_drop (float): 비주얼 프롬프트 토큰 드롭아웃 비율. 기본값 0.0.
+        deep_vpt (bool): 깊은 비주얼 프롬프트 토큰 사용 여부. 기본값 True.
+        mean (tuple): 정규화를 위한 평균값. 기본값 (0.485, 0.456, 0.406).
+        std (tuple): 정규화를 위한 표준편차값. 기본값 (0.229, 0.224, 0.225).
+    """
+    
+    def __init__(self,
+                 truncation=4,
+                 reduction=8,
+                 granularity="fine",
+                 anchor_points="average",
+                 model_name="clip_vit_b_16",
+                 input_size=224,
+                 window_size=224,
+                 stride=224,
+                 prompt_type="word",
+                 dataset_name="qnrf",
+                 num_vpt=32,
+                 vpt_drop=0.,
+                 deep_vpt=True,
+                 mean=(0.485, 0.456, 0.406),
+                 std=(0.229, 0.224, 0.225),
+                 config_dir="configs"):
+        """CLIPEBC 클래스를 설정 매개변수와 함께 초기화합니다."""
+        self.truncation = truncation
+        self.reduction = reduction
+        self.granularity = granularity
+        self.anchor_points_type = anchor_points  # 원래 입력값 저장
+        self.model_name = model_name
+        self.input_size = input_size
+        self.window_size = window_size
+        self.stride = stride
+        self.prompt_type = prompt_type
+        self.dataset_name = dataset_name
+        self.num_vpt = num_vpt
+        self.vpt_drop = vpt_drop
+        self.deep_vpt = deep_vpt
+        self.mean = mean
+        self.std = std
+        self.config_dir = config_dir
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        self.bins = None
+        self.anchor_points = None
+        self.model = None
+        
+        # 초기 설정 로드 및 모델 초기화
+        self._load_config()
+        self._initialize_model()
+        
+    def _load_config(self):
+        """설정 파일을 로드하고 bins와 anchor_points를 설정합니다."""
+        config_path = os.path.join(self.config_dir, f"reduction_{self.reduction}.json")
+        with open(config_path, "r") as f:
+            config = json.load(f)[str(self.truncation)][self.dataset_name]
+        
+        self.bins = config["bins"][self.granularity]
+        self.bins = [(float(b[0]), float(b[1])) for b in self.bins]
+        
+        if self.anchor_points_type == "average":
+            self.anchor_points = config["anchor_points"][self.granularity]["average"]
+        else:
+            self.anchor_points = config["anchor_points"][self.granularity]["middle"]
+        self.anchor_points = [float(p) for p in self.anchor_points]
+        
+    def _initialize_model(self):
+        """CLIP 모델을 초기화합니다."""
+        self.model = get_model(
+            backbone=self.model_name,
+            input_size=self.input_size,
+            reduction=self.reduction,
+            bins=self.bins,
+            anchor_points=self.anchor_points,
+            prompt_type=self.prompt_type,
+            num_vpt=self.num_vpt,
+            vpt_drop=self.vpt_drop,
+            deep_vpt=self.deep_vpt
+        )
+
+        ckpt_path = "assets/CLIP_EBC_nwpu_rmse.pth"
+        ckpt = torch.load(ckpt_path, map_location=device)
+        self.model.load_state_dict(ckpt)
+        self.model = self.model.to(device)
+        self.model.eval()
+
+    def visualize_density_map(self, alpha: float = 0.5, save: bool = False, 
+                            save_path: Optional[str] = None):
+        """
+        현재 저장된 예측 결과를 시각화합니다.
+        
+        Args:
+            alpha (float): density map의 투명도 (0~1). 기본값 0.5
+            save (bool): 시각화 결과를 이미지로 저장할지 여부. 기본값 False
+            save_path (str, optional): 저장할 경로. None일 경우 현재 디렉토리에 자동 생성된 이름으로 저장.
+                기본값 None
+                
+        Returns:
+            Tuple[matplotlib.figure.Figure, np.ndarray]:
+                - density map이 오버레이된 matplotlib Figure 객체
+                - RGB 형식의 시각화된 이미지 배열 (H, W, 3)
+        Raises:
+            ValueError: density_map 또는 processed_image가 None인 경우 (predict 메서드가 실행되지 않은 경우)
+        """
+        if self.density_map is None or self.processed_image is None:
+            raise ValueError("먼저 predict 메서드를 실행하여 예측을 수행해야 합니다.")
+        
+        fig, ax = plt.subplots(dpi=200, frameon=False)
+        ax.imshow(self.processed_image)
+        ax.imshow(self.density_map, cmap="jet", alpha=alpha)
+        ax.axis("off")
+        
+        if save:
+            if save_path is None:
+                timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+                save_path = f"crowd_density_{timestamp}.png"
+            
+            # 여백 제거하고 저장
+            plt.savefig(save_path, bbox_inches='tight', pad_inches=0, dpi=200)
+            print(f"Image saved to: {save_path}")
+        
+        fig.canvas.draw()
+        image_from_plot = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8)
+        image_from_plot = image_from_plot.reshape(fig.canvas.get_width_height()[::-1] + (4,))
+        image_from_plot = image_from_plot[:,:,:3]  # RGB로 변환
+        
+        return fig , image_from_plot
+    
+    def visualize_dots(self, dot_size: int = 20, sigma: float = 1, percentile: float = 97, 
+                    save: bool = False, save_path: Optional[str] = None):
+        """
+        예측된 군중 위치를 점으로 표시하여 시각화합니다.
+        
+        Args:
+            dot_size (int): 점의 크기. 기본값 20
+            sigma (float): Gaussian 필터의 sigma 값. 기본값 1
+            percentile (float): 임계값으로 사용할 백분위수 (0-100). 기본값 97
+            save (bool): 시각화 결과를 이미지로 저장할지 여부. 기본값 False
+            save_path (str, optional): 저장할 경로. None일 경우 현재 디렉토리에 자동 생성된 이름으로 저장.
+                기본값 None
+                
+        Returns:
+            Tuple[matplotlib.backends.backend_agg.FigureCanvasBase, np.ndarray]: 
+                - matplotlib figure의 canvas 객체
+                - RGB 형식의 시각화된 이미지 배열 (H, W, 3)
+        Raises:
+            ValueError: density_map 또는 processed_image가 None인 경우 (predict 메서드가 실행되지 않은 경우)
+        """
+        if self.density_map is None or self.processed_image is None:
+            raise ValueError("먼저 predict 메서드를 실행하여 예측을 수행해야 합니다.")
+            
+        adjusted_pred_count = int(round(self.count))
+        
+        fig, ax = plt.subplots(dpi=200, frameon=False)
+        ax.imshow(self.processed_image)
+        
+        filtered_density = gaussian_filter(self.density_map, sigma=sigma)
+        
+        threshold = np.percentile(filtered_density, percentile)
+        candidate_pixels = np.column_stack(np.where(filtered_density >= threshold))
+        
+        if len(candidate_pixels) > adjusted_pred_count:
+            kmeans = KMeans(n_clusters=adjusted_pred_count, random_state=42, n_init=10)
+            kmeans.fit(candidate_pixels)
+            head_positions = kmeans.cluster_centers_.astype(int)
+        else:
+            head_positions = candidate_pixels
+            
+        y_coords, x_coords = head_positions[:, 0], head_positions[:, 1]
+        ax.scatter(x_coords, y_coords, 
+                    c='red',
+                    s=dot_size,
+                    alpha=1.0,
+                    edgecolors='white',
+                    linewidth=1)
+        
+        ax.axis("off")
+        
+        if save:
+            if save_path is None:
+                timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+                save_path = f"crowd_dots_{timestamp}.png"
+            
+            plt.savefig(save_path, bbox_inches='tight', pad_inches=0, dpi=200)
+            print(f"Image saved to: {save_path}")
+        
+        # Figure를 numpy 배열로 변환
+        fig.canvas.draw()
+        image_from_plot = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8)
+        image_from_plot = image_from_plot.reshape(fig.canvas.get_width_height()[::-1] + (4,))
+        image_from_plot = image_from_plot[:,:,:3]  # RGB로 변환
+        
+        # plt.close(fig)
+        # return image_from_plot
+        return fig.canvas, image_from_plot
+    
+    def _process_image(self, image: Union[str, np.ndarray]) -> torch.Tensor:
+        """
+        이미지를 전처리합니다. 이미지 경로, 넘파이 배열, Streamlit UploadedFile 모두 처리 가능합니다.
+        
+        Args:
+            image: 입력 이미지. 다음 형식 중 하나여야 합니다:
+                - str: 이미지 파일 경로
+                - np.ndarray: (H, W, 3) 형태의 RGB 이미지
+                - UploadedFile: Streamlit의 업로드된 파일
+                    
+        Returns:
+            torch.Tensor: 전처리된 이미지 텐서, shape (1, 3, H, W)
+            
+        Raises:
+            ValueError: 지원하지 않는 이미지 형식이 입력된 경우
+            Exception: 이미지 파일을 열 수 없는 경우
+        """
+        to_tensor = ToTensor()
+        normalize = Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        
+        # 원본 이미지 저장
+        self.original_image = image
+        
+        # 입력 타입에 따른 처리
+        if isinstance(image, str):
+            # 파일 경로인 경우
+            with open(image, "rb") as f:
+                pil_image = Image.open(f).convert("RGB")
+        elif isinstance(image, np.ndarray):
+            # 넘파이 배열인 경우
+            if image.dtype == np.uint8:
+                pil_image = Image.fromarray(image)
+            else:
+                # float 타입인 경우 [0, 1] 범위로 가정하고 변환
+                pil_image = Image.fromarray((image * 255).astype(np.uint8))
+        else:
+            # Streamlit UploadedFile 또는 기타 파일 객체인 경우
+            try:
+                pil_image = Image.open(image).convert("RGB")
+            except Exception as e:
+                raise ValueError(f"지원하지 않는 이미지 형식입니다: {type(image)}") from e
+        
+        # 텐서 변환 및 정규화
+        tensor_image = to_tensor(pil_image)
+        normalized_image = normalize(tensor_image)
+        batched_image = normalized_image.unsqueeze(0)  # (1, 3, H, W)
+        batched_image = batched_image.to(self.device)
+
+        return batched_image
+    def _post_process_image(self, image):
+        """이미지 후처리를 수행합니다."""
+        image = normalize(image, mean=(0., 0., 0.), 
+                        std=(1. / self.std[0], 1. / self.std[1], 1. / self.std[2]))
+        image = normalize(image, mean=(-self.mean[0], -self.mean[1], -self.mean[2]), 
+                        std=(1., 1., 1.))
+        processed_image = to_pil_image(image.squeeze(0))
+        return processed_image
+
+    @torch.no_grad()
+    def predict(self, image: torch.Tensor) -> Image.Image:
+        """
+        모델 출력 이미지의 후처리를 수행합니다.
+        
+        Args:
+            image (torch.Tensor): 후처리할 이미지 텐서, shape (1, 3, H, W)
+            
+        Returns:
+            PIL.Image.Image: 후처리된 PIL 이미지
+            
+        Note:
+            이미지 텐서에 대해 정규화를 역변환하고 PIL 이미지 형식으로 변환합니다.
+            self.mean과 self.std 값을 사용하여 원본 이미지의 스케일로 복원합니다.
+        """
+        processed_image = self._process_image(image)
+        image_height, image_width = processed_image.shape[-2:]
+        processed_image = processed_image.to(self.device)
+        
+        pred_density = sliding_window_predict(self.model, processed_image, 
+                                        self.window_size, self.stride)
+        pred_count = pred_density.sum().item()
+        resized_pred_density = resize_density_map(pred_density, 
+                                                (image_height, image_width)).cpu()
+        
+        self.processed_image = self._post_process_image(processed_image)
+        self.density_map = resized_pred_density.squeeze().numpy()
+        self.count = pred_count
+        
+        return pred_count
+    
+    def crowd_count(self):
+        """
+        가장 최근 예측의 군중 수를 반환합니다.
+        
+        Returns:
+            float: 예측된 군중 수
+            None: 아직 예측이 수행되지 않은 경우
+        """
+        return self.count
+    
+    def get_density_map(self):
+        """
+        가장 최근 예측의 밀도 맵을 반환합니다.
+        
+        Returns:
+            numpy.ndarray: 밀도 맵
+            None: 아직 예측이 수행되지 않은 경우
+        """
+        return self.density_map
+    

custom/clip_ebc_onnx.py

@@ -0,0 +1,465 @@
+import os
+import sys
+import torch
+import numpy as np
+import onnxruntime as ort
+from typing import Union, Tuple, Optional
+from PIL import Image
+import matplotlib.pyplot as plt
+from torchvision.transforms import ToTensor, Normalize
+from torchvision.transforms.functional import normalize, to_pil_image
+import json
+import datetime
+from scipy.ndimage import gaussian_filter
+from sklearn.cluster import KMeans
+import assets
+
+# 프로젝트 루트 디렉토리 설정
+project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.append(project_root)
+
+class ClipEBCOnnx:
+    """
+    CLIP-EBC (Efficient Boundary Counting) ONNX 버전 이미지 처리 클래스입니다.
+    
+    ONNX로 변환된 CLIP 모델을 사용하여 이미지를 처리하며, 슬라이딩 윈도우 예측 기능을 포함한
+    다양한 설정 옵션을 제공합니다.
+    """
+    
+    def __init__(self,
+                 onnx_model_path="assets/CLIP_EBC_nwpu_rmse_onnx.onnx",
+                 truncation=4,
+                 reduction=8,
+                 granularity="fine",
+                 anchor_points="average",
+                 input_size=224,
+                 window_size=224,
+                 stride=224,
+                 dataset_name="qnrf",
+                 mean=(0.485, 0.456, 0.406),
+                 std=(0.229, 0.224, 0.225),
+                 config_dir="configs"):
+        """CLIPEBC ONNX 클래스를 설정 매개변수와 함께 초기화합니다."""
+        self.onnx_model_path = onnx_model_path
+        self.truncation = truncation
+        self.reduction = reduction
+        self.granularity = granularity
+        self.anchor_points_type = anchor_points
+        self.input_size = input_size
+        self.window_size = window_size
+        self.stride = stride
+        self.dataset_name = dataset_name
+        self.mean = mean
+        self.std = std
+        self.config_dir = config_dir
+        
+        # 결과 저장용 변수 초기화
+        self.density_map = None
+        self.processed_image = None
+        self.count = None
+        self.original_image = None
+        
+        # ONNX 추론 세션 설정
+        self.session_options = ort.SessionOptions()
+        self.session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+        
+        # 가능한 경우 GPU 사용
+        self.providers = []
+        if 'CUDAExecutionProvider' in ort.get_available_providers():
+            self.providers.append('CUDAExecutionProvider')
+        self.providers.append('CPUExecutionProvider')
+        
+        # ONNX 런타임 세션 초기화
+        print(f"ONNX 모델 로드 중: {self.onnx_model_path}")
+        self.session = ort.InferenceSession(
+            self.onnx_model_path, 
+            sess_options=self.session_options,
+            providers=self.providers
+        )
+        
+        # 모델의 입력 및 출력 이름 가져오기
+        self.input_name = self.session.get_inputs()[0].name
+        self.output_name = self.session.get_outputs()[0].name
+        
+        print(f"입력 이름: {self.input_name}, 형태: {self.session.get_inputs()[0].shape}")
+        print(f"출력 이름: {self.output_name}, 형태: {self.session.get_outputs()[0].shape}")
+        print(f"실행 제공자: {self.providers}")
+        
+    def _process_image(self, image: Union[str, np.ndarray]) -> np.ndarray:
+        """
+        이미지를 전처리합니다. 이미지 경로, 넘파이 배열, Streamlit UploadedFile 모두 처리 가능합니다.
+        
+        Args:
+            image: 입력 이미지. 다음 형식 중 하나여야 합니다:
+                - str: 이미지 파일 경로
+                - np.ndarray: (H, W, 3) 형태의 RGB 이미지
+                - UploadedFile: Streamlit의 업로드된 파일
+                    
+        Returns:
+            np.ndarray: 전처리된 이미지 배열, shape (1, 3, H, W)
+        """
+        to_tensor = ToTensor()
+        normalize = Normalize(mean=self.mean, std=self.std)
+        
+        # 원본 이미지 저장
+        self.original_image = image
+        
+        # 입력 타입에 따른 처리
+        if isinstance(image, str):
+            # 파일 경로인 경우
+            with open(image, "rb") as f:
+                pil_image = Image.open(f).convert("RGB")
+        elif isinstance(image, np.ndarray):
+            # 넘파이 배열인 경우
+            if image.dtype == np.uint8:
+                pil_image = Image.fromarray(image)
+            else:
+                # float 타입인 경우 [0, 1] 범위로 가정하고 변환
+                pil_image = Image.fromarray((image * 255).astype(np.uint8))
+        else:
+            # Streamlit UploadedFile 또는 기타 파일 객체인 경우
+            try:
+                pil_image = Image.open(image).convert("RGB")
+            except Exception as e:
+                raise ValueError(f"지원하지 않는 이미지 형식입니다: {type(image)}") from e
+        
+        # 텐서 변환 및 정규화
+        tensor_image = to_tensor(pil_image)
+        normalized_image = normalize(tensor_image)
+        batched_image = normalized_image.unsqueeze(0)  # (1, 3, H, W)
+        
+        # numpy로 변환
+        numpy_image = batched_image.numpy()
+        
+        return numpy_image
+    
+    def _post_process_image(self, image_tensor):
+        """이미지 텐서를 PIL 이미지로 변환합니다."""
+        # NumPy 배열을 PyTorch 텐서로 변환
+        if isinstance(image_tensor, np.ndarray):
+            image_tensor = torch.from_numpy(image_tensor)
+            
+        # 정규화 역변환
+        image = normalize(
+            image_tensor,
+            mean=[0., 0., 0.],
+            std=[1./self.std[0], 1./self.std[1], 1./self.std[2]]
+        )
+        
+        image = normalize(
+            image,
+            mean=[-self.mean[0], -self.mean[1], -self.mean[2]],
+            std=[1., 1., 1.]
+        )
+        
+        # 배치 차원 제거 및 PIL 이미지로 변환
+        processed_image = to_pil_image(image.squeeze(0))
+        return processed_image
+
+    def sliding_window_predict(self, image: np.ndarray, window_size: Union[int, Tuple[int, int]], 
+                             stride: Union[int, Tuple[int, int]]) -> np.ndarray:
+        """
+        슬라이딩 윈도우 방식으로 이미지 예측을 수행합니다. 겹치는 영역은 평균값을 사용합니다.
+        
+        Args:
+            image (np.ndarray): 형태가 (1, 3, H, W)인 이미지 배열
+            window_size (int or tuple): 윈도우 크기
+            stride (int or tuple): 윈도우 이동 간격
+            
+        Returns:
+            np.ndarray: 예측된 밀도 맵
+        """
+        # 입력 검증
+        assert len(image.shape) == 4, f"이미지는 4차원 배열이어야 합니다. (1, C, H, W), 현재: {image.shape}"
+        
+        # 윈도우 크기와 스트라이드 설정
+        window_size = (int(window_size), int(window_size)) if isinstance(window_size, (int, float)) else window_size
+        stride = (int(stride), int(stride)) if isinstance(stride, (int, float)) else stride
+        window_size = tuple(window_size)
+        stride = tuple(stride)
+        
+        # 검증
+        assert isinstance(window_size, tuple) and len(window_size) == 2 and window_size[0] > 0 and window_size[1] > 0, \
+            f"윈도우 크기는 양수 정수 튜플 (h, w)이어야 합니다. 현재: {window_size}"
+        assert isinstance(stride, tuple) and len(stride) == 2 and stride[0] > 0 and stride[1] > 0, \
+            f"스트라이드는 양수 정수 튜플 (h, w)이어야 합니다. 현재: {stride}"
+        assert stride[0] <= window_size[0] and stride[1] <= window_size[1], \
+            f"스트라이드는 윈도우 크기보다 작아야 합니다. 현재: {stride}와 {window_size}"
+        
+        image_height, image_width = image.shape[-2:]
+        window_height, window_width = window_size
+        stride_height, stride_width = stride
+        
+        # 슬라이딩 윈도우 수 계산
+        num_rows = int(np.ceil((image_height - window_height) / stride_height) + 1)
+        num_cols = int(np.ceil((image_width - window_width) / stride_width) + 1)
+        
+        # 윈도우 추출
+        windows = []
+        window_positions = []
+        for i in range(num_rows):
+            for j in range(num_cols):
+                x_start, y_start = i * stride_height, j * stride_width
+                x_end, y_end = x_start + window_height, y_start + window_width
+                
+                # 이미지 경계 처리
+                if x_end > image_height:
+                    x_start, x_end = image_height - window_height, image_height
+                if y_end > image_width:
+                    y_start, y_end = image_width - window_width, image_width
+                
+                window = image[:, :, x_start:x_end, y_start:y_end]
+                windows.append(window)
+                window_positions.append((x_start, y_start, x_end, y_end))
+        
+        # 배치 단위로 추론
+        all_preds = []
+        max_batch_size = 8
+        
+        for start_idx in range(0, len(windows), max_batch_size):
+            end_idx = min(start_idx + max_batch_size, len(windows))
+            batch_windows = np.vstack(windows[start_idx:end_idx])  # (batch_size, 3, h, w)
+            
+            # ONNX 추론
+            ort_inputs = {self.input_name: batch_windows}
+            batch_preds = self.session.run([self.output_name], ort_inputs)[0]
+            
+            # Debug 정보
+            # print(f"배치 입력 형태: {batch_windows.shape}, 배치 출력 형태: {batch_preds.shape}")
+            
+            all_preds.extend([batch_preds[i:i+1] for i in range(batch_preds.shape[0])])
+        
+        # 예측 결과를 numpy 배열로 변환
+        preds = np.concatenate(all_preds, axis=0)
+        
+        # 출력 밀도 맵 조립
+        pred_map = np.zeros((preds.shape[1], image_height // self.reduction, image_width // self.reduction), dtype=np.float32)
+        count_map = np.zeros((preds.shape[1], image_height // self.reduction, image_width // self.reduction), dtype=np.float32)
+        
+        idx = 0
+        for i in range(num_rows):
+            for j in range(num_cols):
+                x_start, y_start, x_end, y_end = window_positions[idx]
+                
+                # 출력 영역 계산 (reduction 고려)
+                x_start_out = x_start // self.reduction
+                y_start_out = y_start // self.reduction
+                x_end_out = x_end // self.reduction
+                y_end_out = y_end // self.reduction
+                
+                pred_map[:, x_start_out:x_end_out, y_start_out:y_end_out] += preds[idx]
+                count_map[:, x_start_out:x_end_out, y_start_out:y_end_out] += 1.
+                idx += 1
+        
+        # 겹치는 영역 평균 계산
+        pred_map /= count_map
+        
+        return pred_map
+
+    def resize_density_map(self, density_map: np.ndarray, target_size: Tuple[int, int]) -> np.ndarray:
+        """
+        밀도 맵의 크기를 조정합니다. 총합은 보존됩니다.
+        
+        Args:
+            density_map: 형태가 (C, H, W)인 밀도 맵
+            target_size: 목표 크기 (H', W')
+            
+        Returns:
+            np.ndarray: 크기가 조정된 밀도 맵
+        """
+        from PIL import Image
+        import torch.nn.functional as F
+        import torch
+        
+        # numpy를 torch로 변환
+        if isinstance(density_map, np.ndarray):
+            density_map = torch.from_numpy(density_map)
+        
+        # 배치 차원 추가
+        if density_map.dim() == 3:
+            density_map = density_map.unsqueeze(0)  # (1, C, H, W)
+        
+        current_size = density_map.shape[2:]
+        
+        if current_size[0] == target_size[0] and current_size[1] == target_size[1]:
+            return density_map.squeeze(0).numpy()
+        
+        # 원본 밀도 맵의 총합 계산
+        original_sum = density_map.sum()
+        
+        # 크기 조정 (쌍선형 보간)
+        resized_map = F.interpolate(
+            density_map,
+            size=target_size,
+            mode='bilinear',
+            align_corners=False
+        )
+        
+        # 총합 보존을 위한 스케일링
+        if resized_map.sum() > 0:  # 0으로 나누기 방지
+            resized_map = resized_map * (original_sum / resized_map.sum())
+        
+        return resized_map.squeeze(0).numpy()
+
+    def predict(self, image: Union[str, np.ndarray]) -> float:
+        """
+        이미지에서 군중 계수 예측을 수행합니다.
+        
+        Args:
+            image: 입력 이미지 (경로, 넘파이 배열, 또는 업로드된 파일)
+            
+        Returns:
+            float: 예측된 사람 수
+        """
+        # 이미지 전처리
+        processed_image = self._process_image(image)
+        image_height, image_width = processed_image.shape[-2:]
+        
+        # 슬라이딩 윈도우 예측
+        pred_density = self.sliding_window_predict(
+            processed_image, 
+            self.window_size, 
+            self.stride
+        )
+        
+        # 예측 결과 저장
+        pred_count = pred_density.sum()
+        
+        # 원본 이미지 크기로 밀도 맵 조정
+        resized_pred_density = self.resize_density_map(
+            pred_density, 
+            (image_height, image_width)
+        )
+        
+        # 결과 저장
+        self.processed_image = self._post_process_image(processed_image)
+        self.density_map = resized_pred_density.squeeze()
+        self.count = pred_count
+        
+        return pred_count
+    
+    def visualize_density_map(self, alpha: float = 0.5, save: bool = False, 
+                            save_path: Optional[str] = None):
+        """
+        현재 저장된 예측 결과를 시각화합니다.
+        
+        Args:
+            alpha (float): density map의 투명도 (0~1). 기본값 0.5
+            save (bool): 시각화 결과를 이미지로 저장할지 여부. 기본값 False
+            save_path (str, optional): 저장할 경로. None일 경우 현재 디렉토리에 자동 생성된 이름으로 저장.
+                기본값 None
+                
+        Returns:
+            Tuple[matplotlib.figure.Figure, np.ndarray]:
+                - density map이 오버레이된 matplotlib Figure 객체
+                - RGB 형식의 시각화된 이미지 배열 (H, W, 3)
+        """
+        if self.density_map is None or self.processed_image is None:
+            raise ValueError("먼저 predict 메서드를 실행하여 예측을 수행해야 합니다.")
+        
+        fig, ax = plt.subplots(dpi=200, frameon=False)
+        ax.imshow(self.processed_image)
+        ax.imshow(self.density_map, cmap="jet", alpha=alpha)
+        ax.axis("off")
+        plt.title(f"Count: {self.count:.1f}")
+        
+        if save:
+            if save_path is None:
+                timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+                save_path = f"crowd_density_{timestamp}.png"
+            
+            # 여백 제거하고 저장
+            plt.savefig(save_path, bbox_inches='tight', pad_inches=0, dpi=200)
+            print(f"이미지 저장 완료: {save_path}")
+        
+        fig.canvas.draw()
+        image_from_plot = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8)
+        image_from_plot = image_from_plot.reshape(fig.canvas.get_width_height()[::-1] + (4,))
+        image_from_plot = image_from_plot[:,:,:3]  # RGB로 변환
+        
+        return fig, image_from_plot
+    
+    def visualize_dots(self, dot_size: int = 20, sigma: float = 1, percentile: float = 97, 
+                    save: bool = False, save_path: Optional[str] = None):
+        """
+        예측된 군중 위치를 점으로 표시하여 시각화합니다.
+        
+        Args:
+            dot_size (int): 점의 크기. 기본값 20
+            sigma (float): Gaussian 필터의 sigma 값. 기본값 1
+            percentile (float): 임계값으로 사용할 백분위수 (0-100). 기본값 97
+            save (bool): 시각화 결과를 이미지로 저장할지 여부. 기본값 False
+            save_path (str, optional): 저장할 경로. None일 경우 현재 디렉토리에 자동 생성된 이름으로 저장.
+                기본값 None
+                
+        Returns:
+            Tuple[matplotlib.backends.backend_agg.FigureCanvasBase, np.ndarray]: 
+                - matplotlib figure의 canvas 객체
+                - RGB 형식의 시각화된 이미지 배열 (H, W, 3)
+        """
+        if self.density_map is None or self.processed_image is None:
+            raise ValueError("먼저 predict 메서드를 실행하여 예측을 수행해야 합니다.")
+            
+        adjusted_pred_count = int(round(self.count))
+        
+        fig, ax = plt.subplots(dpi=200, frameon=False)
+        ax.imshow(self.processed_image)
+        
+        filtered_density = gaussian_filter(self.density_map, sigma=sigma)
+        
+        threshold = np.percentile(filtered_density, percentile)
+        candidate_pixels = np.column_stack(np.where(filtered_density >= threshold))
+        
+        if len(candidate_pixels) > adjusted_pred_count:
+            kmeans = KMeans(n_clusters=adjusted_pred_count, random_state=42, n_init=10)
+            kmeans.fit(candidate_pixels)
+            head_positions = kmeans.cluster_centers_.astype(int)
+        else:
+            head_positions = candidate_pixels
+            
+        y_coords, x_coords = head_positions[:, 0], head_positions[:, 1]
+        ax.scatter(x_coords, y_coords, 
+                    c='red',
+                    s=dot_size,
+                    alpha=1.0,
+                    edgecolors='white',
+                    linewidth=1)
+        
+        ax.axis("off")
+        plt.title(f"Count: {self.count:.1f}")
+        
+        if save:
+            if save_path is None:
+                timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+                save_path = f"crowd_dots_{timestamp}.png"
+            
+            plt.savefig(save_path, bbox_inches='tight', pad_inches=0, dpi=200)
+            print(f"이미지 저장 완료: {save_path}")
+        
+        # Figure를 numpy 배열로 변환
+        fig.canvas.draw()
+        image_from_plot = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8)
+        image_from_plot = image_from_plot.reshape(fig.canvas.get_width_height()[::-1] + (4,))
+        image_from_plot = image_from_plot[:,:,:3]  # RGB로 변환
+        
+        return fig.canvas, image_from_plot
+    
+    def crowd_count(self):
+        """
+        가장 최근 예측의 군중 수를 반환합니다.
+        
+        Returns:
+            float: 예측된 군중 수
+            None: 아직 예측이 수행되지 않은 경우
+        """
+        return self.count
+    
+    def get_density_map(self):
+        """
+        가장 최근 예측의 밀도 맵을 반환합니다.
+        
+        Returns:
+            numpy.ndarray: 밀도 맵
+            None: 아직 예측이 수행되지 않은 경우
+        """
+        return self.density_map

custom/clip_ebc_tensorrt.py

@@ -0,0 +1,603 @@
+import os
+import sys
+import torch
+import numpy as np
+import tensorrt as trt
+from typing import Union, Tuple, Optional
+from PIL import Image
+import matplotlib.pyplot as plt
+from torchvision.transforms import ToTensor, Normalize
+from torchvision.transforms.functional import normalize, to_pil_image
+import json
+import datetime
+from scipy.ndimage import gaussian_filter
+from sklearn.cluster import KMeans
+import assets
+
+# 프로젝트 루트 디렉토리 설정
+project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.append(project_root)
+
+class ClipEBCTensorRT:
+    """
+    CLIP-EBC (Efficient Boundary Counting) TensorRT 버전 이미지 처리 클래스입니다.
+    
+    TensorRT로 변환된 CLIP 모델을 사용하여 이미지를 처리하며, 슬라이딩 윈도우 예측 기능을 포함한
+    다양한 설정 옵션을 제공합니다.
+    """
+    
+    def __init__(self,
+                 engine_path="assets/CLIP_EBC_nwpu_rmse_tensorrt.trt",
+                 truncation=4,
+                 reduction=8,
+                 granularity="fine",
+                 anchor_points="average",
+                 input_size=224,
+                 window_size=224,
+                 stride=224,
+                 dataset_name="qnrf",
+                 mean=(0.485, 0.456, 0.406),
+                 std=(0.229, 0.224, 0.225),
+                 config_dir ="configs"):
+        """CLIPEBC TensorRT 클래스를 설정 매개변수와 함께 초기화합니다."""
+        self.engine_path = engine_path
+        self.truncation = truncation
+        self.reduction = reduction
+        self.granularity = granularity
+        self.anchor_points_type = anchor_points
+        self.input_size = input_size
+        self.window_size = window_size
+        self.stride = stride
+        self.dataset_name = dataset_name
+        self.mean = mean
+        self.std = std
+        self.config_dir = config_dir
+        
+        # 결과 저장용 변수 초기화
+        self.density_map = None
+        self.processed_image = None
+        self.count = None
+        self.original_image = None
+        
+        # TensorRT 엔진 로드
+        print(f"TensorRT 엔진 로드 중: {self.engine_path}")
+        self._load_engine()
+        
+        # 입력 및 출력 이름 설정
+        self.input_name = "input"
+        self.output_name = "output"
+        
+        print(f"TensorRT 엔진 초기화 완료")
+        
+    def _load_engine(self):
+        """TensorRT 엔진을 로드합니다."""
+        # TensorRT 로거 생성
+        TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
+        
+        # 런타임 생성
+        self.runtime = trt.Runtime(TRT_LOGGER)
+        
+        # 엔진 파일 로드
+        with open(self.engine_path, 'rb') as f:
+            engine_data = f.read()
+        
+        # 직렬화된 엔진에서 엔진 생성
+        self.engine = self.runtime.deserialize_cuda_engine(engine_data)
+        
+        # 실행 컨텍스트 생성
+        self.context = self.engine.create_execution_context()
+        
+        # TensorRT 10.x에서는 input_binding/output_binding 대신 네트워크 구조를 확인
+        # 입력과 출력을 가져오는 방법이 변경됨
+        self.num_io_tensors = self.engine.num_io_tensors
+        
+        # 입력과 출력 텐서 이름 찾기
+        self.input_tensor_names = []
+        self.output_tensor_names = []
+        
+        print(f"TensorRT 엔진에서 {self.num_io_tensors}개의 IO 텐서를 찾았습니다")
+        
+        for i in range(self.num_io_tensors):
+            name = self.engine.get_tensor_name(i)
+            is_input = self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT
+            
+            if is_input:
+                self.input_tensor_names.append(name)
+            else:
+                self.output_tensor_names.append(name)
+        
+        # 입력과 출력 이름 설정
+        if not self.input_tensor_names:
+            raise ValueError("엔진에서 입력 텐서를 찾을 수 없습니다.")
+        if not self.output_tensor_names:
+            raise ValueError("엔진에서 출력 텐서를 찾을 수 없습니다.")
+        
+        # 기본 입력 및 출력 이름 설정
+        self.input_name = self.input_tensor_names[0]
+        self.output_name = self.output_tensor_names[0]
+        
+        # 입출력 형태 추출
+        self.input_shape = self.engine.get_tensor_shape(self.input_name)
+        self.output_shape = self.engine.get_tensor_shape(self.output_name)
+        
+        print(f"입력 이름: {self.input_name}, 형태: {self.input_shape}")
+        print(f"출력 이름: {self.output_name}, 형태: {self.output_shape}")
+    
+    def _process_image(self, image: Union[str, np.ndarray]) -> np.ndarray:
+        """
+        이미지를 전처리합니다. 이미지 경로, 넘파이 배열, Streamlit UploadedFile 모두 처리 가능합니다.
+        
+        Args:
+            image: 입력 이미지. 다음 형식 �� 하나여야 합니다:
+                - str: 이미지 파일 경로
+                - np.ndarray: (H, W, 3) 형태의 RGB 이미지
+                - UploadedFile: Streamlit의 업로드된 파일
+                    
+        Returns:
+            np.ndarray: 전처리된 이미지 배열, shape (1, 3, H, W)
+        """
+        to_tensor = ToTensor()
+        normalize = Normalize(mean=self.mean, std=self.std)
+        
+        # 원본 이미지 저장
+        self.original_image = image
+        
+        # 입력 타입에 따른 처리
+        if isinstance(image, str):
+            # 파일 경로인 경우
+            with open(image, "rb") as f:
+                pil_image = Image.open(f).convert("RGB")
+        elif isinstance(image, np.ndarray):
+            # 넘파이 배열인 경우
+            if image.dtype == np.uint8:
+                pil_image = Image.fromarray(image)
+            else:
+                # float 타입인 경우 [0, 1] 범위로 가정하고 변환
+                pil_image = Image.fromarray((image * 255).astype(np.uint8))
+        else:
+            # Streamlit UploadedFile 또는 기타 파일 객체인 경우
+            try:
+                pil_image = Image.open(image).convert("RGB")
+            except Exception as e:
+                raise ValueError(f"지원하지 않는 이미지 형식입니다: {type(image)}") from e
+        
+        # 텐서 변환 및 정규화
+        tensor_image = to_tensor(pil_image)
+        normalized_image = normalize(tensor_image)
+        batched_image = normalized_image.unsqueeze(0)  # (1, 3, H, W)
+        
+        # numpy로 변환
+        numpy_image = batched_image.numpy()
+        
+        return numpy_image
+    
+    def _post_process_image(self, image_tensor):
+        """이미지 텐서를 PIL 이미지로 변환합니다."""
+        # NumPy 배열을 PyTorch 텐서로 변환
+        if isinstance(image_tensor, np.ndarray):
+            image_tensor = torch.from_numpy(image_tensor)
+            
+        # 정규화 역변환
+        image = normalize(
+            image_tensor,
+            mean=[0., 0., 0.],
+            std=[1./self.std[0], 1./self.std[1], 1./self.std[2]]
+        )
+        
+        image = normalize(
+            image,
+            mean=[-self.mean[0], -self.mean[1], -self.mean[2]],
+            std=[1., 1., 1.]
+        )
+        
+        # 배치 차원 제거 및 PIL 이미지로 변환
+        processed_image = to_pil_image(image.squeeze(0))
+        return processed_image
+    def _infer_batch(self, batch_input):
+        """
+        TensorRT 엔진을 사용하여 배치 추론을 수행합니다. (수정 버전)
+        """
+        import pycuda.driver as cuda
+        import pycuda.autoinit
+        import numpy as np
+        
+        batch_size = batch_input.shape[0]
+        
+        # 입력의 형태와 데이터 타입 확인
+        input_shape = (batch_size, 3, self.input_size, self.input_size)
+        print(f"입력 배치 형태: {batch_input.shape}, 데이터 타입: {batch_input.dtype}")
+        
+        # 입력 형태 검증
+        if batch_input.shape != input_shape:
+            print(f"경고: 입력 형태 불일치. 예상: {input_shape}, 실제: {batch_input.shape}")
+            # 필요시 형태 수정
+            batch_input = np.resize(batch_input, input_shape)
+        
+        # 데이터 타입 검증
+        if batch_input.dtype != np.float32:
+            print(f"경고: 입력 데이터 타입 불일치. float32로 변환합니다.")
+            batch_input = batch_input.astype(np.float32)
+        
+        # 동적 배치 크기 설정
+        self.context.set_input_shape(self.input_name, input_shape)
+        
+        # 출력 형태 가져오기
+        output_shape = self.context.get_tensor_shape(self.output_name)
+        output_shape = tuple(output_shape)  # 튜플로 변환하여 안전성 보장
+        print(f"출력 형태: {output_shape}")
+        
+        # -1 값을 실제 배치 크기로 대체
+        if output_shape[0] == -1:
+            output_shape = (batch_size,) + output_shape[1:]
+        
+        # 출력 버퍼 준비
+        output = np.empty(output_shape, dtype=np.float32)
+        
+        # 호스트 메모리 준비 (페이지 잠금 메모리 사용)
+        h_input = cuda.pagelocked_empty(batch_input.shape, dtype=np.float32)
+        h_output = cuda.pagelocked_empty(output_shape, dtype=np.float32)
+        
+        # 입력 데이터 복사
+        np.copyto(h_input, batch_input)
+        
+        # 디바이스 메모리 할당
+        d_input = cuda.mem_alloc(h_input.nbytes)
+        d_output = cuda.mem_alloc(h_output.nbytes)
+        
+        # CUDA 스트림 생성
+        stream = cuda.Stream()
+        
+        try:
+            # 메모리 복사 (호스트 -> 디바이스)
+            cuda.memcpy_htod_async(d_input, h_input, stream)
+            
+            # 텐서 주소 설정
+            self.context.set_tensor_address(self.input_name, int(d_input))
+            self.context.set_tensor_address(self.output_name, int(d_output))
+            
+            # 디버깅 정보 (메모리 주소)
+            print(f"입력 메모리 주소: {int(d_input)}, 출력 메모리 주소: {int(d_output)}")
+            
+            # 실행
+            success = self.context.execute_async_v3(stream_handle=stream.handle)
+            if not success:
+                print("TensorRT 실행 실패")
+                return None
+            
+            # 메모리 복사 (디바이스 -> 호스트)
+            cuda.memcpy_dtoh_async(h_output, d_output, stream)
+            
+            # 스트림 동기화
+            stream.synchronize()
+            
+            # 출력 데이터 복사
+            np.copyto(output, h_output)
+            
+            return output
+            
+        except Exception as e:
+            print(f"TensorRT 추론 중 오류 발생: {str(e)}")
+            import traceback
+            traceback.print_exc()
+            return None
+            
+        finally:
+            # 메모리 해제
+            del stream
+            if 'd_input' in locals():
+                d_input.free()
+            if 'd_output' in locals():
+                d_output.free()
+
+    def sliding_window_predict(self, image: np.ndarray, window_size: Union[int, Tuple[int, int]], 
+                             stride: Union[int, Tuple[int, int]]) -> np.ndarray:
+        """
+        슬라이딩 윈도우 방식으로 이미지 예측을 수행합니다. 겹치는 영역은 평균값을 사용합니다.
+        
+        Args:
+            image (np.ndarray): 형태가 (1, 3, H, W)인 이미지 배열
+            window_size (int or tuple): 윈도우 크기
+            stride (int or tuple): 윈도우 이동 간격
+            
+        Returns:
+            np.ndarray: 예측된 밀도 맵
+        """
+        # CUDA 초기화 (처음 사용할 때만)
+        global cuda
+        if 'cuda' not in globals():
+            import pycuda.driver as cuda
+            cuda.init()
+        
+        # 입력 검증
+        assert len(image.shape) == 4, f"이미지는 4차원 배열이어야 합니다. (1, C, H, W), 현재: {image.shape}"
+        
+        # 윈도우 크기와 스트라이드 설정
+        window_size = (int(window_size), int(window_size)) if isinstance(window_size, (int, float)) else window_size
+        stride = (int(stride), int(stride)) if isinstance(stride, (int, float)) else stride
+        window_size = tuple(window_size)
+        stride = tuple(stride)
+        
+        # 검증
+        assert isinstance(window_size, tuple) and len(window_size) == 2 and window_size[0] > 0 and window_size[1] > 0, \
+            f"윈도우 크기는 양수 정수 튜플 (h, w)이어야 합니다. 현재: {window_size}"
+        assert isinstance(stride, tuple) and len(stride) == 2 and stride[0] > 0 and stride[1] > 0, \
+            f"스트라이드는 양수 정수 튜플 (h, w)이어야 합니다. 현재: {stride}"
+        assert stride[0] <= window_size[0] and stride[1] <= window_size[1], \
+            f"스트라이드는 윈도우 크기보다 작아야 합니다. 현재: {stride}와 {window_size}"
+        
+        image_height, image_width = image.shape[-2:]
+        window_height, window_width = window_size
+        stride_height, stride_width = stride
+        
+        # 슬라이딩 윈도우 수 계산
+        num_rows = int(np.ceil((image_height - window_height) / stride_height) + 1)
+        num_cols = int(np.ceil((image_width - window_width) / stride_width) + 1)
+        
+        # 윈도우 추출
+        windows = []
+        window_positions = []
+        for i in range(num_rows):
+            for j in range(num_cols):
+                x_start, y_start = i * stride_height, j * stride_width
+                x_end, y_end = x_start + window_height, y_start + window_width
+                
+                # 이미지 경계 처리
+                if x_end > image_height:
+                    x_start, x_end = image_height - window_height, image_height
+                if y_end > image_width:
+                    y_start, y_end = image_width - window_width, image_width
+                
+                window = image[:, :, x_start:x_end, y_start:y_end]
+                windows.append(window)
+                window_positions.append((x_start, y_start, x_end, y_end))
+        
+        # 배치 단위로 추론
+        all_preds = []
+        max_batch_size = 8
+        
+        for start_idx in range(0, len(windows), max_batch_size):
+            end_idx = min(start_idx + max_batch_size, len(windows))
+            batch_windows = np.vstack(windows[start_idx:end_idx])  # (batch_size, 3, h, w)
+            
+            # TensorRT 추론
+            batch_preds = self._infer_batch(batch_windows)
+            
+            # Debug 정보
+            # print(f"배치 입력 형태: {batch_windows.shape}, 배치 출력 형태: {batch_preds.shape}")
+            
+            all_preds.extend([batch_preds[i:i+1] for i in range(batch_preds.shape[0])])
+        
+        # 예측 결과를 numpy 배열로 변환
+        preds = np.concatenate(all_preds, axis=0)
+        
+        # 출력 밀도 맵 조립
+        pred_map = np.zeros((preds.shape[1], image_height // self.reduction, image_width // self.reduction), dtype=np.float32)
+        count_map = np.zeros((preds.shape[1], image_height // self.reduction, image_width // self.reduction), dtype=np.float32)
+        
+        idx = 0
+        for i in range(num_rows):
+            for j in range(num_cols):
+                x_start, y_start, x_end, y_end = window_positions[idx]
+                
+                # 출력 영역 계산 (reduction 고려)
+                x_start_out = x_start // self.reduction
+                y_start_out = y_start // self.reduction
+                x_end_out = x_end // self.reduction
+                y_end_out = y_end // self.reduction
+                
+                pred_map[:, x_start_out:x_end_out, y_start_out:y_end_out] += preds[idx]
+                count_map[:, x_start_out:x_end_out, y_start_out:y_end_out] += 1.
+                idx += 1
+        
+        # 겹치는 영역 평균 계산
+        pred_map /= count_map
+        
+        return pred_map
+
+    def resize_density_map(self, density_map: np.ndarray, target_size: Tuple[int, int]) -> np.ndarray:
+        """
+        밀도 맵의 크기를 조정합니다. 총합은 보존됩니다.
+        
+        Args:
+            density_map: 형태가 (C, H, W)인 밀도 맵
+            target_size: 목표 크기 (H', W')
+            
+        Returns:
+            np.ndarray: 크기가 조정된 밀도 맵
+        """
+        from PIL import Image
+        import torch.nn.functional as F
+        import torch
+        
+        # numpy를 torch로 변환
+        if isinstance(density_map, np.ndarray):
+            density_map = torch.from_numpy(density_map)
+        
+        # 배치 차원 추가
+        if density_map.dim() == 3:
+            density_map = density_map.unsqueeze(0)  # (1, C, H, W)
+        
+        current_size = density_map.shape[2:]
+        
+        if current_size[0] == target_size[0] and current_size[1] == target_size[1]:
+            return density_map.squeeze(0).numpy()
+        
+        # 원본 밀도 맵의 총합 계산
+        original_sum = density_map.sum()
+        
+        # 크기 조정 (쌍선형 보간)
+        resized_map = F.interpolate(
+            density_map,
+            size=target_size,
+            mode='bilinear',
+            align_corners=False
+        )
+        
+        # 총합 보존을 위한 스케일링
+        if resized_map.sum() > 0:  # 0으로 나누기 방지
+            resized_map = resized_map * (original_sum / resized_map.sum())
+        
+        return resized_map.squeeze(0).numpy()
+
+    def predict(self, image: Union[str, np.ndarray]) -> float:
+        """
+        이미지에서 군중 계수 예측을 수행합니다.
+        
+        Args:
+            image: 입력 이미지 (경로, 넘파이 배열, 또는 업로드된 파일)
+            
+        Returns:
+            float: 예측된 사람 수
+        """
+        # 이미지 전처리
+        processed_image = self._process_image(image)
+        image_height, image_width = processed_image.shape[-2:]
+        
+        # 슬라이딩 윈도우 예측
+        pred_density = self.sliding_window_predict(
+            processed_image, 
+            self.window_size, 
+            self.stride
+        )
+        
+        # 예측 결과 저장
+        pred_count = pred_density.sum()
+        
+        # 원본 이미지 크기로 밀도 맵 조정
+        resized_pred_density = self.resize_density_map(
+            pred_density, 
+            (image_height, image_width)
+        )
+        
+        # 결과 저장
+        self.processed_image = self._post_process_image(processed_image)
+        self.density_map = resized_pred_density.squeeze()
+        self.count = pred_count
+        
+        return pred_count
+    
+    def visualize_density_map(self, alpha: float = 0.5, save: bool = False, 
+                            save_path: Optional[str] = None):
+        """
+        현재 저장된 예측 결과를 시각화합니다.
+        
+        Args:
+            alpha (float): density map의 투명도 (0~1). 기본값 0.5
+            save (bool): 시각화 결과를 이미지로 저장할지 여부. 기본값 False
+            save_path (str, optional): 저장할 경로. None일 경우 현재 디렉토리에 자동 생성된 이름으로 저장.
+                기본값 None
+                
+        Returns:
+            Tuple[matplotlib.figure.Figure, np.ndarray]:
+                - density map이 오버레이된 matplotlib Figure 객체
+                - RGB 형식의 시각화된 이미지 배열 (H, W, 3)
+        """
+        if self.density_map is None or self.processed_image is None:
+            raise ValueError("먼저 predict 메서드를 실행하여 예측을 수행해야 합니다.")
+        
+        fig, ax = plt.subplots(dpi=200, frameon=False)
+        ax.imshow(self.processed_image)
+        ax.imshow(self.density_map, cmap="jet", alpha=alpha)
+        ax.axis("off")
+        plt.title(f"Count: {self.count:.1f}")
+        
+        if save:
+            if save_path is None:
+                timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+                save_path = f"crowd_density_{timestamp}.png"
+            
+            # 여백 제거하고 저장
+            plt.savefig(save_path, bbox_inches='tight', pad_inches=0, dpi=200)
+            print(f"이미지 저장 완료: {save_path}")
+        
+        fig.canvas.draw()
+        image_from_plot = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8)
+        image_from_plot = image_from_plot.reshape(fig.canvas.get_width_height()[::-1] + (4,))
+        image_from_plot = image_from_plot[:,:,:3]  # RGB로 변환
+        
+        return fig, image_from_plot
+    
+    def visualize_dots(self, dot_size: int = 20, sigma: float = 1, percentile: float = 97, 
+                    save: bool = False, save_path: Optional[str] = None):
+        """
+        예측된 군중 위치를 점으로 표시하여 시각화합니다.
+        
+        Args:
+            dot_size (int): 점의 크기. 기본값 20
+            sigma (float): Gaussian 필터의 sigma 값. 기본값 1
+            percentile (float): 임계값으로 사용할 백분위수 (0-100). 기본값 97
+            save (bool): 시각화 결과를 이미지로 저장할지 여부. 기본값 False
+            save_path (str, optional): 저장할 경로. None일 경우 현재 디렉토리에 자동 생성된 이름으로 저장.
+                기본값 None
+                
+        Returns:
+            Tuple[matplotlib.backends.backend_agg.FigureCanvasBase, np.ndarray]: 
+                - matplotlib figure의 canvas 객체
+                - RGB 형식의 시각화된 이미지 배열 (H, W, 3)
+        """
+        if self.density_map is None or self.processed_image is None:
+            raise ValueError("먼저 predict 메서드를 실행하여 예측을 수행해야 합니다.")
+            
+        adjusted_pred_count = int(round(self.count))
+        
+        fig, ax = plt.subplots(dpi=200, frameon=False)
+        ax.imshow(self.processed_image)
+        
+        filtered_density = gaussian_filter(self.density_map, sigma=sigma)
+        
+        threshold = np.percentile(filtered_density, percentile)
+        candidate_pixels = np.column_stack(np.where(filtered_density >= threshold))
+        
+        if len(candidate_pixels) > adjusted_pred_count:
+            kmeans = KMeans(n_clusters=adjusted_pred_count, random_state=42, n_init=10)
+            kmeans.fit(candidate_pixels)
+            head_positions = kmeans.cluster_centers_.astype(int)
+        else:
+            head_positions = candidate_pixels
+            
+        y_coords, x_coords = head_positions[:, 0], head_positions[:, 1]
+        ax.scatter(x_coords, y_coords, 
+                    c='red',
+                    s=dot_size,
+                    alpha=1.0,
+                    edgecolors='white',
+                    linewidth=1)
+        
+        ax.axis("off")
+        plt.title(f"Count: {self.count:.1f}")
+        
+        if save:
+            if save_path is None:
+                timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+                save_path = f"crowd_dots_{timestamp}.png"
+            
+            plt.savefig(save_path, bbox_inches='tight', pad_inches=0, dpi=200)
+            print(f"이미지 저장 완료: {save_path}")
+        
+        # Figure를 numpy 배열로 변환
+        fig.canvas.draw()
+        image_from_plot = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8)
+        image_from_plot = image_from_plot.reshape(fig.canvas.get_width_height()[::-1] + (4,))
+        image_from_plot = image_from_plot[:,:,:3]  # RGB로 변환
+        
+        return fig.canvas, image_from_plot
+    
+    def crowd_count(self):
+        """
+        가장 최근 예측의 군중 수를 반환합니다.
+        
+        Returns:
+            float: 예측된 군중 수
+            None: 아직 예측이 수행되지 않은 경우
+        """
+        return self.count
+    
+    def get_density_map(self):
+        """
+        가장 최근 예측의 밀도 맵을 반환합니다.
+        
+        Returns:
+            numpy.ndarray: 밀도 맵
+            None: 아직 예측이 수행되지 않은 경우
+        """
+        return self.density_map

custom/json2seg.py

@@ -0,0 +1,16 @@
+import json 
+
+def get_segmentation_by_id(target_id, json_file="/home/jungseoik/data/PR/CLIP-EBC/assets/seg.json" ):
+    with open(json_file, "r", encoding="utf-8") as f:
+        data = json.load(f)
+    
+    # annotations 리스트 가져오기
+    annotations = data.get("annotations", [])
+    
+    # annotations 순회하면서 id가 target_id인 항목 찾기
+    for ann in annotations:
+        if ann.get("id") == target_id:
+            return ann.get("segmentation", None)
+    
+    # 해당 id가 없으면 None 반환
+    return None

custom/mock_gen.py

@@ -0,0 +1,80 @@
+import pandas as pd
+import numpy as np
+import random
+def create_mock_data_heatmap():
+    # 기본 구조 생성
+    sections = [f'구역 {i}' for i in range(1, 7)]
+    months = list(range(1, 13))
+    years = list(range(2020, 2025))
+    
+    # 데이터프레임용 리스트 생성
+    data = []
+    for year in years:
+        for section in sections:
+            for month in months:
+                data.append({
+                    'section': section,
+                    'month': month,
+                    'year': year,
+                    'crowd_count': np.random.randint(30000, 500000)
+                })
+    
+    # DataFrame 생성
+    df = pd.DataFrame(data)
+    return df
+
+def create_mock_data_table():
+    mock_data = {
+    'section': [f'구역 {i}' for i in range(1, 7)],
+    'count': np.random.randint(10000, 300000, 6)  
+}
+
+    df = pd.DataFrame(mock_data)
+    return df
+
+def create_mock_data_donut(min_value=10000, max_value=500000):
+    """
+    가상의 인구 이동 데이터를 생성합니다.
+    
+    Returns:
+        tuple: (인바운드 이동 비율, 아웃바운드 이동 비율)
+    """
+    # 랜덤 값 생성 (10000~500000 사이)
+    inbound = random.randint(min_value, max_value)
+    outbound = random.randint(min_value, max_value)
+    
+    # 전체 값 대비 비율 계산 (0-100 사이의 값으로 변환)
+    total = inbound + outbound
+    inbound_percent = round((inbound / total) * 100)
+    outbound_percent = round((outbound / total) * 100)
+    
+    return inbound_percent, outbound_percent
+
+
+def create_mock_data_inout():
+    """
+    방문객 데이터 랜덤 생성
+    - 이번달 방문객: 150,000 ~ 500,000
+    - 오늘 방문객: 5,000 ~ 100,000
+    - delta는 전월/전일 대비 증감량 (-30% ~ +30%)
+    """
+    # 이번달 방문객 (더 큰 범위)
+    monthly_visitors = random.randint(150000, 500000)
+    monthly_delta = int(monthly_visitors * random.uniform(-0.3, 0.3))  # 30% 범위 내 증감
+    
+    # 오늘 방문객 (더 작은 범위)
+    daily_visitors = random.randint(5000, 100000)
+    daily_delta = int(daily_visitors * random.uniform(-0.3, 0.3))  # 30% 범위 내 증감
+    
+    return {
+        'top': {
+            'state': '이번달 방문객',
+            'visitor': monthly_visitors,
+            'delta': monthly_delta
+        },
+        'bottom': {
+            'state': '오늘 방문객',
+            'visitor': daily_visitors,
+            'delta': daily_delta
+        }
+    }

custom/visual.py

@@ -0,0 +1,100 @@
+import altair as alt
+import pandas as pd
+from typing import Tuple, Literal, Union
+# Heatmap
+def make_heatmap(input_df, input_y, input_x, input_color, input_color_theme):
+    heatmap = alt.Chart(input_df).mark_rect().encode(
+            y=alt.Y(f'{input_y}:O', axis=alt.Axis(title="Month", titleFontSize=18, titlePadding=15, titleFontWeight=900, labelAngle=0)),
+            x=alt.X(f'{input_x}:O', axis=alt.Axis(title="", titleFontSize=18, titlePadding=15, titleFontWeight=900, labelAngle=0)),
+            color=alt.Color(f'max({input_color}):Q',
+                             legend=None,
+                             scale=alt.Scale(scheme=input_color_theme)),
+            stroke=alt.value('black'),
+            strokeWidth=alt.value(0.25),
+        ).properties(width=900
+        ).configure_axis(
+        labelFontSize=12,
+        titleFontSize=12
+        ) 
+    # height=300
+    return heatmap
+
+
+# Donut chart
+def make_donut(
+    input_response: float,
+    input_text: str,
+    input_color: Literal['blue', 'green', 'orange', 'red']
+) -> alt.LayerChart:
+    """
+    Altair를 사용하여 지정된 퍼센트, 레이블, 색상 스키마로 도넛 차트를 생성합니다.
+
+    함수 구조:
+    1. 입력 색상에 따른 색상 스키마 정의
+    2. 두 개의 DataFrame 생성:
+        - 퍼센트 표시를 위한 메인 데이터
+        - 전체 원을 위한 배경 데이터
+    3. 세 개의 레이어 생성:
+        - 배경 원 (plot_bg)
+        - 퍼센트 호 (plot)
+        - 중앙 텍스트 표시
+
+    매개변수:
+    ----------
+    input_response : float
+        표시할 퍼센트 값 (0-100 사이)
+    input_text : str
+        차트에 표시할 레이블 텍스트
+    input_color : str
+        사용할 색상 스키마 ('blue', 'green', 'orange', 'red' 중 하나)
+
+    반환값:
+    -------
+    alt.LayerChart
+        배경, 퍼센트 호, 중앙 텍스트가 결합된 Altair 레이어 차트
+
+    사용 예시:
+    ---------
+    >>> chart = make_donut(75, "완료", "blue")
+    >>> chart.save('donut.html')
+    """
+    if input_color == 'blue':
+        chart_color = ['#29b5e8', '#155F7A']
+    if input_color == 'green':
+        chart_color = ['#27AE60', '#12783D']
+    if input_color == 'orange':
+        chart_color = ['#F39C12', '#875A12']
+    if input_color == 'red':
+        chart_color = ['#E74C3C', '#781F16']
+
+    source = pd.DataFrame({
+        "Topic": ['', input_text],
+        "% value": [100-input_response, input_response]
+    })
+    source_bg = pd.DataFrame({
+        "Topic": ['', input_text],
+        "% value": [100, 0]
+    })
+
+    plot = alt.Chart(source).mark_arc(innerRadius=45, cornerRadius=25).encode(
+        theta="% value",
+        color= alt.Color("Topic:N",
+                        scale=alt.Scale(
+                            #domain=['A', 'B'],
+                            domain=[input_text, ''],
+                            # range=['#29b5e8', '#155F7A']),  # 31333F
+                            range=chart_color),
+                        legend=None),
+    ).properties(width=130, height=130)
+
+    text = plot.mark_text(align='center', color="#29b5e8", font="Lato", fontSize=32, fontWeight=700, fontStyle="italic").encode(text=alt.value(f'{input_response} %'))
+    plot_bg = alt.Chart(source_bg).mark_arc(innerRadius=45, cornerRadius=20).encode(
+        theta="% value",
+        color= alt.Color("Topic:N",
+                        scale=alt.Scale(
+                            # domain=['A', 'B'],
+                            domain=[input_text, ''],
+                            range=chart_color),  # 31333F
+                        legend=None),
+    ).properties(width=130, height=130)
+    return plot_bg + plot + text

losses/__init__.py

@@ -0,0 +1,7 @@
+from .dm_loss import DMLoss
+from .dace_loss import DACELoss
+
+__all__ = [
+    "DMLoss",
+    "DACELoss",
+]

losses/bregman_pytorch.py

@@ -0,0 +1,144 @@
+# Code modified from https://github.com/cvlab-stonybrook/DM-Count/blob/master/losses/bregman_pytorch.py
+import torch
+from torch import Tensor
+from torch.cuda.amp import autocast
+from typing import Union, Tuple, Dict
+
+M_EPS = 1e-16
+
+
+@autocast(enabled=True, dtype=torch.float32)  # avoid numerical instability
+def sinkhorn(
+    a: Tensor,
+    b: Tensor,
+    C: Tensor,
+    reg: float = 1e-1,
+    maxIter: int = 1000,
+    stopThr: float = 1e-9,
+    verbose: bool = False,
+    log: bool = True,
+    eval_freq: int = 10,
+    print_freq: int = 200,
+) -> Union[Tensor, Tuple[Tensor, Dict[str, Tensor]]]:
+    """
+    Solve the entropic regularization optimal transport
+    The input should be PyTorch tensors
+    The function solves the following optimization problem:
+
+    .. math::
+        \gamma = arg\min_\gamma <\gamma,C>_F + reg\cdot\Omega(\gamma)
+        s.t.    \gamma 1 = a
+                \gamma^T 1= b
+                \gamma\geq 0
+    where :
+    - C is the (ns,nt) metric cost matrix
+    - :math:`\Omega` is the entropic regularization term :math:`\Omega(\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})`
+    - a and b are target and source measures (sum to 1)
+    The algorithm used for solving the problem is the Sinkhorn-Knopp matrix scaling algorithm as proposed in [1].
+
+    Parameters
+    ----------
+    a : torch.tensor (na,)
+        samples measure in the target domain
+    b : torch.tensor (nb,)
+        samples in the source domain
+    C : torch.tensor (na,nb)
+        loss matrix
+    reg : float
+        Regularization term > 0
+    maxIter : int, optional
+        Max number of iterations
+    stopThr : float, optional
+        Stop threshol on error ( > 0 )
+    verbose : bool, optional
+        Print information along iterations
+    log : bool, optional
+        record log if True
+
+    Returns
+    -------
+    gamma : (na x nb) torch.tensor
+        Optimal transportation matrix for the given parameters
+    log : dict
+        log dictionary return only if log==True in parameters
+
+    References
+    ----------
+    [1] M. Cuturi, Sinkhorn Distances : Lightspeed Computation of Optimal Transport, Advances in Neural Information Processing Systems (NIPS) 26, 2013
+    See Also
+    --------
+    """
+
+    device = a.device
+    na, nb = C.shape
+
+    # a = a.view(-1, 1)
+    # b = b.view(-1, 1)
+
+    assert na >= 1 and nb >= 1, f"C needs to be 2d. Found C.shape = {C.shape}"
+    assert na == a.shape[0] and nb == b.shape[0], f"Shape of a ({a.shape}) or b ({b.shape}) does not match that of C ({C.shape})"
+    assert reg > 0, f"reg should be greater than 0. Found reg = {reg}"
+    assert a.min() >= 0. and b.min() >= 0., f"Elements in a and b should be nonnegative. Found a.min() = {a.min()}, b.min() = {b.min()}"
+
+    if log:
+        log = {"err": []}
+
+    u = torch.ones((na), dtype=a.dtype).to(device) / na
+    v = torch.ones((nb), dtype=b.dtype).to(device) / nb
+
+    K = torch.empty(C.shape, dtype=C.dtype).to(device)
+    torch.div(C, -reg, out=K)
+    torch.exp(K, out=K)
+
+    b_hat = torch.empty(b.shape, dtype=C.dtype).to(device)
+
+    it = 1
+    err = 1
+
+    # allocate memory beforehand
+    KTu = torch.empty(v.shape, dtype=v.dtype).to(device)
+    Kv = torch.empty(u.shape, dtype=u.dtype).to(device)
+
+    while (err > stopThr and it <= maxIter):
+        upre, vpre = u, v
+        # torch.matmul(u, K, out=KTu)
+        KTu = torch.matmul(u.view(1, -1), K).view(-1)
+        v = torch.div(b, KTu + M_EPS)
+        # torch.matmul(K, v, out=Kv)
+        Kv = torch.matmul(K, v.view(-1, 1)).view(-1)
+        u = torch.div(a, Kv + M_EPS)
+
+        if torch.any(torch.isnan(u)) or torch.any(torch.isnan(v)) or \
+                torch.any(torch.isinf(u)) or torch.any(torch.isinf(v)):
+            print("Warning: numerical errors at iteration", it)
+            u, v = upre, vpre
+            break
+
+        if log and it % eval_freq == 0:
+            # we can speed up the process by checking for the error only all
+            # the eval_freq iterations
+            # below is equivalent to:
+            # b_hat = torch.sum(u.reshape(-1, 1) * K * v.reshape(1, -1), 0)
+            # but with more memory efficient
+            b_hat = (torch.matmul(u.view(1, -1), K) * v.view(1, -1)).view(-1)
+            err = (b - b_hat).pow(2).sum().item()
+            # err = (b - b_hat).abs().sum().item()
+            log["err"].append(err)
+
+        if verbose and it % print_freq == 0:
+            print("iteration {:5d}, constraint error {:5e}".format(it, err))
+
+        it += 1
+
+    if log:
+        log["u"] = u
+        log["v"] = v
+        log["alpha"] = reg * torch.log(u + M_EPS)
+        log["beta"] = reg * torch.log(v + M_EPS)
+
+    # transport plan
+    P = u.reshape(-1, 1) * K * v.reshape(1, -1)
+    if log:
+        return P, log
+    else:
+        return P

losses/dace_loss.py

@@ -0,0 +1,70 @@
+import torch
+from torch import nn, Tensor
+from typing import Any, List, Tuple, Dict
+
+from .dm_loss import DMLoss
+from .utils import _reshape_density
+
+
+class DACELoss(nn.Module):
+    def __init__(
+        self,
+        bins: List[Tuple[float, float]],
+        reduction: int,
+        weight_count_loss: float = 1.0,
+        count_loss: str = "mae",
+        **kwargs: Any
+    ) -> None:
+        super().__init__()
+        assert len(bins) > 0, f"Expected at least one bin, got {bins}"
+        assert all([len(b) == 2 for b in bins]), f"Expected all bins to be of length 2, got {bins}"
+        assert all([b[0] <= b[1] for b in bins]), f"Expected all bins to be in increasing order, got {bins}"
+        self.bins = bins
+        self.reduction = reduction
+        self.cross_entropy_fn = nn.CrossEntropyLoss(reduction="none")
+
+        count_loss = count_loss.lower()
+        assert count_loss in ["mae", "mse", "dmcount"], f"Expected count_loss to be one of ['mae', 'mse', 'dmcount'], got {count_loss}"
+        self.count_loss = count_loss
+        if self.count_loss == "mae":
+            self.use_dm_loss = False
+            self.count_loss_fn = nn.L1Loss(reduction="none")
+        elif self.count_loss == "mse":
+            self.use_dm_loss = False
+            self.count_loss_fn = nn.MSELoss(reduction="none")
+        else:
+            self.use_dm_loss = True
+            assert "input_size" in kwargs, f"Expected input_size to be in kwargs when count_loss='dmcount', got {kwargs}"
+            self.count_loss_fn = DMLoss(reduction=reduction, **kwargs)
+
+        self.weight_count_loss = weight_count_loss
+
+    def _bin_count(self, density_map: Tensor) -> Tensor:
+        class_map = torch.zeros_like(density_map, dtype=torch.long)
+        for idx, (low, high) in enumerate(self.bins):
+            mask = (density_map >= low) & (density_map <= high)
+            class_map[mask] = idx
+        return class_map.squeeze(1)  # remove channel dimension
+
+    def forward(self, pred_class: Tensor, pred_density: Tensor, target_density: Tensor, target_points: List[Tensor]) -> Tuple[Tensor, Dict[str, Tensor]]:
+        target_density = _reshape_density(target_density, reduction=self.reduction) if target_density.shape[-2:] != pred_density.shape[-2:] else target_density
+        assert pred_density.shape == target_density.shape, f"Expected pred_density and target_density to have the same shape, got {pred_density.shape} and {target_density.shape}"
+
+        target_class = self._bin_count(target_density)
+
+        cross_entropy_loss = self.cross_entropy_fn(pred_class, target_class).sum(dim=(-1, -2)).mean()
+
+        if self.use_dm_loss:
+            count_loss, loss_info = self.count_loss_fn(pred_density, target_density, target_points)
+            loss_info["ce_loss"] = cross_entropy_loss.detach()
+        else:
+            count_loss = self.count_loss_fn(pred_density, target_density).sum(dim=(-1, -2, -3)).mean()
+            loss_info = {
+                "ce_loss": cross_entropy_loss.detach(),
+                f"{self.count_loss}_loss": count_loss.detach(),
+            }
+
+        loss = cross_entropy_loss + self.weight_count_loss * count_loss
+        loss_info["loss"] = loss.detach()
+
+        return loss, loss_info

losses/dm_loss.py

@@ -0,0 +1,124 @@
+import torch
+from torch import nn, Tensor
+from torch.cuda.amp import autocast
+from typing import List, Any, Tuple, Dict
+
+from .bregman_pytorch import sinkhorn
+from .utils import _reshape_density
+
+EPS = 1e-8
+
+
+class OTLoss(nn.Module):
+    def __init__(
+        self,
+        input_size: int,
+        reduction: int,
+        norm_cood: bool,
+        num_of_iter_in_ot: int = 100,
+        reg: float = 10.0
+    ) -> None:
+        super().__init__()
+        assert input_size % reduction == 0
+
+        self.input_size = input_size
+        self.reduction = reduction
+        self.norm_cood = norm_cood
+        self.num_of_iter_in_ot = num_of_iter_in_ot
+        self.reg = reg
+
+        # coordinate is same to image space, set to constant since crop size is same
+        self.cood = torch.arange(0, input_size, step=reduction, dtype=torch.float32) + reduction / 2
+        self.density_size = self.cood.size(0)
+        self.cood.unsqueeze_(0) # [1, #cood]
+        self.cood = self.cood / input_size * 2 - 1 if self.norm_cood else self.cood
+        self.output_size = self.cood.size(1)
+
+    @autocast(enabled=True, dtype=torch.float32)  # avoid numerical instability
+    def forward(self, pred_density: Tensor, normed_pred_density: Tensor, target_points: List[Tensor]) -> Tuple[Tensor, float, Tensor]:
+        batch_size = normed_pred_density.size(0)
+        assert len(target_points) == batch_size, f"Expected target_points to have length {batch_size}, but got {len(target_points)}"
+        assert self.output_size == normed_pred_density.size(2)
+        device = pred_density.device
+
+        loss = torch.zeros([1]).to(device)
+        ot_obj_values = torch.zeros([1]).to(device)
+        wd = 0 # Wasserstein distance
+        cood = self.cood.to(device)
+        for idx, points in enumerate(target_points):
+            if len(points) > 0:
+                # compute l2 square distance, it should be source target distance. [#gt, #cood * #cood]
+                points = points / self.input_size * 2 - 1 if self.norm_cood else points
+                x = points[:, 0].unsqueeze_(1)  # [#gt, 1]
+                y = points[:, 1].unsqueeze_(1)
+                x_dist = -2 * torch.matmul(x, cood) + x * x + cood * cood # [#gt, #cood]
+                y_dist = -2 * torch.matmul(y, cood) + y * y + cood * cood
+                y_dist.unsqueeze_(2)
+                x_dist.unsqueeze_(1)
+                dist = y_dist + x_dist
+                dist = dist.view((dist.size(0), -1)) # size of [#gt, #cood * #cood]
+
+                source_prob = normed_pred_density[idx][0].view([-1]).detach()
+                target_prob = (torch.ones([len(points)]) / len(points)).to(device)
+                # use sinkhorn to solve OT, compute optimal beta.
+                P, log = sinkhorn(target_prob, source_prob, dist, self.reg, maxIter=self.num_of_iter_in_ot, log=True)
+                beta = log["beta"] # size is the same as source_prob: [#cood * #cood]
+                ot_obj_values += torch.sum(normed_pred_density[idx] * beta.view([1, self.output_size, self.output_size]))
+                # compute the gradient of OT loss to predicted density (pred_density).
+                # im_grad = beta / source_count - < beta, source_density> / (source_count)^2
+                source_density = pred_density[idx][0].view([-1]).detach()
+                source_count = source_density.sum()
+                gradient_1 = (source_count) / (source_count * source_count+ EPS) * beta # size of [#cood * #cood]
+                gradient_2 = (source_density * beta).sum() / (source_count * source_count + EPS) # size of 1
+                gradient = gradient_1 - gradient_2
+                gradient = gradient.detach().view([1, self.output_size, self.output_size])
+                # Define loss = <im_grad, predicted density>. The gradient of loss w.r.t predicted density is im_grad.
+                loss += torch.sum(pred_density[idx] * gradient)
+                wd += torch.sum(dist * P).item()
+
+        return loss, wd, ot_obj_values
+
+
+class DMLoss(nn.Module):
+    def __init__(
+        self,
+        input_size: int,
+        reduction: int,
+        norm_cood: bool = False,
+        weight_ot: float = 0.1,
+        weight_tv: float = 0.01,
+        **kwargs: Any
+    ) -> None:
+        super().__init__()
+        self.ot_loss = OTLoss(input_size, reduction, norm_cood, **kwargs)
+        self.tv_loss = nn.L1Loss(reduction="none")
+        self.count_loss = nn.L1Loss(reduction="mean")
+        self.weight_ot = weight_ot
+        self.weight_tv = weight_tv
+
+    @autocast(enabled=True, dtype=torch.float32)  # avoid numerical instability
+    def forward(self, pred_density: Tensor, target_density: Tensor, target_points: List[Tensor]) -> Tuple[Tensor, Dict[str, Tensor]]:
+        target_density = _reshape_density(target_density, reduction=self.ot_loss.reduction) if target_density.shape[-2:] != pred_density.shape[-2:] else target_density
+        assert pred_density.shape == target_density.shape, f"Expected pred_density and target_density to have the same shape, got {pred_density.shape} and {target_density.shape}"
+
+        pred_count = pred_density.view(pred_density.shape[0], -1).sum(dim=1)
+        normed_pred_density = pred_density / (pred_count.view(-1, 1, 1, 1) + EPS)
+        target_count = torch.tensor([len(p) for p in target_points], dtype=torch.float32).to(target_density.device)
+        normed_target_density = target_density / (target_count.view(-1, 1, 1, 1) + EPS)
+
+        ot_loss, _, _ = self.ot_loss(pred_density, normed_pred_density, target_points)
+
+        tv_loss = (self.tv_loss(normed_pred_density, normed_target_density).sum(dim=(1, 2, 3)) * target_count).mean()
+
+        count_loss = self.count_loss(pred_count, target_count)
+
+        loss = ot_loss * self.weight_ot + tv_loss * self.weight_tv + count_loss
+
+        loss_info = {
+            "loss": loss.detach(),
+            "ot_loss": ot_loss.detach(),
+            "tv_loss": tv_loss.detach(),
+            "count_loss": count_loss.detach(),
+        }
+
+        return loss, loss_info

losses/utils.py

@@ -0,0 +1,9 @@
+from torch import Tensor
+
+
+def _reshape_density(density: Tensor, reduction: int) -> Tensor:
+    assert len(density.shape) == 4, f"Expected 4D (B, 1, H, W) tensor, got {density.shape}"
+    assert density.shape[1] == 1, f"Expected 1 channel, got {density.shape[1]}"
+    assert density.shape[2] % reduction == 0, f"Expected height to be divisible by {reduction}, got {density.shape[2]}"
+    assert density.shape[3] % reduction == 0, f"Expected width to be divisible by {reduction}, got {density.shape[3]}"
+    return density.reshape(density.shape[0], 1, density.shape[2] // reduction, reduction, density.shape[3] // reduction, reduction).sum(dim=(-1, -3))

main.py

@@ -0,0 +1,75 @@
+import argparse
+import os
+from custom.clip_ebc_onnx import ClipEBCOnnx
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='CLIP-EBC Crowd Counting (ONNX)')
+    parser.add_argument('--image', required=True, help='Path to input image')
+    parser.add_argument('--model', default='assets/CLIP_EBC_nwpu_rmse_onnx.onnx', help='Path to ONNX model')
+    parser.add_argument('--visualize', choices=['density', 'dots', 'all', 'none'], 
+                        default='none', help='Visualization type')
+    parser.add_argument('--save', action='store_true', 
+                        help='Save visualization results')
+    parser.add_argument('--output-dir', default='results', 
+                        help='Directory to save results')
+    
+    # 시각화 관련 매개변수
+    parser.add_argument('--alpha', type=float, default=0.5, 
+                        help='Alpha value for density map')
+    parser.add_argument('--dot-size', type=int, default=20, 
+                        help='Dot size for dot visualization')
+    parser.add_argument('--sigma', type=float, default=1, 
+                        help='Sigma value for Gaussian filter')
+    parser.add_argument('--percentile', type=float, default=97, 
+                        help='Percentile threshold for dot visualization')
+    
+
+    
+    return parser.parse_args()
+
+def main():
+    args = parse_args()
+    
+    # 모델 초기화 - ONNX 버전
+    model = ClipEBCOnnx(
+        onnx_model_path=args.model
+    )
+    
+    # 출력 디렉토리 생성
+    if args.save:
+        os.makedirs(args.output_dir, exist_ok=True)
+    
+    # 예측 수행
+    count = model.predict(args.image)
+    print(f"예측된 군중 수: {count:.2f}")
+    
+    # 시각화
+    if args.visualize in ['density', 'all']:
+        save_path = os.path.join(args.output_dir, 'density_map.png') if args.save else None
+        fig, density_map = model.visualize_density_map(
+            alpha=args.alpha,
+            save=args.save,
+            save_path=save_path
+        )
+    
+    if args.visualize in ['dots', 'all']:
+        save_path = os.path.join(args.output_dir, 'dot_map.png') if args.save else None
+        canvas, dot_map = model.visualize_dots(
+            dot_size=args.dot_size,
+            sigma=args.sigma,
+            percentile=args.percentile,
+            save=args.save,
+            save_path=save_path
+        )
+        
+        # matplotlib figure 닫기 (메모리 누수 방지)
+        if args.visualize in ['density', 'all']:
+            import matplotlib.pyplot as plt
+            plt.close(fig)
+        
+        if args.visualize in ['dots', 'all']:
+            import matplotlib.pyplot as plt
+            plt.close(canvas.figure)
+
+if __name__ == "__main__":
+    main()

models/__init__.py

@@ -0,0 +1,49 @@
+from typing import List, Tuple, Optional, Any, Union
+
+from .model import _classifier, _regressor, Classifier, Regressor
+from .clip import _clip_ebc, CLIP_EBC
+import assets
+
+clip_names = ["resnet50", "resnet50x4", "resnet50x16", "resnet50x64", "resnet101", "vit_b_16", "vit_b_32", "vit_l_14"]
+
+
+def get_model(
+    backbone: str,
+    input_size: int,
+    reduction: int,
+    bins: Optional[List[Tuple[float, float]]] = None,
+    anchor_points: Optional[List[float]] = None,
+    **kwargs: Any,
+) -> Union[Regressor, Classifier, CLIP_EBC]:
+    backbone = backbone.lower()
+    if "clip" in backbone:
+        backbone = backbone[5:]
+        assert backbone in clip_names, f"Expected backbone to be in {clip_names}, got {backbone}"
+        return _clip_ebc(
+            backbone=backbone,
+            input_size=input_size,
+            reduction=reduction,
+            bins=bins,
+            anchor_points=anchor_points,
+            **kwargs
+        )
+    elif bins is None and anchor_points is None:
+        return _regressor(
+            backbone=backbone,
+            input_size=input_size,
+            reduction=reduction,
+        )
+    else:
+        assert bins is not None and anchor_points is not None, f"Expected bins and anchor_points to be both None or not None, got {bins} and {anchor_points}"
+        return _classifier(
+            backbone=backbone,
+            input_size=input_size,
+            reduction=reduction,
+            bins=bins,
+            anchor_points=anchor_points,
+        )
+
+
+__all__ = [
+    "get_model",
+]

models/clip/__init__.py

@@ -0,0 +1,7 @@
+from .model import CLIP_EBC, _clip_ebc
+
+
+__all__ = [
+    "CLIP_EBC",
+    "_clip_ebc",
+]

models/clip/_clip/__init__.py

@@ -0,0 +1,273 @@
+import torch
+import os
+from typing import Tuple, Optional, Any, Union
+import json
+
+from .utils import tokenize, transform
+from .prepare import prepare
+from .text_encoder import CLIPTextEncoder
+from .image_encoder import ModifiedResNet, VisionTransformer
+from .model import CLIP
+
+
+curr_dir = os.path.dirname(os.path.abspath(__file__))
+
+clip_model_names = [
+    "clip_resnet50",
+    "clip_resnet101",
+    "clip_resnet50x4",
+    "clip_resnet50x16",
+    "clip_resnet50x64",
+    "clip_vit_b_32",
+    "clip_vit_b_16",
+    "clip_vit_l_14",
+    "clip_vit_l_14_336px",
+]
+
+clip_image_encoder_names = [f"clip_image_encoder_{name[5:]}" for name in clip_model_names]
+clip_text_encoder_names = [f"clip_text_encoder_{name[5:]}" for name in clip_model_names]
+
+
+for name in clip_model_names + clip_image_encoder_names + clip_text_encoder_names:
+    model_weights_path = os.path.join(curr_dir, "weights", f"{name}.pth")
+    model_config_path = os.path.join(curr_dir, "configs", f"{name}.json")
+    if not os.path.exists(os.path.join(curr_dir, "weights", f"{name}.pth")) or not os.path.exists(os.path.join(curr_dir, "configs", f"{name}.json")):
+        prepare()
+        break
+
+
+for name in clip_model_names + clip_image_encoder_names + clip_text_encoder_names:
+    assert os.path.exists(os.path.join(curr_dir, "weights", f"{name}.pth")), f"Missing {name}.pth in weights folder. Please run models/clip/prepare.py to download the weights."
+    assert os.path.exists(os.path.join(curr_dir, "configs", f"{name}.json")), f"Missing {name}.json in configs folder. Please run models/clip/prepare.py to download the configs."
+
+
+def _clip(name: str, input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    with open(os.path.join(curr_dir, "configs", f"clip_{name}.json"), "r") as f:
+        config = json.load(f)
+
+    model = CLIP(
+        embed_dim=config["embed_dim"],
+        # vision
+        image_resolution=config["image_resolution"],
+        vision_layers=config["vision_layers"],
+        vision_width=config["vision_width"],
+        vision_patch_size=config["vision_patch_size"],
+        # text
+        context_length=config["context_length"],
+        vocab_size=config["vocab_size"],
+        transformer_width=config["transformer_width"],
+        transformer_heads=config["transformer_heads"],
+        transformer_layers=config["transformer_layers"]
+    )
+    state_dict = torch.load(os.path.join(curr_dir, "weights", f"clip_{name}.pth"), map_location="cpu")
+    model.load_state_dict(state_dict, strict=True)
+
+    if input_size is not None:
+        input_size = (input_size, input_size) if isinstance(input_size, int) else input_size
+        if name.startswith("vit"):
+            model.visual.adjust_pos_embed(*input_size)
+
+    return model
+
+
+def _resnet(
+    name: str,
+    reduction: int = 32,
+    features_only: bool = False,
+    out_indices: Optional[Tuple[int, ...]] = None,
+    **kwargs: Any
+) -> ModifiedResNet:
+    with open(os.path.join(curr_dir, "configs", f"clip_image_encoder_{name}.json"), "r") as f:
+        config = json.load(f)
+    model = ModifiedResNet(
+        layers=config["vision_layers"],
+        output_dim=config["embed_dim"],
+        input_resolution=config["image_resolution"],
+        width=config["vision_width"],
+        heads=config["vision_heads"],
+        features_only=features_only,
+        out_indices=out_indices,
+        reduction=reduction
+    )
+    state_dict = torch.load(os.path.join(curr_dir, "weights", f"clip_image_encoder_{name}.pth"), map_location="cpu")
+    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+    if len(missing_keys) > 0 or len(unexpected_keys) > 0:
+        print(f"Missing keys: {missing_keys}")
+        print(f"Unexpected keys: {unexpected_keys}")
+    else:
+        print(f"All keys matched successfully.")
+
+    return model
+
+
+def _vit(name: str, features_only: bool = False, input_size: Optional[Union[int, Tuple[int, int]]] = None, **kwargs: Any) -> VisionTransformer:
+    with open(os.path.join(curr_dir, "configs", f"clip_image_encoder_{name}.json"), "r") as f:
+        config = json.load(f)
+    model = VisionTransformer(
+        input_resolution=config["image_resolution"],
+        patch_size=config["vision_patch_size"],
+        output_dim=config["embed_dim"],
+        width=config["vision_width"],
+        layers=config["vision_layers"],
+        heads=config["vision_heads"],
+        features_only=features_only
+    )
+    state_dict = torch.load(os.path.join(curr_dir, "weights", f"clip_image_encoder_{name}.pth"), map_location="cpu")
+    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+    if len(missing_keys) > 0 or len(unexpected_keys) > 0:
+        print(f"Missing keys: {missing_keys}")
+        print(f"Unexpected keys: {unexpected_keys}")
+    else:
+        print(f"All keys matched successfully.")
+
+    if input_size is not None:
+        input_size = (input_size, input_size) if isinstance(input_size, int) else input_size
+        model.adjust_pos_embed(*input_size)
+    return model
+
+
+def _text_encoder(name: str) -> CLIPTextEncoder:
+    with open(os.path.join(curr_dir, "configs", f"clip_text_encoder_{name}.json"), "r") as f:
+        config = json.load(f)
+    model = CLIPTextEncoder(
+        embed_dim=config["embed_dim"],
+        context_length=config["context_length"],
+        vocab_size=config["vocab_size"],
+        transformer_width=config["transformer_width"],
+        transformer_heads=config["transformer_heads"],
+        transformer_layers=config["transformer_layers"]
+    )
+    state_dict = torch.load(os.path.join(curr_dir, "weights", f"clip_text_encoder_{name}.pth"), map_location="cpu")
+    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+    if len(missing_keys) > 0 or len(unexpected_keys) > 0:
+        print(f"Missing keys: {missing_keys}")
+        print(f"Unexpected keys: {unexpected_keys}")
+    else:
+        print(f"All keys matched successfully.")
+
+    return model
+
+
+
+# CLIP models
+def resnet50_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("resnet50", input_size)
+
+def resnet101_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("resnet101", input_size)
+
+def resnet50x4_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("resnet50x4", input_size)
+
+def resnet50x16_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("resnet50x16", input_size)
+
+def resnet50x64_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("resnet50x64", input_size)
+
+def vit_b_32_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("vit_b_32", input_size)
+
+def vit_b_16_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("vit_b_16", input_size)
+
+def vit_l_14_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("vit_l_14", input_size)
+
+def vit_l_14_336px_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("vit_l_14_336px", input_size)
+
+
+# CLIP image encoders
+def resnet50_img(features_only: bool = False, out_indices: Optional[Tuple[int, ...]] = None, **kwargs: Any) -> ModifiedResNet:
+    return _resnet("resnet50", features_only=features_only, out_indices=out_indices, **kwargs)
+
+def resnet101_img(features_only: bool = False, out_indices: Optional[Tuple[int, ...]] = None, **kwargs: Any) -> ModifiedResNet:
+    return _resnet("resnet101", features_only=features_only, out_indices=out_indices, **kwargs)
+
+def resnet50x4_img(features_only: bool = False, out_indices: Optional[Tuple[int, ...]] = None, **kwargs: Any) -> ModifiedResNet:
+    return _resnet("resnet50x4", features_only=features_only, out_indices=out_indices, **kwargs)
+
+def resnet50x16_img(features_only: bool = False, out_indices: Optional[Tuple[int, ...]] = None, **kwargs: Any) -> ModifiedResNet:
+    return _resnet("resnet50x16", features_only=features_only, out_indices=out_indices, **kwargs)
+
+def resnet50x64_img(features_only: bool = False, out_indices: Optional[Tuple[int, ...]] = None, **kwargs: Any) -> ModifiedResNet:
+    return _resnet("resnet50x64", features_only=features_only, out_indices=out_indices, **kwargs)
+
+def vit_b_32_img(features_only: bool = False, input_size: Optional[Union[int, Tuple[int, int]]] = None, **kwargs: Any) -> VisionTransformer:
+    return _vit("vit_b_32", features_only=features_only, input_size=input_size, **kwargs)
+
+def vit_b_16_img(features_only: bool = False, input_size: Optional[Union[int, Tuple[int, int]]] = None, **kwargs: Any) -> VisionTransformer:
+    return _vit("vit_b_16", features_only=features_only, input_size=input_size, **kwargs)
+
+def vit_l_14_img(features_only: bool = False, input_size: Optional[Union[int, Tuple[int, int]]] = None, **kwargs: Any) -> VisionTransformer:
+    return _vit("vit_l_14", features_only=features_only, input_size=input_size, **kwargs)
+
+def vit_l_14_336px_img(features_only: bool = False, input_size: Optional[Union[int, Tuple[int, int]]] = None, **kwargs: Any) -> VisionTransformer:
+    return _vit("vit_l_14_336px", features_only=features_only, input_size=input_size, **kwargs)
+
+
+# CLIP text encoders
+def resnet50_txt() -> CLIPTextEncoder:
+    return _text_encoder("resnet50")
+
+def resnet101_txt() -> CLIPTextEncoder:
+    return _text_encoder("resnet101")
+
+def resnet50x4_txt() -> CLIPTextEncoder:
+    return _text_encoder("resnet50x4")
+
+def resnet50x16_txt() -> CLIPTextEncoder:
+    return _text_encoder("resnet50x16")
+
+def resnet50x64_txt() -> CLIPTextEncoder:
+    return _text_encoder("resnet50x64")
+
+def vit_b_32_txt() -> CLIPTextEncoder:
+    return _text_encoder("vit_b_32")
+
+def vit_b_16_txt() -> CLIPTextEncoder:
+    return _text_encoder("vit_b_16")
+
+def vit_l_14_txt() -> CLIPTextEncoder:
+    return _text_encoder("vit_l_14")
+
+def vit_l_14_336px_txt() -> CLIPTextEncoder:
+    return _text_encoder("vit_l_14_336px")
+
+
+__all__ = [
+    # utils
+    "tokenize",
+    "transform",
+    # clip models
+    "resnet50_clip",
+    "resnet101_clip",
+    "resnet50x4_clip",
+    "resnet50x16_clip",
+    "resnet50x64_clip",
+    "vit_b_32_clip",
+    "vit_b_16_clip",
+    "vit_l_14_clip",
+    "vit_l_14_336px_clip",
+    # clip image encoders
+    "resnet50_img",
+    "resnet101_img",
+    "resnet50x4_img",
+    "resnet50x16_img",
+    "resnet50x64_img",
+    "vit_b_32_img",
+    "vit_b_16_img",
+    "vit_l_14_img",
+    "vit_l_14_336px_img",
+    # clip text encoders
+    "resnet50_txt",
+    "resnet101_txt",
+    "resnet50x4_txt",
+    "resnet50x16_txt",
+    "resnet50x64_txt",
+    "vit_b_32_txt",
+    "vit_b_16_txt",
+    "vit_l_14_txt",
+    "vit_l_14_336px_txt",
+]

models/clip/_clip/blocks.py

@@ -0,0 +1,137 @@
+import torch
+from torch import nn, Tensor
+import torch.nn.functional as F
+from collections import OrderedDict
+from typing import Optional, Iterable
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16."""
+
+    def forward(self, x: Tensor):
+        orig_type = x.dtype
+        ret = super().forward(x.type(torch.float32))
+        return ret.type(orig_type)
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x: Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, d_model: int, n_head: int, attn_mask: Tensor = None):
+        super().__init__()
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.ln_1 = LayerNorm(d_model)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, d_model * 4)),
+            ("gelu", QuickGELU()),
+            ("c_proj", nn.Linear(d_model * 4, d_model))
+        ]))
+        self.ln_2 = LayerNorm(d_model)
+        self.attn_mask = attn_mask
+
+    def attention(self, x: Tensor):
+        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
+        return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = x + self.attention(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self, width: int, layers: int, heads: int, attn_mask: Tensor = None):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)])
+
+    def forward(self, x: Tensor):
+        return self.resblocks(x)
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1):
+        super().__init__()
+
+        # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
+        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.relu2 = nn.ReLU(inplace=True)
+
+        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
+
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu3 = nn.ReLU(inplace=True)
+
+        self.downsample = None
+        self.stride = stride
+
+        if stride > 1 or inplanes != planes * Bottleneck.expansion:
+            # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
+            self.downsample = nn.Sequential(OrderedDict([
+                ("-1", nn.AvgPool2d(stride)),
+                ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
+                ("1", nn.BatchNorm2d(planes * self.expansion))
+            ]))
+
+    def forward(self, x: Tensor):
+        identity = x
+
+        out = self.relu1(self.bn1(self.conv1(x)))
+        out = self.relu2(self.bn2(self.conv2(out)))
+        out = self.avgpool(out)
+        out = self.bn3(self.conv3(out))
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu3(out)
+        return out
+
+
+class AttentionPool2d(nn.Module):
+    def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(torch.randn(spacial_dim + 1, embed_dim) / embed_dim ** 0.5)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
+        self.num_heads = num_heads
+
+    def forward(self, x):
+        x = x.flatten(start_dim=2).permute(2, 0, 1)  # NCHW -> (HW)NC
+        x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
+        x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (HW+1)NC
+        x, _ = F.multi_head_attention_forward(
+            query=x[:1], key=x, value=x,
+            embed_dim_to_check=x.shape[-1],
+            num_heads=self.num_heads,
+            q_proj_weight=self.q_proj.weight,
+            k_proj_weight=self.k_proj.weight,
+            v_proj_weight=self.v_proj.weight,
+            in_proj_weight=None,
+            in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
+            bias_k=None,
+            bias_v=None,
+            add_zero_attn=False,
+            dropout_p=0,
+            out_proj_weight=self.c_proj.weight,
+            out_proj_bias=self.c_proj.bias,
+            use_separate_proj_weight=True,
+            training=self.training,
+            need_weights=False
+        )
+        return x.squeeze(0)

models/clip/_clip/bpe_simple_vocab_16e6.txt.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:924691ac288e54409236115652ad4aa250f48203de50a9e4722a6ecd48d6804a
+size 1356917

models/clip/_clip/configs/clip_image_encoder_resnet101.json

@@ -0,0 +1,13 @@
+{
+    "embed_dim": 512,
+    "image_resolution": 224,
+    "vision_layers": [
+        3,
+        4,
+        23,
+        3
+    ],
+    "vision_width": 64,
+    "vision_patch_size": null,
+    "vision_heads": 32
+}

models/clip/_clip/configs/clip_image_encoder_resnet50.json

@@ -0,0 +1,13 @@
+{
+    "embed_dim": 1024,
+    "image_resolution": 224,
+    "vision_layers": [
+        3,
+        4,
+        6,
+        3
+    ],
+    "vision_width": 64,
+    "vision_patch_size": null,
+    "vision_heads": 32
+}

models/clip/_clip/configs/clip_image_encoder_resnet50x16.json

@@ -0,0 +1,13 @@
+{
+    "embed_dim": 768,
+    "image_resolution": 384,
+    "vision_layers": [
+        6,
+        8,
+        18,
+        8
+    ],
+    "vision_width": 96,
+    "vision_patch_size": null,
+    "vision_heads": 48
+}

models/clip/_clip/configs/clip_image_encoder_resnet50x4.json

@@ -0,0 +1,13 @@
+{
+    "embed_dim": 640,
+    "image_resolution": 288,
+    "vision_layers": [
+        4,
+        6,
+        10,
+        6
+    ],
+    "vision_width": 80,
+    "vision_patch_size": null,
+    "vision_heads": 40
+}

models/clip/_clip/configs/clip_image_encoder_resnet50x64.json

@@ -0,0 +1,13 @@
+{
+    "embed_dim": 1024,
+    "image_resolution": 448,
+    "vision_layers": [
+        3,
+        15,
+        36,
+        10
+    ],
+    "vision_width": 128,
+    "vision_patch_size": null,
+    "vision_heads": 64
+}

models/clip/_clip/configs/clip_image_encoder_vit_b_16.json

@@ -0,0 +1,8 @@
+{
+    "embed_dim": 512,
+    "image_resolution": 224,
+    "vision_layers": 12,
+    "vision_width": 768,
+    "vision_patch_size": 16,
+    "vision_heads": 12
+}

models/clip/_clip/configs/clip_image_encoder_vit_b_32.json

@@ -0,0 +1,8 @@
+{
+    "embed_dim": 512,
+    "image_resolution": 224,
+    "vision_layers": 12,
+    "vision_width": 768,
+    "vision_patch_size": 32,
+    "vision_heads": 12
+}

models/clip/_clip/configs/clip_image_encoder_vit_l_14.json

@@ -0,0 +1,8 @@
+{
+    "embed_dim": 768,
+    "image_resolution": 224,
+    "vision_layers": 24,
+    "vision_width": 1024,
+    "vision_patch_size": 14,
+    "vision_heads": 16
+}

models/clip/_clip/configs/clip_image_encoder_vit_l_14_336px.json

@@ -0,0 +1,8 @@
+{
+    "embed_dim": 768,
+    "image_resolution": 336,
+    "vision_layers": 24,
+    "vision_width": 1024,
+    "vision_patch_size": 14,
+    "vision_heads": 16
+}

models/clip/_clip/configs/clip_resnet101.json

@@ -0,0 +1,17 @@
+{
+    "embed_dim": 512,
+    "image_resolution": 224,
+    "vision_layers": [
+        3,
+        4,
+        23,
+        3
+    ],
+    "vision_width": 64,
+    "vision_patch_size": null,
+    "context_length": 77,
+    "vocab_size": 49408,
+    "transformer_width": 512,
+    "transformer_heads": 8,
+    "transformer_layers": 12
+}

models/clip/_clip/configs/clip_resnet50.json

@@ -0,0 +1,17 @@
+{
+    "embed_dim": 1024,
+    "image_resolution": 224,
+    "vision_layers": [
+        3,
+        4,
+        6,
+        3
+    ],
+    "vision_width": 64,
+    "vision_patch_size": null,
+    "context_length": 77,
+    "vocab_size": 49408,
+    "transformer_width": 512,
+    "transformer_heads": 8,
+    "transformer_layers": 12
+}

models/clip/_clip/configs/clip_resnet50x16.json

@@ -0,0 +1,17 @@
+{
+    "embed_dim": 768,
+    "image_resolution": 384,
+    "vision_layers": [
+        6,
+        8,
+        18,
+        8
+    ],
+    "vision_width": 96,
+    "vision_patch_size": null,
+    "context_length": 77,
+    "vocab_size": 49408,
+    "transformer_width": 768,
+    "transformer_heads": 12,
+    "transformer_layers": 12
+}

models/clip/_clip/configs/clip_resnet50x4.json

@@ -0,0 +1,17 @@
+{
+    "embed_dim": 640,
+    "image_resolution": 288,
+    "vision_layers": [
+        4,
+        6,
+        10,
+        6
+    ],
+    "vision_width": 80,
+    "vision_patch_size": null,
+    "context_length": 77,
+    "vocab_size": 49408,
+    "transformer_width": 640,
+    "transformer_heads": 10,
+    "transformer_layers": 12
+}

models/clip/_clip/configs/clip_resnet50x64.json

@@ -0,0 +1,17 @@
+{
+    "embed_dim": 1024,
+    "image_resolution": 448,
+    "vision_layers": [
+        3,
+        15,
+        36,
+        10
+    ],
+    "vision_width": 128,
+    "vision_patch_size": null,
+    "context_length": 77,
+    "vocab_size": 49408,
+    "transformer_width": 1024,
+    "transformer_heads": 16,
+    "transformer_layers": 12
+}

models/clip/_clip/configs/clip_text_encoder_resnet101.json

@@ -0,0 +1,8 @@
+{
+    "embed_dim": 512,
+    "context_length": 77,
+    "vocab_size": 49408,
+    "transformer_width": 512,
+    "transformer_heads": 8,
+    "transformer_layers": 12
+}

models/clip/_clip/configs/clip_text_encoder_resnet50.json

@@ -0,0 +1,8 @@
+{
+    "embed_dim": 1024,
+    "context_length": 77,
+    "vocab_size": 49408,
+    "transformer_width": 512,
+    "transformer_heads": 8,
+    "transformer_layers": 12
+}

models/clip/_clip/configs/clip_text_encoder_resnet50x16.json

@@ -0,0 +1,8 @@
+{
+    "embed_dim": 768,
+    "context_length": 77,
+    "vocab_size": 49408,
+    "transformer_width": 768,
+    "transformer_heads": 12,
+    "transformer_layers": 12
+}

models/clip/_clip/configs/clip_text_encoder_resnet50x4.json

@@ -0,0 +1,8 @@
+{
+    "embed_dim": 640,
+    "context_length": 77,
+    "vocab_size": 49408,
+    "transformer_width": 640,
+    "transformer_heads": 10,
+    "transformer_layers": 12
+}

models/clip/_clip/configs/clip_text_encoder_resnet50x64.json

@@ -0,0 +1,8 @@
+{
+    "embed_dim": 1024,
+    "context_length": 77,
+    "vocab_size": 49408,
+    "transformer_width": 1024,
+    "transformer_heads": 16,
+    "transformer_layers": 12
+}

models/clip/_clip/configs/clip_text_encoder_vit_b_16.json

@@ -0,0 +1,8 @@
+{
+    "embed_dim": 512,
+    "context_length": 77,
+    "vocab_size": 49408,
+    "transformer_width": 512,
+    "transformer_heads": 8,
+    "transformer_layers": 12
+}

models/clip/_clip/configs/clip_text_encoder_vit_b_32.json

@@ -0,0 +1,8 @@
+{
+    "embed_dim": 512,
+    "context_length": 77,
+    "vocab_size": 49408,
+    "transformer_width": 512,
+    "transformer_heads": 8,
+    "transformer_layers": 12
+}

models/clip/_clip/configs/clip_text_encoder_vit_l_14.json

@@ -0,0 +1,8 @@
+{
+    "embed_dim": 768,
+    "context_length": 77,
+    "vocab_size": 49408,
+    "transformer_width": 768,
+    "transformer_heads": 12,
+    "transformer_layers": 12
+}

models/clip/_clip/configs/clip_text_encoder_vit_l_14_336px.json

@@ -0,0 +1,8 @@
+{
+    "embed_dim": 768,
+    "context_length": 77,
+    "vocab_size": 49408,
+    "transformer_width": 768,
+    "transformer_heads": 12,
+    "transformer_layers": 12
+}

models/clip/_clip/configs/clip_vit_b_16.json

@@ -0,0 +1,12 @@
+{
+    "embed_dim": 512,
+    "image_resolution": 224,
+    "vision_layers": 12,
+    "vision_width": 768,
+    "vision_patch_size": 16,
+    "context_length": 77,
+    "vocab_size": 49408,
+    "transformer_width": 512,
+    "transformer_heads": 8,
+    "transformer_layers": 12
+}

models/clip/_clip/configs/clip_vit_b_32.json

@@ -0,0 +1,12 @@
+{
+    "embed_dim": 512,
+    "image_resolution": 224,
+    "vision_layers": 12,
+    "vision_width": 768,
+    "vision_patch_size": 32,
+    "context_length": 77,
+    "vocab_size": 49408,
+    "transformer_width": 512,
+    "transformer_heads": 8,
+    "transformer_layers": 12
+}