Buckets:
| """Local CPU inference with canonical labels and original-image coordinates.""" | |
| import hashlib | |
| import json | |
| from pathlib import Path | |
| import numpy as np | |
| from .preprocess import load_rgb | |
| from .tiling import iter_tiles, stitch | |
| from .fusion import fuse | |
| def semantic_array(result): | |
| return result.semantic_mask.data.cpu().numpy() | |
| class Model: | |
| def __init__(self, model_dir, *, detection_profile=None): | |
| from .runtime import configure_local_settings | |
| configure_local_settings() | |
| import torch | |
| from ultralytics import YOLO | |
| self.directory = Path(model_dir).resolve() | |
| self.config = json.loads((self.directory / 'model.json').read_text()) | |
| import ultralytics | |
| if ultralytics.__version__ != self.config['framework_version']: | |
| raise RuntimeError('install the pinned framework version from model/requirements.lock') | |
| torch.set_num_threads(4) | |
| self.models = {} | |
| for task, record in self.config['components'].items(): | |
| path = (self.directory / record['path']).resolve() | |
| if not path.is_relative_to(self.directory) or hashlib.sha256(path.read_bytes()).hexdigest() != record['sha256']: | |
| raise ValueError('model checkpoint identity mismatch') | |
| self.models[task] = YOLO(str(path), task='obb' if task == 'detection' else 'semantic') | |
| self.mapping = {int(k):v for k,v in self.config['components']['detection']['source_class_mapping'].items()} | |
| self.thresholds = self.config['detection_thresholds'] | |
| self.profile = None | |
| self.abstained = set() | |
| if detection_profile is not None: | |
| profile_path = (self.directory / detection_profile).resolve() | |
| if not profile_path.is_relative_to(self.directory): | |
| raise ValueError('detection profile must remain inside the model directory') | |
| profile = json.loads(profile_path.read_text()) | |
| if profile['checkpoint_sha256'] != self.config['components']['detection']['sha256']: | |
| raise ValueError('detection profile does not match the checkpoint') | |
| supported = set(self.config['detection_class_ids']) | |
| abstained = set(profile['abstained_class_ids']) | |
| thresholds = profile['thresholds'] | |
| active = {int(k) for k in thresholds} | |
| if (abstained & active) or (abstained | active) != supported or set(thresholds) != {str(c) for c in active}: | |
| raise ValueError('profile must explicitly cover or abstain on every detection class') | |
| if any(isinstance(v,bool) or not np.isfinite(v) or not 0 <= v <= 1 for v in thresholds.values()): | |
| raise ValueError('invalid detection profile threshold') | |
| self.thresholds, self.abstained, self.profile = thresholds, abstained, profile | |
| self.names = {int(k):v for k,v in self.config['classes'].items()} | |
| def predict(self, source, *, confidence=None, tile_size=1024, overlap=.25): | |
| import torch | |
| if confidence is not None and getattr(self,'profile',None) is not None: | |
| raise ValueError('confidence override cannot be combined with a calibrated detection profile') | |
| if confidence is not None and (isinstance(confidence,bool) or not np.isfinite(confidence) or not 0 <= confidence <= 1): | |
| raise ValueError('confidence must be between zero and one') | |
| image, valid = load_rgb(source) | |
| h,w = image.shape[:2] | |
| mask = np.full((h,w),255,np.uint8) | |
| weights = np.zeros((h,w),np.uint16) | |
| objects = [] | |
| for tile in iter_tiles(image,tile_size,overlap): | |
| tile_valid = np.zeros((1024,1024),bool) | |
| tile_valid[:tile.height,:tile.width] = valid[tile.y:tile.y+tile.height,tile.x:tile.x+tile.width] | |
| if not tile_valid.any(): | |
| continue | |
| pixels = tile.image[:,:,::-1].copy() | |
| with torch.inference_mode(): | |
| detection = self.models['detection'].predict(pixels,imgsz=1024,device='cpu',conf=.001,max_det=300,save=False,verbose=False)[0] | |
| segmentation = self.models['segmentation'].predict(pixels,imgsz=1024,device='cpu',save=False,verbose=False)[0] | |
| for points,cls,score in zip(detection.obb.xyxyxyxy.cpu().numpy(),detection.obb.cls.cpu().numpy(),detection.obb.conf.cpu().numpy()): | |
| category = self.mapping.get(int(cls)) | |
| if category is None or category in getattr(self,'abstained',set()) or score < (confidence if confidence is not None else self.thresholds[str(category)]): | |
| continue | |
| center = points.mean(axis=0).astype(int) | |
| if not (0 <= center[0] < tile.width and 0 <= center[1] < tile.height and tile_valid[center[1],center[0]]): | |
| continue | |
| points = points + np.array([tile.x,tile.y]) | |
| low = np.maximum(points.min(axis=0),[0,0]);high=np.minimum(points.max(axis=0),[w,h]) | |
| objects.append({'_tile':(tile.x,tile.y),'category_id':category,'class':self.names[category],'confidence':float(score), | |
| 'obb_polygon':points.reshape(-1).tolist(),'bbox':[float(low[0]),float(low[1]),float(high[0]-low[0]),float(high[1]-low[1])]}) | |
| raw = semantic_array(segmentation) | |
| ids = np.asarray(self.config['segmentation_class_ids'],np.uint8) | |
| if raw.ndim != 2 or not np.issubdtype(raw.dtype,np.integer) or np.any((raw < 0) | (raw >= len(ids))): | |
| raise ValueError('invalid contiguous semantic output') | |
| stitch(mask,weights,tile,ids[raw],tile_valid) | |
| return {'objects':fuse(objects,self.config['tiling']['nms_iou']), 'regions':mask, | |
| 'width':w,'height':h,'coordinate_system':'source_image_pixels','ignore_id':255, | |
| 'detection_policy':{'name':self.profile['name'] if getattr(self,'profile',None) else 'validation_f1', | |
| 'abstained_class_ids':sorted(getattr(self,'abstained',set())), | |
| 'certified_99_precision':False}} | |
Xet Storage Details
- Size:
- 6.07 kB
- Xet hash:
- 4c2f87728a569d8a59fabf0c12222b193b5fbcf98285a2b69d91f845aa5f4f41
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.