Upload 9 files

app.py

@@ -1,7 +1,18 @@
 import gradio as gr
 
+from mhr.predict_tools import MHRStoneRecognizeMgr, MHRVedioSimpleCuter
+
+mgr = MHRStoneRecognizeMgr(
+    whole_pkl = "./whole_model.pkl",
+    hole_pkl = "./hole_model.pkl",
+    vedio_cutter = MHRVedioSimpleCuter(),
+)
+
 def greet(v):
-    return "Hello !"
+    print(v)
+    ret = mgr.recognize(v)
+    print(ret)
+    return ret
 
 demo = gr.Interface(
     fn=greet,

hole_model.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a01490c907afe2d131f31c577f6bee4736ada5ae0c0f7de50ab71da7be2e42df
+size 98027

mhr/common.py

@@ -0,0 +1,28 @@
+import numpy as np # linear algebra
+import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)
+
+import os
+
+import random
+    
+# import 
+import torch
+import torch.nn as nn
+import torch.nn.functional as F 
+import torchvision
+import torchvision.transforms as T
+#from torch.utils.data import DataLoader, SubsetRandomSampler
+#from torch.utils.tensorboard import SummaryWriter
+import pandas as pd
+import numpy as np
+import cv2 as cv
+from PIL import Image,ImageDraw,ImageFont
+import matplotlib.pyplot as plt
+
+from sklearn import datasets
+from sklearn.cluster import KMeans
+#from sklearn.externals import joblib
+import pickle
+# other
+import tr
+

mhr/config.py

@@ -0,0 +1,63 @@
+from mhr.custom_transform import *
+from mhr.custom_net import *
+from mhr.custom_dateaset import *
+
+# vedio data
+data_dir = "/app/data/"
+
+# model
+whole_pkl_file = '/share_dir/mhr_data/whole_model.pkl'
+hole_pkl_file = '/share_dir/mhr_data/hole_model.pkl'
+speed_ratio = 5
+
+
+tsfm_whole = T.Compose([
+    T.CenterCrop([720,720]),
+    T.Resize([360,360]),
+    T.ToTensor(),
+    T.Grayscale(num_output_channels=1),
+])
+tsfm_whole4cv = T.Compose([
+    CV2PIL(),
+    T.CenterCrop([720,720]),
+    T.Resize([360,360]),
+    T.ToTensor(),
+    T.Grayscale(num_output_channels=1),
+])
+#ds = MyDataset(set_name='whole', root=data_dir, transform=tsfm_whole, no_label=True)
+#print("whole:", len(ds))
+
+
+tsfm_hole = T.Compose([
+    PIL2CV(),
+    CvtColor(cv.COLOR_BGR2GRAY),
+    Threshold(125,255,cv.THRESH_BINARY),
+    CV2PIL(),
+    T.Resize([28,28]),
+    T.ToTensor(),
+    T.Grayscale(num_output_channels=1),
+])
+tsfm_hole4cv = T.Compose([
+    CvtColor(cv.COLOR_BGR2GRAY),
+    Threshold(125,255,cv.THRESH_BINARY),
+    CV2PIL(),
+    T.Resize([28,28]),
+    T.ToTensor(),
+    T.Grayscale(num_output_channels=1),
+])
+#ds = MyDataset(set_name='hole', root=data_dir, transform=tsfm_hole, no_label=True)
+#print("hole:", len(ds))
+    
+
+tsfm_skill = T.Compose([
+    T.ToTensor(),
+    T.Grayscale(num_output_channels=1),
+])
+tsfm_skill4cv = T.Compose([
+    CV2PIL(),
+    T.ToTensor(),
+    T.Grayscale(num_output_channels=1),
+])
+#ds = MyDataset(set_name='skill', root=data_dir, transform=tsfm_skill, no_label=True)
+#print("skill:", len(ds))
+

mhr/custom_dateaset.py

