mistralai-distributed / data /object_detection.py
changgyu's picture
Upload 19 files
668bf5d verified
raw
history blame contribute delete
972 Bytes
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