Spaces:
No application file
No application file
import os | |
from concurrent.futures import ProcessPoolExecutor | |
def detect_objects(path): | |
from ultralytics import YOLO # κ° νλ‘μΈμ€ λ΄μμ import | |
model = YOLO("yolov8l.pt") # κ° νλ‘μΈμ€ λ΄μμ λ‘λ | |
fname = os.path.basename(path) | |
result = model(path) | |
labels = set([model.names[int(box.cls)] for box in result[0].boxes]) | |
return {"frame": fname, "objects": list(labels)} | |
def detect_all_objects(frame_dir): | |
""" | |
Args: | |
frame_dir (str): νλ μ μ΄λ―Έμ§(.jpg)λ€μ΄ μ μ₯λ ν΄λ κ²½λ‘ | |
Returns: | |
List[Dict[str, Any]]: [{'frame': νμΌλͺ , 'objects': ['person', 'car', ...]}, ...] | |
""" | |
img_paths = sorted([ | |
os.path.join(frame_dir, fname) | |
for fname in os.listdir(frame_dir) | |
if fname.endswith(".jpg") | |
]) | |
with ProcessPoolExecutor() as executor: | |
results = list(executor.map(detect_objects, img_paths)) | |
return results | |