Spaces:
Sleeping
Sleeping
| import PIL.Image as Image | |
| import PIL.ExifTags | |
| import gradio as gr | |
| from ultralytics import YOLO | |
| # Update the path to the YOLOv8 weights file | |
| weights_path = "best (2).pt" | |
| model = YOLO(weights_path) | |
| def correct_image_orientation(img): | |
| """Correct image orientation based on EXIF metadata.""" | |
| try: | |
| exif = img._getexif() | |
| if exif is not None: | |
| orientation_tag = [key for key, value in PIL.ExifTags.TAGS.items() if value == 'Orientation'] | |
| if orientation_tag: | |
| orientation = exif.get(orientation_tag[0]) | |
| if orientation == 3: | |
| img = img.rotate(180, expand=True) | |
| elif orientation == 6: | |
| img = img.rotate(270, expand=True) | |
| elif orientation == 8: | |
| img = img.rotate(90, expand=True) | |
| except (AttributeError, KeyError, IndexError): | |
| # If there is no EXIF data or any error occurs, just return the image as is | |
| pass | |
| return img | |
| def predict_image(img, conf_threshold, iou_threshold): | |
| """Predicts and plots labeled objects in an image using YOLOv8 model with adjustable confidence and IOU thresholds. """ | |
| # Correct the image orientation first | |
| img = correct_image_orientation(img) | |
| # Perform prediction using YOLOv8 model | |
| results = model.predict( | |
| source=img, | |
| conf=conf_threshold, | |
| iou=iou_threshold, | |
| show_labels=True, | |
| show_conf=True, | |
| imgsz=640, | |
| ) | |
| # Convert the results back to a PIL image | |
| if results: | |
| im_array = results[0].plot() | |
| im = Image.fromarray(im_array) | |
| # Ensure the output image is horizontally oriented | |
| im = im.rotate(0, expand=True) | |
| return im | |
| return img | |
| iface = gr.Interface( | |
| fn=predict_image, | |
| inputs=[ | |
| gr.inputs.Image(type="pil", label="Upload Image"), | |
| gr.inputs.Slider(minimum=0, maximum=1, default=0.25, label="Confidence threshold"), | |
| gr.inputs.Slider(minimum=0, maximum=1, default=0.45, label="IoU threshold") | |
| ], | |
| outputs=gr.outputs.Image(type="pil", label="Result"), | |
| title="U-SUB DETECTRON", | |
| description="""Upload images for inference. The Ultralytics YOLOv8n model is used by default. | |
| DO NOT TAKE PICTURE VERTICALLY, current bug prevents model detection due to random rotations. TAKE PICTURE HORIZONTALLY! (input integral image should have bottom of integral facing downwards or leftwards)""", | |
| examples=[ | |
| ["p1.png", 0.25, 0.45], | |
| ["p2.png", 0.25, 0.45], | |
| ["test1 (1).png", 0.25, 0.45], | |
| ["p3.jpeg", 0.25, 0.45], | |
| ["p4.jpg", 0.25, 0.45], | |
| ] | |
| ) | |
| if __name__ == '__main__': | |
| iface.launch() | |