find.py

import torch
torch.cuda.set_device(0)


# In[2]:


import torch
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.data.datasets import register_coco_instances
from detectron2.data import DatasetCatalog
import cv2
import pycocotools.mask as mask_util
import numpy as np
from tqdm import tqdm
from detectron2.config import get_cfg
from detectron2 import model_zoo
from fastcore.all import *
from ensemble_boxes import *


# In[3]:


config_yaml_path = "COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml"
config_yaml_path = "Misc/cascade_mask_rcnn_R_50_FPN_3x.yaml"
config_yaml_path = "Misc/cascade_mask_rcnn_X_152_32x8d_FPN_IN5k_gn_dconv.yaml"
val_json_path  = '/media/data/ben/semi_all/val_1.json'
#val_json_path  = '/media/data/ben/k/sar/k5_corrected/val_0.json'
val_image_path = '/media/data/ben/data'
all_models_path = '/media/data/ben/temp/1_b2/*.pth'

print(all_models_path)
n_class        = 10

CLASS_CONFIDENCE_THRESHOLDS = [.15, .50, .15]
CLASS_CONFIDENCE_THRESHOLDS = [.31, .52, .48]
CLASS_CONFIDENCE_THRESHOLDS = [.56, .71, .27]
MIN_PIXELS = [75, 150, 75]
MIN_PIXELS = [0, 0, 0]
MIN_PIXELS = [150, 50, 0]
MIN_PIXELS = [175, 75, 75]
#MIN_PIXELS = [175, 0, 0]

IOU_TH = 0.4

register_coco_instances('sartorius_val',{},val_json_path, val_image_path)


# In[5]:


val_ds = DatasetCatalog.get('sartorius_val')

cfg = get_cfg()

if not torch.cuda.is_available():
    cfg.MODEL.DEVICE='cpu'

print("Mode: ", cfg.MODEL.DEVICE)

cfg.merge_from_file(model_zoo.get_config_file(config_yaml_path))

cfg.INPUT.MASK_FORMAT='bitmask'
cfg.MODEL.ROI_HEADS.NUM_CLASSES = n_class
cfg.TEST.DETECTIONS_PER_IMAGE = 1000


# 
# cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 256

# SIZE = 704 
# cfg.INPUT.MAX_SIZE_TEST = SIZE
# cfg.INPUT.MAX_SIZE_TRAIN = SIZE
# cfg.INPUT.MIN_SIZE_TEST = SIZE
# cfg.INPUT.MIN_SIZE_TRAIN = (SIZE)
# cfg.MODEL.WEIGHTS = weight_path
# predictor = DefaultPredictor(cfg)


# In[6]:


def precision_at(threshold, iou):
    matches = iou > threshold
    if matches.shape == (0,):
        matches = np.array([[]])
    true_positives = np.sum(matches, axis=1) == 1  # Correct objects
    false_positives = np.sum(matches, axis=0) == 0  # Missed objects
    false_negatives = np.sum(matches, axis=1) == 0  # Extra objects
    return np.sum(true_positives), np.sum(false_positives), np.sum(false_negatives)


def convert_contours_to_masks(contours, image_size=(520,704)):
    masks = []    
    for c in contours:
        c = np.asarray(c)
        c = np.squeeze(c)
        c = c.reshape(-1,2)
        mask = np.zeros(image_size, np.uint8)
        mask = cv2.drawContours(mask, [c], -1, 255, -1)
        mask = mask.astype(np.bool)
        masks.append(mask)
    return np.array(masks)


def score(pred_masks, targ, coco_annotation_segmentation='rle'):
#     pred_masks = pred['instances'].pred_masks.cpu().numpy()
    
    if coco_annotation_segmentation == 'rle':
    # if annotation in rle
        enc_preds = [mask_util.encode(np.asarray(p, order='F')) for p in pred_masks]
        enc_targs = list(map(lambda x:x['segmentation'], targ['annotations']))
    
    else:
    # if annotation in polygon
        enc_preds = [mask_util.encode(np.asarray(p, order='F')) for p in pred_masks]
        enc_targs = list(map(lambda x:x['segmentation'], targ['annotations']))
        enc_targs = convert_contours_to_masks(enc_targs)
        enc_targs = [mask_util.encode(np.asarray(p, order='F')) for p in enc_targs]

    ious = mask_util.iou(enc_preds, enc_targs, [0]*len(enc_targs))
    prec = []
    for t in np.arange(0.5, 1.0, 0.05):
        tp, fp, fn = precision_at(t, ious)
        p = tp / (tp + fp + fn)
        prec.append(p)
    return np.mean(prec)


