Spaces:
Sleeping
Sleeping
File size: 5,670 Bytes
e0388c0 e479bc2 e0388c0 c7ef42a e0388c0 c7ef42a e479bc2 e0388c0 c7ef42a e0388c0 c7ef42a e0388c0 c7ef42a e0388c0 c7ef42a e0388c0 c7ef42a e0388c0 c7ef42a e0388c0 c7ef42a e0388c0 c7ef42a e0388c0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 |
import os
import sys
import io
import base64
import threading
import requests
from flask import Flask, request, jsonify, send_from_directory
from PIL import Image
import torch
import supervision as sv
from ultralytics import YOLO
from rfdetr import RFDETRNano
# Ensure local 'rfdetr' folder is found if present
sys.path.insert(0, os.getcwd())
app = Flask(__name__, static_folder="static")
# --- Constants & Configuration ---
# Map Class IDs to Names (Common for both models if they share the dataset)
CLASS_MAP = {0: 'Gun', 1: 'Explosive', 2: 'Grenade', 3: 'Knife'}
# Weight Paths
RF_WEIGHTS_URL = "https://huggingface.co/Subh775/Threat-Detection-RFDETR/resolve/main/checkpoint_best_total.pth"
RF_WEIGHTS_PATH = "/tmp/rfdetr_best.pth"
YOLO_WEIGHTS_URL = "https://huggingface.co/Subh775/Threat-Detection-YOLOv8n/resolve/main/weights/best.pt"
YOLO_WEIGHTS_PATH = "/tmp/yolov8_best.pt"
# Global Model Instances
models = {
"rf": None,
"yolo": None
}
# --- Utilities ---
def download_if_missing(url, path):
"""Downloads file from URL if it doesn't exist locally."""
if not os.path.exists(path):
print(f"[INFO] Downloading weights: {path}...")
try:
r = requests.get(url, stream=True)
r.raise_for_status()
with open(path, "wb") as f:
for chunk in r.iter_content(chunk_size=8192):
f.write(chunk)
print("[INFO] Download complete.")
except Exception as e:
print(f"[ERROR] Failed to download {url}: {e}")
def get_models():
"""Lazy loader: initializes models only if they aren't ready."""
# 1. Load RF-DETR
if models["rf"] is None and RFDETRNano:
download_if_missing(RF_WEIGHTS_URL, RF_WEIGHTS_PATH)
try:
print("[INFO] Loading RF-DETR Nano...")
models["rf"] = RFDETRNano(pretrain_weights=RF_WEIGHTS_PATH)
except Exception as e:
print(f"[ERROR] RF-DETR Init Failed: {e}")
# 2. Load YOLOv8
if models["yolo"] is None:
download_if_missing(YOLO_WEIGHTS_URL, YOLO_WEIGHTS_PATH)
try:
print("[INFO] Loading YOLOv8...")
models["yolo"] = YOLO(YOLO_WEIGHTS_PATH)
except Exception as e:
print(f"[ERROR] YOLO Init Failed: {e}")
return models["rf"], models["yolo"]
def img_to_base64(img):
"""Encodes PIL Image to Base64 string."""
buf = io.BytesIO()
img.save(buf, format="JPEG", quality=85)
return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode('utf-8')
def base64_to_img(data_str):
"""Decodes Base64 string to PIL Image."""
if "base64," in data_str:
data_str = data_str.split("base64,")[1]
return Image.open(io.BytesIO(base64.b64decode(data_str))).convert("RGB")
def annotate_image(image, detections):
"""
Annotates an image with bounding boxes and labels using Supervision.
Expects detections to be a supervision.Detections object.
"""
# Initialize annotators
box_annotator = sv.BoxAnnotator(thickness=2)
label_annotator = sv.LabelAnnotator(text_scale=0.5, text_padding=4)
# Generate labels: "ClassName Confidence"
labels = []
for class_id, conf in zip(detections.class_id, detections.confidence):
name = CLASS_MAP.get(class_id, str(class_id))
labels.append(f"{name} {conf:.2f}")
# Apply annotations
annotated = image.copy()
annotated = box_annotator.annotate(scene=annotated, detections=detections)
annotated = label_annotator.annotate(scene=annotated, detections=detections, labels=labels)
return annotated
# --- Routes ---
@app.route('/')
def index():
return send_from_directory('static', 'index.html')
@app.route('/predict', methods=['POST'])
def predict():
try:
data = request.json
if not data or 'image' not in data:
return jsonify({"error": "No image data provided"}), 400
# Parse inputs
raw_image = base64_to_img(data['image'])
conf_threshold = float(data.get('conf', 0.25))
# Ensure models are loaded
rf_model, yolo_model = get_models()
# --- Run RF-DETR ---
rf_result_b64 = data['image'] # Fallback to original
if rf_model:
try:
# Predict -> Returns Supervision Detections
detections = rf_model.predict(raw_image, threshold=conf_threshold)
annotated_rf = annotate_image(raw_image, detections)
rf_result_b64 = img_to_base64(annotated_rf)
except Exception as e:
print(f"RF-DETR Inference Error: {e}")
# --- Run YOLOv8 ---
yolo_result_b64 = data['image'] # Fallback to original
if yolo_model:
try:
# Predict -> Returns Ultralytics Results -> Convert to Supervision
results = yolo_model(raw_image, conf=conf_threshold, verbose=False)[0]
detections = sv.Detections.from_ultralytics(results)
annotated_yolo = annotate_image(raw_image, detections)
yolo_result_b64 = img_to_base64(annotated_yolo)
except Exception as e:
print(f"YOLO Inference Error: {e}")
# Return JSON
return jsonify({
"rfdetr": {"image": rf_result_b64},
"yolov8": {"image": yolo_result_b64}
})
except Exception as e:
print(f"Server Error: {e}")
return jsonify({"error": str(e)}), 500
if __name__ == '__main__':
# Pre-load models in background to speed up first request
threading.Thread(target=get_models).start()
app.run(host='0.0.0.0', port=7860) |