Files
segment-anything-ui-gpu/segment_anything_ui/utils/precompute_folder.py
AI-team\cyhan b436a81091 initial_tune
2025-05-12 11:23:49 +09:00

27 lines
748 B
Python

import glob
import os
import cv2
import numpy as np
import torch
import rich
from PIL import Image
import safetensors
from segment_anything import sam_model_registry
from segment_anything_ui.modeling.storable_sam import StorableSam
from segment_anything_ui.config import Config
config = Config()
sam = sam_model_registry[config.get_sam_model_name()](checkpoint=config.default_weights)
allowed_extensions = [".jpg", ".png", ".tif", ".tiff"]
def load_images_from_folder(folder_path):
images = []
for filename in os.listdir(folder_path):
allowed_extensions = [".jpg", ".png"]
if any(filename.endswith(ext) for ext in allowed_extensions):
img = Image.open(os.path.join(folder_path, filename))
return images