reynsh/Flight_Defect_Detection

1import gradio as gr2import torch3from ultralyticsplus import YOLO, render_result4from ultralytics.nn.tasks import DetectionModel5import torch.serialization6from torch.nn.modules.container import Sequential7 8# Add all necessary classes to safe globals9torch.serialization.add_safe_globals([10    DetectionModel,11    Sequential,12    torch.nn.Module,13    torch.nn.Conv2d,14    torch.nn.BatchNorm2d,15    torch.nn.LeakyReLU,16    torch.nn.Upsample,17    torch.nn.MaxPool2d,18    torch.nn.Sequential19])20 21# torch.hub.download_url_to_file(22#     'https://mattpearsonaviation.com/wp-content/uploads/2017/12/IMG_0560.jpg', 'one.jpg')23# # torch.hub.download_url_to_file(24# #     'https://cdn.airplane-pictures.net/images/uploaded-images/2011/11/25/169465.jpg', 'two.jpg')25# torch.hub.download_url_to_file(26#     'https://imgproc.airliners.net/photos/airliners/7/1/9/0767917.jpg?v=v40', 'three.jpg')27 28 29def yoloV8_func(image: gr.Image = None,30                image_size: gr.Slider = 640,31                conf_threshold: gr.Slider = 0.4,32                iou_threshold: gr.Slider = 0.50):33    """This function performs YOLOv8 object detection on the given image.34 35    Args:36        image (gr.inputs.Image, optional): Input image to detect objects on. Defaults to None.37        image_size (gr.inputs.Slider, optional): Desired image size for the model. Defaults to 640.38        conf_threshold (gr.inputs.Slider, optional): Confidence threshold for object detection. Defaults to 0.4.39        iou_threshold (gr.inputs.Slider, optional): Intersection over Union threshold for object detection. Defaults to 0.50.40    """41    # Trained dataset42    model_path = "best.pt"43    model = YOLO(model_path)44 45   46    results = model.predict(image,47                            conf=conf_threshold,48                            iou=iou_threshold,49                            imgsz=image_size)50 51 52    box = results[0].boxes53    print("Object type:", box.cls)54    print("Coordinates:", box.xyxy)55    print("Probability:", box.conf)56 57 58    render = render_result(model=model, image=image, result=results[0])59    return render60 61 62inputs = [63    gr.Image(type="filepath", label="Input Image"),64    gr.Slider(minimum=320, maximum=1280, value=640,65                     step=32, label="Image Size"),66    gr.Slider(minimum=0.0, maximum=1.0, value=0.25,67                     step=0.05, label="Confidence Threshold"),68    gr.Slider(minimum=0.0, maximum=1.0, value=0.45,69                     step=0.05, label="IOU Threshold"),70]71 72 73outputs = gr.Image(type="filepath", label="Output Image")74 75title = "Aircraft Damage Detection"76 77yolo_app = gr.Interface(78    fn=yoloV8_func,79    inputs=inputs,80    outputs=outputs,81    title=title,82)83 84# Launching Gradio interface 85yolo_app.launch(share=True, debug=True)86