SteelAwsm/constellationclassification

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.1  # Lower confidence threshold14 15def detect_constellation_type(image):16    # Convert image to numpy array17    image_np = np.array(image)18 19    # Perform inference20    results = model(image_np)21 22    # Get all predictions and their confidence scores23    detection_info = []24    for *box, confidence, class_idx in results.xyxy[0]:25        if confidence >= 0.1:  # Filter based on confidence threshold26            label = results.names[int(class_idx)]27            detection_info.append(f"Label: {label}")28 29    # Render bounding boxes on image (modifies image in place)30    results.render()31 32    # Convert the result back to PIL Image (YOLO stores rendered images in results.imgs)33    detected_img = Image.fromarray(results.imgs[0])34 35    # Return the image with all predictions and the corresponding info36    return detected_img, "\n".join(detection_info)37 38# Create the Gradio interface39interface = gr.Interface(fn=detect_constellation_type,40                         inputs=gr.Image(type="pil"),41                         outputs=[gr.Image(type="pil"), gr.Textbox()],42                         title="Constellation Type Detection",43                         description="Upload an image of the sky and detect The constellation inside it!.")44 45# Launch the interface46interface.launch()