def nms_predictions(classes, scores, bboxes, masks, 
                    iou_th, shape=(520, 704)):
#     print(len(classes))
    he, wd = shape[0], shape[1]
    boxes_list = [[x[0] / wd, x[1] / he, x[2] / wd, x[3] / he]
                  for x in bboxes]
    scores_list = [x for x in scores]
    labels_list = [x for x in classes]
    
    nms_bboxes, nms_scores, nms_classes = nms(
        boxes=[boxes_list], 
        scores=[scores_list], 
        labels=[labels_list], 
        weights=None,
        iou_thr=iou_th
    )
    nms_masks = []
    for s in nms_scores:
        nms_masks.append(masks[scores.index(s)])

    return nms_classes, nms_scores, nms_masks


def get_pred_class_mode(pred_classes, wanted_pred_class_list=[0,1,2], default_pred_class=0):
    pc = pred_classes.detach().cpu().numpy()
    unique, counts = np.unique(pc, return_counts=True)
    classes = dict(zip(unique, counts))
    classes = dict(sorted(classes.items(), key=lambda item: item[1], reverse=True))
    for cls in classes.keys():
        if cls in wanted_pred_class_list:
            return cls
    return default_pred_class


def get_mask(im, iou_th, class_confidence_thresholds, min_pixels):
    pred = predictor(im)
    # pred_class = torch.mode(pred['instances'].pred_classes)[0]
    pred_class = get_pred_class_mode(pred['instances'].pred_classes)
    take = pred['instances'].scores >= class_confidence_thresholds[pred_class]

    pred_masks = pred['instances'].pred_masks[take].cpu().numpy()
    classes = pred['instances'].pred_classes[take].cpu().numpy().tolist()
    scores = pred['instances'].scores[take].cpu().numpy().tolist()
    bboxes = pred['instances'].pred_boxes[take].tensor.cpu().numpy().tolist()
    
    if len(classes):
        classes, scores, pred_masks = nms_predictions(classes, scores, bboxes, pred_masks, iou_th=iou_th)

    final_masks = []
    used = np.zeros(im.shape[:2], dtype=int) 
    for mask in pred_masks:
        mask = mask * (1-used)
        if mask.sum() >= min_pixels[pred_class]: # skip predictions with small area
            final_masks.append(mask)

    return np.array(final_masks).astype(np.bool)


def no_tqdm(r):
    return r


def score_all(iou_th=IOU_TH, class_confidence_thresholds=CLASS_CONFIDENCE_THRESHOLDS, min_pixels=MIN_PIXELS, wanted_class=None, use_tqdm=True):
    scores = []

    if use_tqdm:
        tmp_tqdm = tqdm
    else:
        tmp_tqdm = no_tqdm
    
    for item in tmp_tqdm(val_ds):
        class_id = int(item['annotations'][0]['category_id'])
        if wanted_class is not None and class_id != wanted_class:
            continue
        im =  cv2.imread(item['file_name'])
        pred_mask = get_mask(im, iou_th==iou_th, class_confidence_thresholds=class_confidence_thresholds, min_pixels=min_pixels)
        
        sc = score(pred_mask, item)
        scores.append(sc)
    return np.mean(scores)


# In[8]:




from glob import glob

model_paths = glob(all_models_path)


predictor = None
model_scores = {}
for model_path in tqdm(model_paths):
    predictor = None
    model_name = os.path.basename(model_path)
    cfg.MODEL.WEIGHTS = model_path

    SIZE = 1333
    cfg.INPUT.MAX_SIZE_TEST = SIZE
    cfg.INPUT.MAX_SIZE_TRAIN = SIZE
    cfg.INPUT.MIN_SIZE_TEST = SIZE
    cfg.INPUT.MIN_SIZE_TRAIN = (SIZE,)


    predictor = DefaultPredictor(cfg)
    r = score_all(use_tqdm=False)
#     print(model_name, ": ", r)
    model_scores[model_name] = r
    del predictor
    torch.cuda.empty_cache()
    


# In[12]:

log_txt = []
model_scores = dict(sorted(model_scores.items(), key=lambda item: item[1], reverse=True))
for k, v in model_scores.items():
    print(k, " :", v)
    log_txt.append(str(k) + " : " + str(v))

import os
dname = os.path.dirname(val_json_path)
pname = os.path.basename(val_json_path).replace('.json', 'model_score.txt')
pname = os.path.join(dname, pname)
with open(pname, 'w') as f:
    f.write("\n".join(log_txt))