HelloMe / detector /gesture_detector.py
Maulidaaa's picture
Update detector/gesture_detector.py
ae173e0 verified
from ultralytics import YOLO
from PIL import Image
import cv2
import os
from datetime import datetime
# Load YOLO model
model = YOLO("best.pt")
# Folder aman untuk Hugging Face
SAVE_DIR = "/tmp/detected_images"
os.makedirs(SAVE_DIR, exist_ok=True)
# Base URL Hugging Face kamu (ganti sesuai nama Space kamu)
BASE_URL = "https://maulidaaa-hellome.hf.space/detected_images"
def detect_gesture(image_file):
image = Image.open(image_file.stream).convert("RGB")
results = model.predict(image, conf=0.3)
detections = []
if results[0].boxes is not None:
for box in results[0].boxes:
class_id = int(box.cls[0])
confidence = float(box.conf[0])
label = model.names[class_id]
x1, y1, x2, y2 = box.xyxy[0].tolist()
detections.append({
"label": label,
"confidence": round(confidence, 2),
"bbox": {
"x1": round(x1, 2),
"y1": round(y1, 2),
"x2": round(x2, 2),
"y2": round(y2, 2)
}
})
# Simpan gambar hasil
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
filename = f"detection_{timestamp}.jpg"
saved_path = os.path.join(SAVE_DIR, filename)
cv2.imwrite(saved_path, results[0].plot())
return {
"saved_to": f"{BASE_URL}/{filename}",
"detections": detections,
"total_detections": len(detections)
}