| import subprocess, sys
|
|
|
| subprocess.check_call([
|
| sys.executable, '-m', 'pip', 'install',
|
| '--quiet', '--no-cache-dir',
|
| 'ultralytics>=8.3.0',
|
| 'pillow>=10.0.0',
|
| 'numpy>=1.24.0'
|
| ])
|
| import sys
|
| import threading
|
| import gradio as gr
|
| import spaces
|
| from PIL import Image
|
| import numpy as np
|
|
|
| def _suppress_fd_warnings(unraisable):
|
| if isinstance(unraisable.exc_value, ValueError) and 'Invalid file descriptor' in str(unraisable.exc_value):
|
| return
|
| sys.__unraisablehook__(unraisable)
|
| sys.unraisablehook = _suppress_fd_warnings
|
|
|
| from ultralytics import YOLO
|
|
|
| _models = {}
|
|
|
| def _load_triage():
|
| if 'triage' not in _models:
|
| print('Loading triage model...')
|
| _models['triage'] = YOLO('natural_seabed.pt', task='classify')
|
| return _models['triage']
|
|
|
| def _load_detector(key):
|
| if key not in _models:
|
| print(f'Loading {key} model...')
|
| _models[key] = YOLO(f'{key}.pt')
|
| return _models[key]
|
|
|
| DETECTOR_KEYS = ['shipwreck', 'subpipes', 'mines', 'cylinder', 'ghostvision']
|
|
|
| @spaces.GPU
|
| def predict_sonar(image, detector_choice, conf_cutoff):
|
| if image is None:
|
| return None, {'error': 'No image provided'}
|
|
|
| pil_img = image if isinstance(image, Image.Image) else Image.fromarray(image)
|
| w, h = pil_img.size
|
|
|
|
|
| triage = _load_triage()
|
| t_res = triage.predict(pil_img, imgsz=224, verbose=False)[0]
|
| probs = t_res.probs
|
| prob_clean = float(probs.data[0])
|
| prob_debris = float(probs.data[1])
|
| is_anomaly = prob_debris >= 0.35
|
|
|
|
|
| detections = []
|
| models_to_run = (
|
| DETECTOR_KEYS if detector_choice in ('all', 'auto')
|
| else [detector_choice.lower()]
|
| )
|
|
|
| annotated_img = np.array(pil_img)
|
| for m_key in models_to_run:
|
| if m_key not in DETECTOR_KEYS:
|
| continue
|
| model = _load_detector(m_key)
|
| imgsz = 1024 if m_key == 'shipwreck' else (1536 if m_key == 'cylinder' else 640)
|
| det_conf = max(conf_cutoff, 0.28 if m_key == 'shipwreck' else (0.32 if m_key == 'subpipes' else conf_cutoff))
|
|
|
| res = model.predict(pil_img, conf=det_conf, imgsz=imgsz, verbose=False)[0]
|
| if m_key == models_to_run[0]:
|
| annotated_img = res.plot()
|
|
|
| for box in res.boxes:
|
| cid = int(box.cls[0])
|
| score = float(box.conf[0])
|
| x1, y1, x2, y2 = [float(c) for c in box.xyxy[0].tolist()]
|
| bw = max(0.0, x2 - x1)
|
| bh = max(0.0, y2 - y1)
|
| detections.append({
|
| 'model': m_key, 'class_id': cid,
|
| 'label': model.names[cid],
|
| 'confidence': round(score, 4),
|
| 'bbox': {
|
| 'x1': round(x1,1), 'y1': round(y1,1),
|
| 'x2': round(x2,1), 'y2': round(y2,1),
|
| 'w': round(bw,1), 'h': round(bh,1),
|
| 'pct_x': round((x1/w)*100,2), 'pct_y': round((y1/h)*100,2),
|
| 'pct_w': round((bw/w)*100,2), 'pct_h': round((bh/h)*100,2),
|
| }
|
| })
|
|
|
| return annotated_img, {
|
| 'triage': {
|
| 'status': 'ABNORMALITY FLAGGED' if is_anomaly else 'CLEAN SEABED',
|
| 'is_anomaly': is_anomaly,
|
| 'prob_debris': round(prob_debris, 4),
|
| 'prob_clean': round(prob_clean, 4),
|
| },
|
| 'detections_count': len(detections),
|
| 'detections': detections,
|
| }
|
|
|
| demo = gr.Interface(
|
| fn=predict_sonar,
|
| inputs=[
|
| gr.Image(type='pil', label='Side-Scan Sonar Swath'),
|
| gr.Dropdown(choices=['all'] + DETECTOR_KEYS, value='shipwreck', label='Detector'),
|
| gr.Slider(0.05, 0.95, value=0.28, step=0.01, label='Confidence Cutoff'),
|
| ],
|
| outputs=[
|
| gr.Image(type='numpy', label='Annotated Swath'),
|
| gr.JSON(label='Pipeline Intelligence'),
|
| ],
|
| title='SonarOps Acoustic Intelligence Engine',
|
| description='Dual-Stage AI: Natural Seabed Triage + High-Resolution Target Detectors.',
|
| api_name='predict',
|
| )
|
|
|
| def _preload_all():
|
| print('Pre-loading all 6 models on CPU...')
|
| _load_triage()
|
| for key in DETECTOR_KEYS:
|
| _load_detector(key)
|
| print('All 6 models ready!')
|
|
|
| threading.Thread(target=_preload_all, daemon=True).start()
|
|
|
| demo.queue()
|
| demo.launch(ssr_mode=False)
|
|
|
