import cv2
from torchvision.models import resnet
import torch
from torch import nn
import os
import numpy as np
import torchvision
defvideo_capture(pattern='array'):capture = cv2.VideoCapture(0)_, frame = capture.read()if pattern=='array':return np.array(frame)else:return frame
defsimeple_detron(dir='detectron_system',pattern='train',label=1,pre_dict=None,opeartion=None,retrain=False):''':describe:this function is used to control function in 'operation'.basic principle is:label=CONV(your image)action=operation(label):param:dir: weigh saved pathpattern: record your info, or detect who you arelabel: record labelpre_dict: {label:name}operation: your own system would like to be controlretrain: train from head with out loading the weights before'''assert pattern in['train','val'],print('only support train or val')model = resnet.resnet34(pretrained=True)ifnot os.path.exists(dir):os.mkdir(dir)if pattern=='train':del(model.fc)model.fc=nn.Linear(in_features=512,out_features=10)ifnot retrain:model.load_state_dict(torch.load(dir+'/'+'person.pth'))while(1):image=video_capture(pattern='image')cv2.imshow('photo_shop',image)if cv2.waitKey(30)==ord('q'):breakimage=torchvision.transforms.Compose([torchvision.transforms.ToTensor(),torchvision.transforms.Resize((224,224))])(image)optim=torch.optim.Adam(lr=0.001,params=model.parameters())for i inrange(100):pre=model(image.unsqueeze(0).repeat(2,1,1,1))loss=nn.CrossEntropyLoss()(pre,torch.Tensor([label,label]).view(2,1).squeeze(-1).long())optim.zero_grad()loss.backward()optim.step()torch.save(model.state_dict(),dir+'/'+'person.pth')if pattern=='val':del(model.fc)model.fc = nn.Linear(in_features=512, out_features=10)model.load_state_dict(torch.load(dir+'/'+'person.pth'))model=model.eval()while(1):image = video_capture(pattern='image')cv2.imshow('photo_shop', image)if cv2.waitKey(30)==ord('q'):breakimage=torchvision.transforms.Compose([torchvision.transforms.ToTensor(),torchvision.transforms.Resize((224,224))])(image)pre=model(image.unsqueeze(0))if pre_dict:pre=pre_dict[pre.argmax(dim=-1).cpu().detach().item()]print(f'this person is {pre}')if opeartion:opeartion(pre)if __name__=='__main__':print('*************This is an simple detection system**************')print('*************you can assign a function to this system so that make this system become a real cotrol system')print('*************for example: simeple_detron(opeartion= your_function)')print('lets start!')while(1):pattern=input('train is to map your shape to your name, val is to check who you are: ')if pattern=='train':label =input('0~10: used for registration ')ifeval(label)<0oreval(label)>10:continueelse:label=int(eval(label))else:label=Nonesure=bool(eval(input('are you sure? 0 or 1 ')))if sure:breakdict_ ={1:'my friend'}simeple_detron(pattern=pattern,label=label,pre_dict=dict_)