Spaces:
Sleeping
Sleeping
| import cv2 as cv | |
| import numpy as np | |
| import gradio as gr | |
| from huggingface_hub import hf_hub_download | |
| from yolox import YoloX | |
| # Download YOLOX model from Hugging Face (optional fallback) | |
| model_path = hf_hub_download( | |
| repo_id="opencv/object_detection_yolox", | |
| filename="object_detection_yolox_2022nov.onnx" | |
| ) | |
| # Initialize YOLOX model | |
| model = YoloX( | |
| modelPath=model_path, | |
| confThreshold=0.5, | |
| nmsThreshold=0.5, | |
| objThreshold=0.5, | |
| backendId=cv.dnn.DNN_BACKEND_OPENCV, | |
| targetId=cv.dnn.DNN_TARGET_CPU | |
| ) | |
| classes = ('person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', | |
| 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', | |
| 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', | |
| 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', | |
| 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', | |
| 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', | |
| 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', | |
| 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', | |
| 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', | |
| 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', | |
| 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', | |
| 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', | |
| 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', | |
| 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush') | |
| def letterbox(srcimg, target_size=(640, 640)): | |
| padded_img = np.ones((target_size[0], target_size[1], 3), dtype=np.float32) * 114.0 | |
| ratio = min(target_size[0] / srcimg.shape[0], target_size[1] / srcimg.shape[1]) | |
| resized_img = cv.resize(srcimg, (int(srcimg.shape[1] * ratio), int(srcimg.shape[0] * ratio)), interpolation=cv.INTER_LINEAR).astype(np.float32) | |
| padded_img[:int(srcimg.shape[0] * ratio), :int(srcimg.shape[1] * ratio)] = resized_img | |
| return padded_img, ratio | |
| def unletterbox(bbox, scale): | |
| return bbox / scale | |
| def visualize(dets, image, scale): | |
| res_img = image.copy() | |
| h, w = res_img.shape[:2] | |
| font_scale = max(0.5, min(w, h) / 640.0 * 0.5) | |
| thickness = max(1, int(font_scale * 2)) | |
| for det in dets: | |
| box = unletterbox(det[:4], scale).astype(np.int32) | |
| score = det[-2] | |
| cls_id = int(det[-1]) | |
| x0, y0, x1, y1 = box | |
| label = '{}:{:.1f}%'.format(classes[cls_id], score * 100) | |
| cv.rectangle(res_img, (x0, y0), (x1, y1), (0, 255, 0), thickness) | |
| (tw, th), _ = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, font_scale, thickness) | |
| cv.rectangle(res_img, (x0, y0), (x0 + tw + 2, y0 + th + 4), (255, 255, 255), -1) | |
| cv.putText(res_img, label, (x0, y0 + th), cv.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 0), thickness) | |
| return res_img | |
| def detect_objects(input_image): | |
| bgr = cv.cvtColor(input_image, cv.COLOR_RGB2BGR) | |
| input_blob, scale = letterbox(cv.cvtColor(bgr, cv.COLOR_BGR2RGB)) | |
| results = model.infer(input_blob) | |
| if results is None or len(results) == 0: | |
| return input_image | |
| vis_image = visualize(results, bgr, scale) | |
| return cv.cvtColor(vis_image, cv.COLOR_BGR2RGB) | |
| # Gradio interface | |
| demo = gr.Interface( | |
| fn=detect_objects, | |
| inputs=gr.Image(type="numpy", label="Upload Image"), | |
| outputs=gr.Image(type="numpy", label="Detected Objects"), | |
| title="YOLOX Object Detection (OpenCV + ONNX)", | |
| description="Upload an image to detect objects using YOLOX ONNX model and OpenCV DNN.", | |
| allow_flagging="never" | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |