Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from ultralytics import YOLO | |
| import numpy as np | |
| import cv2 | |
| model = YOLO('best-seg.pt') | |
| CLASS_NAMES = {0: 'Capped Cell', 1: 'Failed Cell', 2: 'Matured Cell', 3: 'Open Cell', 4: 'Semi-Matured Cell'} | |
| MATURITY_MAP = { | |
| 'Open Cell': {'days': 10, 'percentage': 10, 'description': 'Elongated, open-ended; larva is seen (3-5 days old)'}, | |
| 'Capped Cell': {'days': 7, 'percentage': 40, 'description': 'Partially sealed cell; transition stage (4-6 days old)'}, | |
| 'Semi-Matured Cell': {'days': 5, 'percentage': 70, 'description': 'Uniform color (5-8 days old)'}, | |
| 'Matured Cell': {'days': 2, 'percentage': 95, 'description': 'Conical tip dark; ready to hatch'}, | |
| 'Failed Cell': {'days': 0, 'percentage': 0, 'description': 'Dead cell; failed process'} | |
| } | |
| def analyze(image): | |
| results = model(image, conf=0.25, iou=0.45) | |
| annotated = image.copy() | |
| detections = [] | |
| distribution = {'open': 0, 'capped': 0, 'mature': 0, 'semiMature': 0, 'failed': 0} | |
| color_map = { | |
| 'Open Cell': (255, 165, 0), | |
| 'Capped Cell': (255, 255, 0), | |
| 'Semi-Matured Cell': (225, 105, 65), | |
| 'Matured Cell': (128, 0, 128), | |
| 'Failed Cell': (255, 99, 71) | |
| } | |
| for r in results: | |
| boxes = r.boxes | |
| masks = r.masks | |
| if boxes is not None: | |
| for i, box in enumerate(boxes): | |
| cls = int(box.cls[0]) | |
| conf = float(box.conf[0]) | |
| x1, y1, x2, y2 = map(int, box.xyxy[0].tolist()) | |
| class_name = CLASS_NAMES.get(cls, 'Unknown') | |
| maturity_info = MATURITY_MAP.get(class_name, {'days': 0, 'percentage': 0, 'description': 'Unknown'}) | |
| color = color_map.get(class_name, (255, 255, 255)) | |
| cv2.rectangle(annotated, (x1, y1), (x2, y2), color, 2) | |
| cv2.putText(annotated, f"{class_name} {conf:.0%}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2) | |
| if masks is not None and i < len(masks.data): | |
| mask = masks.data[i].cpu().numpy() | |
| mask_resized = cv2.resize(mask, (annotated.shape[1], annotated.shape[0])) | |
| annotated[mask_resized > 0.5] = annotated[mask_resized > 0.5] * 0.5 + np.array(color) * 0.5 | |
| detections.append({ | |
| 'id': i + 1, | |
| 'type': class_name, | |
| 'confidence': round(conf * 100), | |
| 'bbox': [x1, y1, x2-x1, y2-y1], | |
| 'percentage': maturity_info['percentage'], | |
| 'estimatedHatchingDays': maturity_info['days'], | |
| 'description': maturity_info['description'] | |
| }) | |
| if 'Semi-Matured' in class_name: | |
| distribution['semiMature'] += 1 | |
| elif 'Matured' in class_name: | |
| distribution['mature'] += 1 | |
| elif 'Capped' in class_name: | |
| distribution['capped'] += 1 | |
| elif 'Open' in class_name: | |
| distribution['open'] += 1 | |
| elif 'Failed' in class_name: | |
| distribution['failed'] += 1 | |
| recommendations = [] | |
| if distribution['mature'] > 0: | |
| recommendations.append(f"Monitor {distribution['mature']} mature cell(s) for emergence within 2-3 days") | |
| if distribution['failed'] > 0: | |
| recommendations.append(f"Remove {distribution['failed']} failed cell(s) to prevent disease") | |
| if len(detections) > 5: | |
| recommendations.append('High queen cell count - consider swarm prevention measures') | |
| result = { | |
| 'totalQueenCells': len(detections), | |
| 'cells': detections, | |
| 'maturityDistribution': distribution, | |
| 'recommendations': recommendations if recommendations else ['Continue regular monitoring'] | |
| } | |
| return annotated, result | |
| demo = gr.Interface( | |
| fn=analyze, | |
| inputs=gr.Image(type="numpy"), | |
| outputs=[gr.Image(label="Detected Queen Cells"), gr.JSON(label="Analysis Results")], | |
| title="iBrood Queen Cell Analyzer", | |
| api_name="analyze" | |
| ) | |
| demo.launch() | |