@@ -0,0 +1,26 @@
+from mhr.common import *
+from mhr.custom_transform import *
+
+# define dataset
+class MyDataset(torch.utils.data.Dataset): #torch.utils.data.Dataset
+    def __init__(self, set_name, root='.', transform=lambda x: x, no_label=False):
+        super(MyDataset,self).__init__()
+        self._root = root
+        self._transform = transform
+        self._df = pd.read_csv(self._root+'/label/'+set_name+'.csv')
+        self._df = self._df
+        self._no_label = no_label
+
+    def __getitem__(self, index):
+        if self._no_label:
+            name, label = self._df.loc[index]
+            label = name
+        else:
+            name, label = self._df[~self._df['label'].isna()].loc[index]
+        return self._transform(Image.open(self._root+'/image/'+name)), label
+
+    def __len__(self): # return size
+        if self._no_label:
+            return len(self._df)
+        else:
+            return len(self._df[~self._df['label'].isna()])

mhr/custom_net.py

@@ -0,0 +1,184 @@
+from mhr.common import *
+from mhr.custom_transform import *
+
+def getGaborFilters(ksize, n_output, sigma_ratio_func, theta_ratio_func, lamda_ratio_func, gamma=0.5, psi=0, show=False):
+    filters = []
+    sigma = np.pi/2.0  # gaussian window width
+    theta = np.pi/2.0  # direction of cosine (raid)
+    lamda = np.pi/2.0  # wavelength of cosine
+    for i in range(n_output):
+        #print(i, sigma_ratio_func(i), theta_ratio_func(i), lamda_ratio_func(i))
+        kernel = cv.getGaborKernel((ksize,ksize), 
+                                   sigma*sigma_ratio_func(i),
+                                   theta*theta_ratio_func(i),
+                                   lamda*lamda_ratio_func(i), gamma, psi, ktype=cv.CV_32F)
+        filters.append(kernel)
+    return filters
+
+
+class TorchModelSaver:
+    def __init__(self):
+        pass
+    def save(self, model, path):
+        torch.save(model, path)
+    def load(self, path):
+        return torch.load(path)
+
+class SklearnModelSaver:
+    def __init__(self):
+        pass
+    def save(self, model, path):
+        pickle.dumps(model, path)
+    def load(self, path):
+        return pickle.loads(path)    
+
+# for whole
+class GaborFeatureNet(nn.Module):
+    def __init__(self, num_classes, show_filters=False, show_images=False):
+        super(GaborFeatureNet, self).__init__()
+        # config of gabor filters
+        ksize = 20
+        n_output = 12
+        sigma_func = lambda x: (x//4)/2+1
+        theta_func = lambda x: (x%4)/2
+        lamda_func = lambda x: x//4+1
+        filters = getGaborFilters(ksize, n_output, 
+                              sigma_func, theta_func, lamda_func
+                              )
+        self.conv1 = torch.nn.Conv2d(1, n_output, (ksize,ksize),stride=1, bias=False)
+        self.conv1.weight.data = torch.Tensor(filters).unsqueeze(1)
+        self.pool1 = nn.Sigmoid()
+        self.pool2 = nn.MaxPool2d(5)
+        self.pool3 = nn.MaxPool2d(2)
+        #set_parameter_requires_grad(self.features, True)#�̶�������ȡ�����
+        for p in self.parameters():
+            p.requires_grad = False
+        self.classifier = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(12*34*34 , 1024),
+            nn.ReLU(),
+            nn.Linear(1024, 1024),
+            nn.Dropout(0.5),
+            nn.ReLU(),
+            nn.Linear(1024, num_classes)
+        )
+        
+    def forward(self, img):
+        img = self.conv1(img)
+        img = self.pool1(img)
+        img = self.pool2(img) + self.pool2(-1*img)
+        img = self.pool3(img)
+        img = self.classifier(img)
+        return img
+
+class WholeModelMgr:
+    def __init__(self, num_classes):
+        self.model = GaborFeatureNet(num_classes)
+        self.saver = TorchModelSaver()
+    def save(self, path):
+        self.saver.save(self.model, path)
+    def load(self, path):
+        self.model = self.saver.load(path)
+
+
+# for hole
+class GaborFeatureGen(nn.Module):
+    def __init__(self, num_classes, show_filters=False, show_images=False):
+        super(GaborFeatureGen, self).__init__()
+        # config of gabor filters
+        ksize = 20
+        n_output = 12
+        sigma_func = lambda x: ((x//4)/2+2)/4
+        theta_func = lambda x: (x%4)/2
+        lamda_func = lambda x: (x//4+1)/2
+        self.show_filters = show_filters
+        self.show_images = show_images
+        self.wins, self.dirs = self._get_wins_dirs(n_output, theta_func)
+        filters = getGaborFilters(ksize, n_output, 
+                                sigma_func, theta_func, lamda_func, 
+                                psi=np.pi/2)
+        self.conv1 = torch.nn.Conv2d(1, n_output, (ksize,ksize),stride=1,padding='same', bias=False)
+        self.conv1.weight.data = torch.Tensor(filters).unsqueeze(1)
+        self.pool1 = nn.Sigmoid()
+        self.pool2 = nn.MaxPool2d(1)
+        self.pool3 = nn.MaxPool2d(2)
+        if show_filters:
+            self._show_img(self.wins, self.dirs, filters)
+        
+    def forward(self, img):
+        img = self.conv1(img)
+        img = self.pool1(img)
+        #img = self.pool2(img) + self.pool2(-1*img)
+        img = self.pool3(img)
+        if self.show_images:
+            self._show_img(self.wins, self.dirs, img[0])
+        return nn.Flatten()(img)
+    
+    def _show_img(self, wins,dirs,imgs):
+        plt.figure(1)
+        for i in range(len(imgs)):
+            plt.subplot(wins, dirs, i+1)
+            if type(imgs[i]) is np.ndarray:
+                plt.imshow(imgs[i], cmap=plt.get_cmap('gray'))
+            else:
+                plt.imshow(T.functional.to_pil_image(imgs[i]), cmap=plt.get_cmap('gray'))
+        plt.show()
+    
+    def _get_wins_dirs(self, n_output, theta_func):
+        dirs = len(set([ theta_func(i) for i in range(n_output) ]))
+        return n_output//dirs, dirs
+
+class HoleModelMgr:
+    def __init__(self, n_clusters):
+        self.feat_model = GaborFeatureGen(0)
+        self.model = KMeans(n_clusters)
+        self.saver = SklearnModelSaver()
+    def save(self, path):
+        self.saver.save(self.model, path)
+    def load(self, path):
+        self.model = self.saver.load(path)
+
+
+
+# for skill
+class MyTrRecognizeNet(torch.nn.Module):
+    def __init__(self, image_padding):
+        super(MyTrRecognizeNet, self).__init__()
+        ip = image_padding if image_padding is not None else 0
+        nm_p = [5, 22+16*0]
+        nm_h = 18 
+        nm_w = 16*10
+        lv_p = [29, 193]
+        lv_h = 18
+        lv_w = 44
+        #print([ x-ip for x in nm_p], nm_h+ip*2, nm_w+ip*2)
+        #print([ x-ip for x in lv_p], lv_h+ip*2, lv_w+ip*2)
+        self.tsfm_nm = T.Compose([
+            TensorCut([ x-ip for x in nm_p], nm_h+ip*2, nm_w+ip*2), # ����
+        ])
+        self.tsfm_lv = T.Compose([
+            TensorCut([ x-ip for x in lv_p], lv_h+ip*2, lv_w+ip*2), # ����
+        ])
+
+    def batch_forward(self, imgs):
+        return [ self.single_forward(img) for img in imgs ]        
+    def single_forward(self, img):
+        img_nm = self.tsfm_nm(img.clone().detach())
+        img_lv = self.tsfm_lv(img.clone().detach())
+        nm = tr.recognize(T.functional.to_pil_image(img_nm))
+        lv = tr.recognize(T.functional.to_pil_image(img_lv))
+        return nm,lv
+    def forward(self, img):
+        if len(img.shape) == 4:
+            return self.batch_forward(img)
+        elif len(img.shape) == 3:
+            return self.single_forward(img)
+    
+    def forward_bak(self, img):
+        if len(img.shape) == 4:
+            img = img.squeeze(0)
+        img_nm = self.tsfm_nm(img.clone().detach())
+        img_lv = self.tsfm_lv(img.clone().detach())
+        nm = tr.recognize(T.functional.to_pil_image(img_nm))
+        lv = tr.recognize(T.functional.to_pil_image(img_lv))
+        return nm,lv

mhr/custom_transform.py

@@ -0,0 +1,87 @@
+from mhr.common import *
+
+class PIL2CV(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+    def forward(self, img):
+        return cv.cvtColor(np.asarray(img), cv.COLOR_RGB2BGR)
+
+class CV2PIL(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+    def forward(self, img):
+        return Image.fromarray(cv.cvtColor(img, cv.COLOR_BGR2RGB))
+
+class Tensor2CV(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+    def forward(self, img):
+        if len(img.shape) == 4:
+            img = img.squeeze(0)
+        img = img.mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).type(torch.uint8).numpy()
+        img = cv.cvtColor(img, cv.COLOR_RGB2BGR)
+        return img
+
+
+class CvtColor(torch.nn.Module):
+    def __init__(self, cvt):
+        super().__init__()
+        self._cvt = cvt
+    def forward(self, img):
+        return cv.cvtColor(img, self._cvt)
+
+class GaussianBlur(torch.nn.Module):
+    def __init__(self, kernel, sigma):
+        super().__init__()
+        self._kernel = kernel
+        self._sigma = sigma
+    def forward(self, img):
+        return cv.GaussianBlur(img, self._kernel, self._sigma)
+
+class EqualizeHist(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+    def forward(self, img):
+        return cv.equalizeHist(img)
+
+class SobelBitwiseOrXY(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+    def forward(self, img):
+        sobelx = cv.Sobel(img, cv.CV_64F, 1, 0)
+        sobely = cv.Sobel(img, cv.CV_64F, 0, 1)
+        sobelx = cv.convertScaleAbs(sobelx)
+        sobely = cv.convertScaleAbs(sobely)
+        return cv.bitwise_or(sobelx, sobely)
+
+class Threshold(torch.nn.Module):
+    def __init__(self, thresh, maxval, tt):
+        super().__init__()
+        self._thresh = thresh
+        self._maxval = maxval
+        self._tt = tt
+    def forward(self, img):
+        _, tmp = cv.threshold(img, self._thresh, self._maxval, self._tt)
+        return tmp
+
+class Cut(torch.nn.Module):
+    def __init__(self, point, offsetx, offsety):
+        super().__init__()
+        self._p = point
+        self._offsetx = offsetx
+        self._offsety = offsety
+    def forward(self, img):
+        p = self._p
+        return img[p[0]:p[0]+self._offsetx, p[1]:p[1]+self._offsety ]
+
+
+
+class TensorCut(torch.nn.Module):
+    def __init__(self, point, offsetx, offsety):
+        super().__init__()
+        self._p = point
+        self._offsetx = offsetx
+        self._offsety = offsety
+    def forward(self, img):
+        p = self._p
+        return img[:, p[0]:p[0]+self._offsetx, p[1]:p[1]+self._offsety ]

mhr/predict_tools.py

@@ -0,0 +1,127 @@
+from mhr.config import *
+
+
+class MHRVedioCuter:
+    def __init__(self, speed_ratio=1):
+        ## config
+        #self.part_pos = {
+        #    'pos':[(628,350),(993,565)],
+        #    'page': [(781,575),(848,600)],
+        #    'hole': [(1167,197),(1250,227)],
+        #    'skill': [(1010,260),(1254,600)],
+        #}
+        self.part_pos = {
+            'hole': [1166,200,28,26],
+            'skill': [1014,264,240,50],
+        }
+        self.speed_ratio = speed_ratio
+
+
+    def iter(self, name):
+        vc = cv.VideoCapture(name)
+        fps = vc.get(cv.CAP_PROP_FPS)
+        print("vedio:", vc.isOpened(), fps)
+        label = "00:00:{:05.2f}({})"
+        rval = True
+        idx=0
+        while rval: 
+            rval, img = vc.read()
+            idx+=1
+            if rval and idx%self.speed_ratio == 0:
+                yield self._cut_whole(img), self._cut_hole(img), self._cut_skill(img), label.format(idx/fps, idx)
+
+        vc.release()
+
+    def _cut_whole(self, img):
+        #pos_w, pos_h, w, h = self.part_pos['skill']
+        return img
+
+    def _cut_hole(self, img):
+        pos_w, pos_h, w, h = self.part_pos['hole']
+        return [ img[pos_h:pos_h+h, pos_w+w*i:pos_w+w*i+w] for i in range(3) ]
+
+    def _cut_skill(self, img):
+        pos_w, pos_h, w, h = self.part_pos['skill']
+        return [ img[pos_h+h*i:pos_h+h*i+h, pos_w:pos_w+w] for i in range(7) ]
+
+class MHRVedioSimpleCuter(MHRVedioCuter):
+    def __init__(self):
+        super(MHRVedioSimpleCuter, self).__init__()
+
+    def iter(self, v):
+        for img in v:
+            yield self._cut_whole(img), self._cut_hole(img), self._cut_skill(img), label.format(idx/fps, idx)
+
+
+
+class MHRStoneRecognizeMgr:
+    def __init__(self, whole_pkl, hole_pkl, vedio_cutter):
+        self.mapping_hole = [0,2,1,3,4]
+        cp = torch.load(whole_pkl)
+        self.whole_model = GaborFeatureNet(num_classes=2)
+        self.whole_model.load_state_dict(cp['model'])
+        if torch.cuda.is_available():
+            self.whole_model = self.whole_model.cuda()
+        #self.whole_model = torch.load(whole_pkl)
+        self.hole_feat_model = GaborFeatureGen(0)
+        if torch.cuda.is_available():
+            self.hole_feat_model = self.hole_feat_model.cuda()
+        self.hole_model = pickle.load(open(hole_pkl, 'rb'))
+        self.skill_model = MyTrRecognizeNet(image_padding=2)
+        self._vedio_cutter = vedio_cutter
+
+    def recognize(self, vname, fname=None):
+        def dump(rr):
+            return '_'.join([ str(x) for x in rr[1] ]) + "|" + '|'.join([ x[0]+":"+x[1] for x in rr[2] ])
+        results = []
+        i=0
+        for data in self._vedio_cutter.iter(vname):
+            i+=1
+            #if i < 200 and i > 250: 
+            #    continue
+            #if len(results) == 1:
+            #    break
+            #whole
+            data_whole = tsfm_whole4cv(data[0])
+            data_whole = data_whole.unsqueeze(0)
+            if torch.cuda.is_available():
+                data_whole = data_whole.cuda()
+            ret = self.whole_model(data_whole)
+            if ret[0][1] - ret[0][0] < 2:
+                continue
+            ret_whole = True
+            #new hole
+            data_hole = torch.cat([ tsfm_hole4cv(item).unsqueeze(0) for item in data[1] ], dim=0)
+            if torch.cuda.is_available():
+                data_hole = data_hole.cuda()
+            output = self.hole_feat_model(data_hole)
+            df = pd.DataFrame(output.tolist())
+            ret_hole = list(self.hole_model.predict(df))
+            #new skill
+            data_skill = torch.cat([ tsfm_skill4cv(item).unsqueeze(0) for item in data[2] ], dim=0)
+            ret = self.skill_model(data_skill)
+            ret_skill = [ (x[0][0], x[1][0][-1]) for x in filter(lambda sk: sk[0][1] > 0.9 and sk[1][1] > 0.9, ret) ]
+            # reuslt
+            result = [data[3], ret_hole, ret_skill]
+            if len(results) > 0 and dump(results[-1]) == dump(result):
+                continue
+            results.append(result)
+        if fname:
+            with open(fname, 'w') as f:
+                for result in results:
+                    line = result[0]
+                    line += ','
+                    line += ','.join([ str(x) for x in result[1] ])
+                    line += ','
+                    line += ','.join([ x[0]+","+x[1] for x in result[2] ])
+                    line += '\n'
+                    f.write(line)
+        return results
+
+
+
+mgr = MHRStoneRecognizeMgr(
+    whole_pkl = whole_pkl_file,
+    hole_pkl = hole_pkl_file,
+    vedio_cutter = MHRVedioCuter(speed_ratio),
+)

whole_model.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:267011c5b8bb860dd61a4b3e769582d5f69296f9432a51b32b93ea8b168ccc72
+size 61054629