Spaces:
Running
Running
| import urllib.request | |
| import tempfile | |
| OPTIMAL_NMS_THRESHOLD = 0.7 | |
| _model_config_url = "https://huggingface.co/TZProject/final_tz_segmentor/resolve/main/final_model_config.yaml" | |
| def get_set_up(): | |
| import torch | |
| TORCH_VERSION = ".".join(torch.__version__.split(".")[:2]) | |
| CUDA_VERSION = torch.__version__.split("+")[-1] | |
| print("torch: ", TORCH_VERSION, "; cuda: ", CUDA_VERSION) | |
| print(f'GPU available: {torch.cuda.is_available()}') | |
| print(torch.cuda.get_device_capability()) | |
| # print("detectron2:", detectron2.__version__) | |
| def load_model(): | |
| # return None | |
| import torch | |
| from detectron2.engine import DefaultPredictor | |
| from detectron2.config import get_cfg | |
| ## define relevant parameters | |
| cfg = get_cfg() | |
| with tempfile.NamedTemporaryFile(suffix=".yaml") as tmp: | |
| print(tmp.name) | |
| urllib.request.urlretrieve(_model_config_url, filename=tmp.name) | |
| cfg.merge_from_file(tmp.name) | |
| if not torch.cuda.is_available(): | |
| cfg.MODEL.DEVICE = "cpu" | |
| else: | |
| cfg.MODEL.DEVICE = 'cuda' | |
| ## when rerouting to use the final model (final_tz_segmentor) USE_FED_LOSS has to be set to False | |
| ## this setting requires the training data to calculate class imbalance that the app will not have access to and cause a runtime error | |
| ## some messages will appear when using the model that certain weights are not being used | |
| ## but these are used during training and not inference and shouldn't affect the model performance | |
| ## code below | |
| cfg.MODEL.ROI_BOX_HEAD.USE_FED_LOSS = False | |
| predictor = DefaultPredictor(cfg) | |
| return predictor | |
| def mask_nms(masks, scores, nms_threshold=OPTIMAL_NMS_THRESHOLD): | |
| """ | |
| Runs class agnostic NMS on masks/segmentations instead of the bounding boxes. | |
| :param masks: (list float) List of coordinates that make up the mask output from the model. | |
| :param scores: (list float) List of corresponding confidence scores given to each mask. | |
| :param nms_threshold: (float) Threshold to apply mask-based class agnostic NMS. | |
| :return masks_kept (list float): List of masks kept after applying NMS. | |
| """ | |
| import supervision as sv | |
| from shapely.geometry.polygon import Polygon | |
| polygons = [] | |
| for mask in masks: | |
| contour = sv.mask_to_polygons(mask) | |
| if len(contour) > 0: | |
| polygons.append(Polygon(contour[0])) | |
| else: | |
| polygons.append(Polygon([])) | |
| order = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True) | |
| masks_kept = [] | |
| while order: | |
| i = order.pop(0) | |
| masks_kept.append(i) | |
| for j in order: | |
| # Calculate the IoU between the two polygons | |
| intersection = polygons[i].intersection(polygons[j]).area | |
| union = polygons[i].union(polygons[j]).area | |
| iou = intersection / union | |
| # Remove masks with IoU greater than the threshold | |
| if iou > nms_threshold: | |
| order.remove(j) | |
| return masks_kept | |
| def apply_nms(prediction, mask=False, cls_agnostic_nms=OPTIMAL_NMS_THRESHOLD): | |
| from torchvision.ops import nms | |
| from detectron2.structures import Instances | |
| if mask: | |
| nms_indices = mask_nms(prediction["instances"].pred_masks.numpy(), | |
| prediction["instances"]._fields["scores"], cls_agnostic_nms) | |
| else: | |
| nms_indices = nms(prediction["instances"].pred_boxes.tensor, | |
| prediction["instances"].scores, cls_agnostic_nms) | |
| pred = {"instances": Instances(image_size=prediction["instances"].image_size, | |
| pred_boxes=prediction["instances"].pred_boxes[nms_indices], | |
| scores=prediction["instances"].scores[nms_indices], | |
| pred_classes=prediction["instances"].pred_classes[nms_indices], | |
| pred_masks=prediction["instances"].pred_masks[nms_indices])} | |
| return pred | |
| if __name__ == '__main__': | |
| # get_set_up() | |
| load_model() |