Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import cv2 | |
| import numpy as np | |
| from ultralytics import YOLO | |
| from PIL import Image | |
| import io | |
| import os | |
| # --- Load models --- | |
| print("Loading models...") | |
| try: | |
| model_v1 = YOLO("v1.pt") | |
| model_v1l = YOLO("v1l.pt") | |
| print("Models loaded successfully!") | |
| except Exception as e: | |
| print(f"Error loading models: {e}") | |
| raise | |
| def ensemble_predict(image, conf_threshold=0.25): | |
| """ | |
| Run ensemble inference on input image using v1 and v1l models | |
| Returns annotated image and detection results | |
| """ | |
| if image is None: | |
| return None, "No image provided" | |
| # Convert PIL Image to OpenCV format | |
| if isinstance(image, Image.Image): | |
| image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) | |
| else: | |
| image_cv = image | |
| # Predict dengan kedua model | |
| results_v1 = model_v1.predict( | |
| source=image_cv, conf=conf_threshold, save=False, verbose=False | |
| )[0] | |
| results_v1l = model_v1l.predict( | |
| source=image_cv, conf=conf_threshold, save=False, verbose=False | |
| )[0] | |
| # Ensemble: voting class dengan weighted confidence | |
| ensemble_results = [] | |
| for box_v1 in results_v1.boxes: | |
| cls_v1 = int(box_v1.cls[0]) | |
| conf_v1 = float(box_v1.conf[0]) | |
| xyxy_v1 = box_v1.xyxy[0].cpu().numpy() | |
| # Cek deteksi terkait di v1l (dengan IoU threshold) | |
| best_match = None | |
| best_iou = 0 | |
| for box_v1l in results_v1l.boxes: | |
| xyxy_v1l = box_v1l.xyxy[0].cpu().numpy() | |
| # Hitung IoU | |
| x1_min, y1_min, x1_max, y1_max = xyxy_v1 | |
| x2_min, y2_min, x2_max, y2_max = xyxy_v1l | |
| inter_x1 = max(x1_min, x2_min) | |
| inter_y1 = max(y1_min, y2_min) | |
| inter_x2 = min(x1_max, x2_max) | |
| inter_y2 = min(y1_max, y2_max) | |
| if inter_x2 > inter_x1 and inter_y2 > inter_y1: | |
| inter_area = (inter_x2 - inter_x1) * (inter_y2 - inter_y1) | |
| box1_area = (x1_max - x1_min) * (y1_max - y1_min) | |
| box2_area = (x2_max - x2_min) * (y2_max - y2_min) | |
| union_area = box1_area + box2_area - inter_area | |
| iou = inter_area / union_area if union_area > 0 else 0 | |
| if iou > best_iou: | |
| best_iou = iou | |
| best_match = box_v1l | |
| if best_iou > 0.3: # Jika ada overlap significant | |
| cls_v1l = int(best_match.cls[0]) | |
| conf_v1l = float(best_match.conf[0]) | |
| # Voting: jika kedua model setuju class, gunakan weighted avg confidence | |
| if cls_v1 == cls_v1l: | |
| final_cls = cls_v1 | |
| final_conf = (conf_v1 + conf_v1l) / 2 | |
| else: | |
| # Jika beda class, ambil yang confidence-nya lebih tinggi | |
| if conf_v1 >= conf_v1l: | |
| final_cls = cls_v1 | |
| final_conf = conf_v1 | |
| else: | |
| final_cls = cls_v1l | |
| final_conf = conf_v1l | |
| else: | |
| # Jika tidak ada match, gunakan v1 saja | |
| final_cls = cls_v1 | |
| final_conf = conf_v1 | |
| ensemble_results.append((xyxy_v1, final_cls, final_conf)) | |
| # --- Create annotated image --- | |
| annotated = image_cv.copy() | |
| # Color mapping untuk setiap class | |
| colors = [(0, 128, 0), (128, 0, 0), (0, 0, 128), (128, 128, 0)] | |
| # Scale font based on image size | |
| image_height = image_cv.shape[0] | |
| base_font_scale = max(0.6, image_height / 1000.0) # Scale with image height | |
| base_thickness = max(2, int(image_height / 500.0)) | |
| results_text = "=== Ensemble Prediction Results ===\n" | |
| for xyxy, cls, conf in ensemble_results: | |
| x1, y1, x2, y2 = xyxy.astype(int) | |
| label = model_v1.names[cls] | |
| text = f"{label} {conf*100:.1f}%" | |
| results_text += f"Label: {label}\nConfidence: {conf:.3f}\nBBox: [{x1}, {y1}, {x2}, {y2}]\n\n" | |
| # Pilih warna berdasarkan class | |
| color = colors[cls % len(colors)] | |
| # Draw bounding box with scaled thickness | |
| cv2.rectangle(annotated, (x1, y1), (x2, y2), color, base_thickness) | |
| # Draw label background with scaled font | |
| font = cv2.FONT_HERSHEY_SIMPLEX | |
| text_size = cv2.getTextSize(text, font, base_font_scale, base_thickness)[0] | |
| cv2.rectangle( | |
| annotated, | |
| (x1, y1 - text_size[1] - 8), | |
| (x1 + text_size[0] + 8, y1), | |
| color, | |
| -1, | |
| ) | |
| cv2.putText( | |
| annotated, | |
| text, | |
| (x1 + 4, y1 - 4), | |
| font, | |
| base_font_scale, | |
| (255, 255, 255), | |
| base_thickness, | |
| ) | |
| if len(ensemble_results) == 0: | |
| results_text += "No detections found." | |
| # Convert BGR to RGB for display | |
| annotated_rgb = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB) | |
| return annotated_rgb, results_text | |
| # --- Create Gradio Interface --- | |
| with gr.Blocks(title="YonkersNet") as demo: | |
| gr.Markdown("# YonkersNet") | |
| gr.Markdown( | |
| "A Yolo model trained for detecting anime breast size, using ensemble method." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| image_input = gr.Image(label="Upload Image", type="pil") | |
| conf_slider = gr.Slider( | |
| minimum=0.0, | |
| maximum=1.0, | |
| value=0.25, | |
| step=0.05, | |
| label="Confidence Threshold", | |
| ) | |
| gr.Markdown("Sometimes at the rare moment the model isn't really confidence yet to predict. Lowering the threshold can be solve that.") | |
| predict_btn = gr.Button("Run Detection", variant="primary") | |
| with gr.Column(): | |
| image_output = gr.Image(label="Detection Result") | |
| results_output = gr.Textbox(label="Detection Results", lines=10) | |
| predict_btn.click( | |
| fn=ensemble_predict, inputs=[image_input, conf_slider], outputs=[image_output, results_output] | |
| ) | |
| if __name__ == "__main__": | |
| demo.queue().launch(share=True) | |