Spaces:
Running
Running
File size: 3,200 Bytes
f891951 2f96362 f891951 572116d f891951 572116d f891951 572116d a566b11 572116d 2f96362 572116d f891951 6a26293 572116d f891951 2f96362 f7bc51e f891951 2f96362 0226845 f891951 2f96362 6fb47b9 a04e86f f891951 2f96362 f891951 a566b11 f891951 19c3114 f891951 572116d a566b11 f891951 572116d |
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 |
import gradio as gr
import PIL.Image as Image
from ultralytics import YOLO
import torch
import os
# Load model once at startup
print("Loading YOLOv8m Defence model...")
model = YOLO("yolov8m_defence.pt")
# Set device and optimize for CPU inference
if torch.cuda.is_available():
device = 'cuda'
print("Using GPU acceleration")
else:
device = 'cpu'
print("Using CPU inference")
torch.set_num_threads(2)
model.to(device)
def predict_image(img, conf_threshold, iou_threshold):
"""Predicts objects in an image using YOLOv8m Defence model with adjustable confidence and IOU thresholds."""
try:
results = model.predict(
source=img,
conf=conf_threshold,
iou=iou_threshold,
show_labels=True,
show_conf=True,
imgsz=640,
verbose=False,
device=device
)
for r in results:
im_array = r.plot()
im = Image.fromarray(im_array[..., ::-1])
return im
except Exception as e:
print(f"Error during prediction: {e}")
return img
# Custom CSS for font styling
css = """
body, .gradio-container {
font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif !important;
}
.gr-button {
font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif !important;
font-weight: 500 !important;
}
.gr-box h1 {
font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif !important;
font-weight: 600 !important;
}
.gr-box p {
font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif !important;
}
"""
# Create interface
iface = gr.Interface(
fn=predict_image,
inputs=[
gr.Image(type="pil", label="Upload Image"),
gr.Slider(minimum=0.1, maximum=1.0, value=0.4, step=0.05, label="Confidence threshold"),
gr.Slider(minimum=0.1, maximum=1.0, value=0.5, step=0.05, label="IoU threshold"),
],
outputs=gr.Image(type="pil", label="Detection Results"),
title="Defence Object Detection",
description="""
Upload images to detect military and civilian vehicles, aircraft, and ships using our fine-tuned YOLOv8m model.
**Model Card:** [spencercdz/YOLOv8m_defence](https://huggingface.co/spencercdz/YOLOv8m_defence)
**Detectable Objects (18 categories):** Aircraft (cargo, commercial, fighter, helicopter, etc.),
Vehicles (car, truck, tank, bus, van), Ships (cargo, yacht, cruise, warship, sailboat),
and specialized items (drone, missile).
**Performance:** Running on Free Tier - inference may take up to 40 seconds per image.
""",
examples=[
["examples/test1.jpg", 0.4, 0.5],
["examples/test2.jpg", 0.4, 0.5],
["examples/test3.jpg", 0.4, 0.5],
["examples/test4.jpg", 0.4, 0.5],
["examples/test5.jpg", 0.4, 0.5],
["examples/test6.jpg", 0.4, 0.5],
["examples/test7.jpg", 0.4, 0.5],
["examples/test8.jpg", 0.4, 0.5],
["examples/test9.jpg", 0.4, 0.5],
["examples/test10.jpg", 0.4, 0.5],
],
css=css,
cache_examples=True,
allow_flagging="never"
)
if __name__ == "__main__":
iface.launch(share=True) |