Spaces:
Sleeping
Sleeping
| import os | |
| import cv2 | |
| import PIL.Image as Image | |
| import gradio as gr | |
| import numpy as np | |
| import math | |
| from pathlib import Path | |
| from ultralytics import ASSETS, YOLO | |
| DIR_NAME = Path(os.path.dirname(__file__)) | |
| DETECTION_MODEL_n = os.path.join(DIR_NAME, 'models', 'YOLOv8-N_CNO_Detection.pt') | |
| DETECTION_MODEL_s = os.path.join(DIR_NAME, 'models', 'YOLOv8-S_CNO_Detection.pt') | |
| DETECTION_MODEL_m = os.path.join(DIR_NAME, 'models', 'YOLOv8-M_CNO_Detection.pt') | |
| DETECTION_MODEL_l = os.path.join(DIR_NAME, 'models', 'YOLOv8-L_CNO_Detection.pt') | |
| DETECTION_MODEL_x = os.path.join(DIR_NAME, 'models', 'YOLOv8-X_CNO_Detection.pt') | |
| # MODEL = os.path.join(DIR_NAME, 'models', 'YOLOv8-M_CNO_Detection.pt') | |
| # model = YOLO(MODEL) | |
| def predict_image(name, model, img, conf_threshold, iou_threshold): | |
| """Predicts and plots labeled objects in an image using YOLOv8 model with adjustable confidence and IOU thresholds.""" | |
| gr.Info("Starting process") | |
| # gr.Warning("Name is empty") | |
| if name == "": | |
| gr.Warning("Name is empty") | |
| if model == 'YOLOv8-N': | |
| CNO_model = YOLO(DETECTION_MODEL_n) | |
| elif model == 'YOLOv8-S': | |
| CNO_model = YOLO(DETECTION_MODEL_s) | |
| elif model == 'YOLOv8-M': | |
| CNO_model = YOLO(DETECTION_MODEL_m) | |
| elif model == 'YOLOv8-L': | |
| CNO_model = YOLO(DETECTION_MODEL_l) | |
| else: | |
| CNO_model = YOLO(DETECTION_MODEL_x) | |
| results = CNO_model.predict( | |
| source=img, | |
| conf=conf_threshold, | |
| iou=iou_threshold, | |
| show_labels=False, | |
| show_conf=False, | |
| imgsz=512, | |
| max_det=1200 | |
| ) | |
| for r in results: | |
| CNO = len(r.boxes) | |
| CNO_coor = np.empty([CNO, 2], dtype=int) | |
| for j in range(CNO): | |
| # w = r.boxes.xywh[j][2] | |
| # h = r.boxes.xywh[j][3] | |
| # area = (math.pi * w * h / 4) * 20 * 20 / (512 * 512) | |
| # total_area += area | |
| # bbox_img = r.orig_img | |
| x = round(r.boxes.xywh[j][0].item()) | |
| y = round(r.boxes.xywh[j][1].item()) | |
| x1 = round(r.boxes.xyxy[j][0].item()) | |
| y1 = round(r.boxes.xyxy[j][1].item()) | |
| x2 = round(r.boxes.xyxy[j][2].item()) | |
| y2 = round(r.boxes.xyxy[j][3].item()) | |
| CNO_coor[j] = [x, y] | |
| cv2.rectangle(r.orig_img, (x1, y1), (x2, y2), (0, 255, 0), 1) | |
| im_array = r.orig_img | |
| im = Image.fromarray(im_array[..., ::-1]) | |
| CNO_count = "CNO Count: " + str(CNO) | |
| return CNO_count, im | |
| iface = gr.Interface( | |
| fn=predict_image, | |
| inputs=[ | |
| gr.Textbox(label="User Name"), | |
| gr.Radio(["YOLOv8-N", "YOLOv8-S", "YOLOv8-M", "YOLOv8-L", "YOLOv8-X"], value="YOLOv8-M"), | |
| # gr.Image(type="filepath", label="Upload Image"), | |
| gr.File(file_types=["image"], file_count="multiple", label="Upload Image"), | |
| gr.Slider(minimum=0, maximum=1, value=0.2, label="Confidence threshold"), | |
| gr.Slider(minimum=0, maximum=1, value=0.5, label="IoU threshold") | |
| ], | |
| outputs=[gr.Label(label="Analysis Results"), gr.Image(type="pil", label="Result")], | |
| title="AFM AI Analysis", | |
| description="Upload images for inference. The YOLOv8-M model is used by default.", | |
| theme=gr.themes.Default() | |
| ) | |
| if __name__ == '__main__': | |
| iface.launch() | |