SteelAwsm/testgrad

1import gradio as gr2import torch3from PIL import Image4import numpy as np5 6# Define the path to the model weights7model_path = 'models/best.pt'8 9# Load the trained YOLOv7 model10model = torch.hub.load('WongKinYiu/yolov7', 'custom', model_path)11 12# Set the confidence threshold for detection13model.conf = 0.02  # Lower confidence threshold14 15def detect_moon_phase(image):16    # Convert image to numpy array17    image_np = np.array(image)18 19    # Perform inference20    results = model(image_np)21 22    # Find the prediction with the highest confidence23    highest_confidence = 024    best_prediction = None25    for i, (*box, confidence, class_idx) in enumerate(results.xyxy[0]):26        if confidence >= 0.025 and confidence > highest_confidence:  # Check confidence threshold27            highest_confidence = confidence28            best_prediction = (i, box, confidence, class_idx)29 30    if best_prediction is not None:31        i, box, confidence, class_idx = best_prediction32        33        # Filter out all other boxes except the highest confidence one34        results.xyxy[0] = results.xyxy[0][i:i+1]  # Keep only the top prediction35        36        # Render bounding boxes on image (modifies image in place)37        results.render()38 39        # Extract label and confidence for the top prediction40        label = results.names[int(class_idx)]41        detection_info = f"Label: {label}"42    else:43        detection_info = "No predictions"44 45    # Convert the result back to PIL Image (YOLO stores rendered images in results.imgs)46    detected_img = Image.fromarray(results.imgs[0])47 48    return detected_img, detection_info49 50# Create the Gradio interface51interface = gr.Interface(fn=detect_moon_phase,52                         inputs=gr.Image(type="pil"),53                         outputs=[gr.Image(type="pil"), gr.Textbox()],54                         title="Moon Phase Detection",55                         description="Upload an image of the moon and detect its phase using YOLOv7.")56 57# Launch the interface58interface.launch()59