Code/Inference.py

""" 
inference on single image for MaskRCNN (FROM DETECTRON) + DLC 

two step, pretrained MaskRCNN, then DLC

"""
import cv2 
import torch

import sys
sys.path.append("Repositories/DeepLabCut-live")

import deeplabcut as dlc
from dlclive import DLCLive, Processor

import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm
import os
import shutil

import torchvision
from torchvision.transforms import transforms as transforms

import pickle

import detectron2

# import some common detectron2 utilities
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog, DatasetCatalog

import cv2



COCO_INSTANCE_CATEGORY_NAMES = [
    '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
    'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign',
    'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
    'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A', 'N/A',
    'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
    'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
    'bottle', 'N/A', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
    'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
    'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table',
    'N/A', 'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
    'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book',
    'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]


def Process_Crop(Crop, CropSize):
    """Crop image and pad, if too big, will scale down """
    # import ipdb;ipdb.set_trace()
    if Crop.shape[0] > CropSize[0] or Crop.shape[1] > CropSize[1]: #Crop is bigger, scale down
        ScaleProportion = min(CropSize[0]/Crop.shape[0],CropSize[1]/Crop.shape[1])
        
        width_scaled = int(Crop.shape[1] * ScaleProportion)
        height_scaled = int(Crop.shape[0] * ScaleProportion)
        Crop = cv2.resize(Crop, (width_scaled,height_scaled), interpolation=cv2.INTER_LINEAR)  # resize image

        # Points2D = {k:[v[0]*ScaleProportion,v[1]*ScaleProportion] for k,v in Points2D.items()}
    else:
        ScaleProportion = 1
        
    if Crop.shape[0] %2 ==0:
        #Shape is even number
        YPadTop = int((CropSize[1] - Crop.shape[0])/2)
        YPadBot = int((CropSize[1] - Crop.shape[0])/2)
    else:
        YPadTop = int( ((CropSize[1] - Crop.shape[0])/2)-0.5)
        YPadBot = int(((CropSize[1] - Crop.shape[0])/2)+0.5)
    ##Padding:
    if Crop.shape[1] %2 ==0:
        #Shape is even number
        XPadLeft = int((CropSize[0] - Crop.shape[1])/2)
        XPadRight= int((CropSize[0] - Crop.shape[1])/2)
    else:
        XPadLeft =  int(((CropSize[0] - Crop.shape[1])/2)-0.5)
        XPadRight= int(((CropSize[0] - Crop.shape[1])/2)+0.5)



    OutImage = cv2.copyMakeBorder(Crop, YPadTop,YPadBot,XPadLeft,XPadRight,cv2.BORDER_CONSTANT,value=[0,0,0])
    
    return OutImage,ScaleProportion, YPadTop,XPadLeft



def DLCInference(Crop,dlc_liveObj,CropSize):
    """Inference for DLC"""

    ###Scale crop if image bigger than cropsize
    # import ipdb;ipdb.set_trace()
    if Crop.shape[0] > CropSize[0] or Crop.shape[1] > CropSize[1]: #Image bigger than crop size, scale down
        ScaleRatio = min([CropSize[0]/Crop.shape[0], CropSize[1]/Crop.shape[1]])
        ScaleWidth = round(Crop.shape[1] * ScaleRatio)
        ScaleHeight = round(Crop.shape[0]*ScaleRatio)
        resizedCrop = cv2.resize(Crop, (ScaleWidth,ScaleHeight), interpolation=cv2.INTER_LINEAR)  # resize image
        ScaleUpRatio = 1/ScaleRatio #ratio to scale keypoints back up to original
        # import ipdb;ipdb.set_trace()
    else:
        resizedCrop = Crop
        ScaleUpRatio = 1
    # cv2.imwrite(filename="tempresize.jpg", img=resizedCrop)
    # cv2.imwrite(filename="temp.jpg", img=Crop)
    if dlc_liveObj.sess == None: #if first time, init
        DLCPredict2D = dlc_liveObj.init_inference(resizedCrop)

    DLCPredict2D= dlc_liveObj.get_pose(resizedCrop)
    DLCPredict2D[:,0] = DLCPredict2D[:,0]*ScaleUpRatio
    DLCPredict2D[:,1] = DLCPredict2D[:,1]*ScaleUpRatio

    return DLCPredict2D


def VisualizeAll(frame, box, DLCPredict2D,ScaleBBox, imsize):
    """Visualize all stuff"""
    colourList = [(0,255,255),(255,0 ,255),(128,0,128),(255,192,203),(255, 255, 0),(0, 0 , 255 ),(205,133,63),(0,255,0),(255,0,0)]

    ##Order: Lshoulder, Rshoulder, topKeel,botKeel,Tail,Beak,Nose,Leye,Reye
    ##Points:
    PlotPoints = []
    for x,point in enumerate(DLCPredict2D):
        roundPoint = [round(point[0]+box[0]),round(point[1]+box[1])]
        cv2.circle(frame,roundPoint,1,colourList[x], 5)
        PlotPoints.append(roundPoint)

    cv2.rectangle(frame,(round(box[0]),round(box[1])),(round(box[2]),round(box[3])),[0,0,255],3)


    return frame, PlotPoints


def Inference(frame,predictor,dlc_liveObj,ScaleBBox=1,Dilate=5,DLCThreshold=0.3):
    """Loop through video for SAM, save framewise info"""

    InferFrame = frame.copy()
    outputs = predictor(InferFrame)["instances"].to("cpu")
    CropSize = (320,320)

    # import ipdb;ipdb.set_trace()
    imsize = [frame.shape[1],frame.shape[0]]

    BirdIndex = np.where(outputs.pred_classes.numpy() == 14)[0] #14 is ID for bird
    BirdBBox = outputs.pred_boxes[BirdIndex].tensor.numpy()
    # import ipdb;ipdb.set_trace()
    BirdMasks = (outputs.pred_masks>0.95).numpy()[BirdIndex]

    for x in range(BirdBBox.shape[0]):
        # import ipdb;ipdb.set_trace()

        bbox = list(BirdBBox[x])
        Mask = BirdMasks[x]>0
        Mask = np.array(Mask,dtype=np.uint8)
        # show_anns(frame, Mask)

        if Dilate > 0:
            DilateKernel = np.ones((Dilate,Dilate),np.uint8)
            Mask = cv2.dilate(Mask,DilateKernel,iterations = 3) 

        # import ipdb;ipdb.set_trace()
        Mask = np.array(Mask,dtype=np.uint8)
        Mask = Mask.reshape(imsize[1],imsize[0],1)
        Crop = cv2.bitwise_and(InferFrame, InferFrame, mask=Mask)

        # cv2.imwrite(filename="temp.jpg", img = Crop)

        ##change box to XYWH to scale up
        bbox = [bbox[0],bbox[1],bbox[2]-bbox[0],bbox[3]-bbox[1]]
        ScaleWidth = ((ScaleBBox * bbox[2])/2)-(bbox[2]/2)
        ScaleHeight = ((ScaleBBox * bbox[3])/2)-(bbox[3]/2)
        # import ipdb;ipdb.set_trace()
        # BirdCrop = frame[round(bbox[1]):round(bbox[3]),round(bbox[0]):round(bbox[2])] #bbox is XYWH

        x1 = round(bbox[0]-ScaleWidth) if round(bbox[0]-ScaleWidth)>0 else 0
        y1 = round(bbox[1]-ScaleHeight)if round(bbox[1]-ScaleHeight)>0 else 0
        x2 = round(bbox[0]+bbox[2]+ScaleWidth) if round(bbox[0]+bbox[2]+ScaleWidth) < imsize[0] else imsize[0]
        y2 = round(bbox[1]+bbox[3]+ScaleHeight)if round(bbox[1]+bbox[3]+ScaleHeight) < imsize[1] else imsize[1]
        bbox = [x1,y1,x2,y2]
        BirdCrop = Crop[y1:y2,x1:x2] #bbox is XYWH

        DLCPredict2D= DLCInference(BirdCrop,dlc_liveObj,CropSize)

        frame, PlotPoints = VisualizeAll(frame, bbox, DLCPredict2D,ScaleBBox,imsize)

    if BirdBBox.shape[0] == 0:
        DLCPredict2D= DLCInference(InferFrame,dlc_liveObj,CropSize)
        bbox = [0,0,0,0]
        frame, PlotPoints = VisualizeAll(frame, bbox, DLCPredict2D,ScaleBBox,imsize)

    return frame