123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758 |
- from pathlib import Path
- import matplotlib.path as mat
- from utils.general import strtolstl
- from utils.general import compute_IOU
- from torchvision import transforms
- from PIL import Image
- import torch
- import torch.nn.functional as F
- import numpy as np
- from ultralytics.engine.results import Results
- mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
- test = transforms.Compose([transforms.Resize((224,224)),
- #transforms.CenterCrop(224),
- transforms.ToTensor(),
- transforms.Normalize(mean=mean, std=std)
- ])
- def clapre(modelcla,claimg,clapoint):
- imgten = torch.stack(claimg,dim=0)
- clapoint = torch.stack(clapoint,dim=0)
- imgten = imgten.to(0)
- result = modelcla(imgten)
- result = F.softmax(result)
- print(result)
- index = result.argmax(1)
- index = index.cpu().numpy()
- index = np.argwhere(index<5)
- index = index.reshape(-1)
- print(index)
- if len(index)>0:
- print(clapoint[index])
- return clapoint[index]
- else:
- return None
- class Helmet:
- def __init__(self):
- self.flag = False
- def getflag(self,det,persondet,annotator,fence=0,point=None,names=None,rname=None,num=1):
- #print(type(det))
- self.flag = False
- for *xyxy, conf, cls in reversed(det):
- c = int(cls)
- labelname = names[c]
- if labelname in rname:
- if fence == 1:
- pointa = strtolstl(point)
- for poi in pointa:
- p1 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3),
- xyxy[3].cpu().item())
- p2 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3 * 2),
- xyxy[3].cpu().item())
- pt = [p1, p2]
- inflag = mat.Path(poi).contains_points(pt)
- if inflag.any():
- if persondet is None:
- self.flag = True
- # c = int(cls) # integer class
- # label = f'{self.names[c]} {conf:.2f}'
- label = None
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- else:
- for person in persondet:
- personinflag = mat.Path(person).contains_points(pt)
- if personinflag.any():
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- self.flag = True
- else:
- if persondet is None:
- self.flag = True
- # c = int(cls) # integer class
- # label = f'{self.names[c]} {conf:.2f}'
- label = None
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- else:
- for person in persondet:
- p1 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3),
- xyxy[3].cpu().item())
- p2 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3 * 2),
- xyxy[3].cpu().item())
- pt = [p1, p2]
- personinflag = mat.Path(person).contains_points(pt)
- if personinflag.any():
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- self.flag = True
- return self.flag
- class Uniform:
- def __init__(self):
- self.flag = False
- def getflag(self, det, persondet,annotator, fence=0, point=None, names=None, rname=None,num=1):
- self.flag = False
- for *xyxy, conf, cls in reversed(det):
- c = int(cls)
- labelname = names[c]
- if labelname in rname:
- if fence == 1:
- pointa = strtolstl(point)
- for poi in pointa:
- p1 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3),
- xyxy[3].cpu().item())
- p2 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3 * 2),
- xyxy[3].cpu().item())
- pt = [p1, p2]
- inflag = mat.Path(poi).contains_points(pt)
- if inflag.any():
- if persondet is None:
- self.flag = True
- # c = int(cls) # integer class
- # label = f'{self.names[c]} {conf:.2f}'
- label = None
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- else:
- for person in persondet:
- p1 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3),
- xyxy[3].cpu().item())
- p2 = (
- int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3 * 2),
- xyxy[3].cpu().item())
- pt = [p1, p2]
- personinflag = mat.Path(person).contains_points(pt)
- if personinflag.any():
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- self.flag = True
- else:
- if persondet is None:
- self.flag = True
- # c = int(cls) # integer class
- # label = f'{self.names[c]} {conf:.2f}'
- label = None
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- else:
- for person in persondet:
- p1 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3),
- xyxy[3].cpu().item())
- p2 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3 * 2),
- xyxy[3].cpu().item())
- pt = [p1, p2]
- personinflag = mat.Path(person).contains_points(pt)
- if personinflag.any():
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- self.flag = True
- return self.flag
- class Fall:
- def __init__(self):
- self.flag = False
- def getflag(self, det, persondet,annotator, fence=0, point=None, names=None, rname=None,num=1):
- self.flag = False
- for *xyxy, conf, cls in reversed(det):
- c = int(cls)
- labelname = names[c]
- if labelname in rname:
- if fence == 1:
- pointa = strtolstl(point)
- for poi in pointa:
- p1 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3),
- xyxy[3].cpu().item())
- p2 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3 * 2),
- xyxy[3].cpu().item())
- pt = [p1, p2]
- inflag = mat.Path(poi).contains_points(pt)
- if inflag.any():
- if persondet is None:
- self.flag = True
- # c = int(cls) # integer class
- # label = f'{self.names[c]} {conf:.2f}'
- label = None
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- else:
- for person in persondet:
- personinflag = mat.Path(person).contains_points(pt)
- if personinflag.any():
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- self.flag = True
- else:
- if persondet is None:
- self.flag = True
- # c = int(cls) # integer class
- # label = f'{self.names[c]} {conf:.2f}'
- label = None
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- else:
- for person in persondet:
- p1 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3),
- xyxy[3].cpu().item())
- p2 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3 * 2),
- xyxy[3].cpu().item())
- pt = [p1, p2]
- personinflag = mat.Path(person).contains_points(pt)
- if personinflag.any():
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- self.flag = True
- return self.flag
- class Personcount:
- def __init__(self):
- self.flag = False
- def getflag(self, det,persondet, annotator, fence=0, point=None, names=None, rname=None,num=1):
- self.flag = False
- detnum = 0
- for *xyxy, conf, cls in reversed(det):
- c = int(cls)
- labelname = names[c]
- if labelname in rname:
- if fence == 1:
- pointa = strtolstl(point)
- for poi in pointa:
- p1 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3),
- xyxy[3].cpu().item())
- p2 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3 * 2),
- xyxy[3].cpu().item())
- pt = [p1, p2]
- inflag = mat.Path(poi).contains_points(pt)
- if inflag.any():
- if persondet is None:
- #self.flag = True
- detnum = detnum+1
- # c = int(cls) # integer class
- # label = f'{self.names[c]} {conf:.2f}'
- label = None
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- else:
- for person in persondet:
- personinflag = mat.Path(person).contains_points(pt)
- if personinflag.any():
- detnum = detnum+1
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- #self.flag = True
- else:
- if persondet is None:
- #self.flag = True
- detnum = detnum+1
- # c = int(cls) # integer class
- # label = f'{self.names[c]} {conf:.2f}'
- label = None
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- else:
- for person in persondet:
- p1 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3),
- xyxy[3].cpu().item())
- p2 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3 * 2),
- xyxy[3].cpu().item())
- pt = [p1, p2]
- personinflag = mat.Path(person).contains_points(pt)
- if personinflag.any():
- detnum = detnum+1
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- #self.flag = True
- if detnum >= num:
- self.flag = True
- return self.flag
-
- class Arm:
- def __init__(self):
- self.flag = False
- def getflag(self, det, persondet,annotator, fence=0, point=None, names=None, rname=None,num=1):
- self.flag = False
- for *xyxy, conf, cls in reversed(det):
- c = int(cls)
- labelname = names[c]
- if labelname in rname:
- if fence == 1:
- pointa = strtolstl(point)
- for poi in pointa:
- p1 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3),
- xyxy[3].cpu().item())
- p2 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3 * 2),
- xyxy[3].cpu().item())
- pt = [p1, p2]
- inflag = mat.Path(poi).contains_points(pt)
- if inflag.any():
- if persondet is None:
- self.flag = True
- # c = int(cls) # integer class
- # label = f'{self.names[c]} {conf:.2f}'
- label = None
- annotator.box_label(xyxy, None, color=(0,0,255))
- else:
- for person in persondet:
- personinflag = mat.Path(person).contains_points(pt)
- if personinflag.any():
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- self.flag = True
- else:
- if persondet is None:
- self.flag = True
- # c = int(cls) # integer class
- # label = f'{self.names[c]} {conf:.2f}'
- label = None
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- else:
- for person in persondet:
- p1 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3),
- xyxy[3].cpu().item())
- p2 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3 * 2),
- xyxy[3].cpu().item())
- pt = [p1, p2]
- personinflag = mat.Path(person).contains_points(pt)
- if personinflag.any():
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- self.flag = True
- return self.flag
-
- class Bag:
- def __init__(self):
- self.flag = False
- def getflag(self, det, persondet,annotator, fence=0, point=None, names=None, rname=None,num=1):
- self.flag = False
- for *xyxy, conf, cls in reversed(det):
- c = int(cls)
- labelname = names[c]
- if labelname in rname:
- if fence == 1:
- pointa = strtolstl(point)
- for poi in pointa:
- p1 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3),
- xyxy[3].cpu().item())
- p2 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3 * 2),
- xyxy[3].cpu().item())
- pt = [p1, p2]
- inflag = mat.Path(poi).contains_points(pt)
- if inflag.any():
- if persondet is None:
- self.flag = True
- # c = int(cls) # integer class
- # label = f'{self.names[c]} {conf:.2f}'
- label = None
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- else:
- for person in persondet:
- personinflag = mat.Path(person).contains_points(pt)
- if personinflag.any():
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- self.flag = True
- else:
- if persondet is None:
- self.flag = True
- # c = int(cls) # integer class
- # label = f'{self.names[c]} {conf:.2f}'
- label = None
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- else:
- for person in persondet:
- p1 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3),
- xyxy[3].cpu().item())
- p2 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3 * 2),
- xyxy[3].cpu().item())
- pt = [p1, p2]
- personinflag = mat.Path(person).contains_points(pt)
- if personinflag.any():
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- self.flag = True
- return self.flag
-
- class Cross:
- def __init__(self):
- self.flag = False
- def getflag(self, det, persondet,annotator, fence=0, point=None, names=None, rname=None,num=1):
- self.flag = False
- detnum = 0
- for *xyxy, conf, cls in reversed(det):
- c = int(cls)
- labelname = names[c]
- if labelname in rname:
- if fence == 1:
- pointa = strtolstl(point)
- for poi in pointa:
- p1 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3),
- xyxy[3].cpu().item())
- p2 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3 * 2),
- xyxy[3].cpu().item())
- pt = [p1, p2]
- inflag = mat.Path(poi).contains_points(pt)
- if inflag.any():
- if persondet is None:
- detnum = detnum+1
- self.flag = True
- # c = int(cls) # integer class
- # label = f'{self.names[c]} {conf:.2f}'
- label = None
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- else:
- for person in persondet:
- personinflag = mat.Path(person).contains_points(pt)
- if personinflag.any():
- detnum = detnum+1
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- self.flag = True
- else:
- if persondet is None:
- self.flag = True
- # c = int(cls) # integer class
- # label = f'{self.names[c]} {conf:.2f}'
- detnum = detnum+1
- label = None
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- else:
- for person in persondet:
- p1 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3),
- xyxy[3].cpu().item())
- p2 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3 * 2),
- xyxy[3].cpu().item())
- pt = [p1, p2]
- personinflag = mat.Path(person).contains_points(pt)
- if personinflag.any():
- detnum = detnum +1
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- self.flag = True
- return self.flag
- class Extinguisher:
- def __init__(self):
- self.flag = False
- def getflag(self, det, persondet,annotator, fence=0, point=None, names=None, rname=None,num=1):
- self.flag = False
- for *xyxy, conf, cls in reversed(det):
- c = int(cls)
- labelname = names[c]
- if labelname in rname:
- if fence == 1:
- pointa = strtolstl(point)
- for poi in pointa:
- p1 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3),
- xyxy[3].cpu().item())
- p2 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3 * 2),
- xyxy[3].cpu().item())
- pt = [p1, p2]
- inflag = mat.Path(poi).contains_points(pt)
- if inflag.any():
- if persondet is None:
- self.flag = True
- # c = int(cls) # integer class
- # label = f'{self.names[c]} {conf:.2f}'
- label = None
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- else:
- for person in persondet:
- personinflag = mat.Path(person).contains_points(pt)
- if personinflag.any():
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- self.flag = True
- else:
- if persondet is None:
- self.flag = True
- # c = int(cls) # integer class
- # label = f'{self.names[c]} {conf:.2f}'
- label = None
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- else:
- for person in persondet:
- p1 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3),
- xyxy[3].cpu().item())
- p2 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3 * 2),
- xyxy[3].cpu().item())
- pt = [p1, p2]
- personinflag = mat.Path(person).contains_points(pt)
- if personinflag.any():
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- self.flag = True
- return self.flag
- class Persontre:
- def __init__(self):
- self.flag = False
- self.classifier_model = torch.load('/home/h3c/yolo/persontrecls.pt')
- def getflag(self, det, persondet,annotator, fence=0, point=None, names=None, rname=None,num=1,im0=None):
- self.flag = False
- dirp = {}
- dirf = {}
- for *xyxy, conf, cls in reversed(det):
- c = int(cls)
- labelname = names[c]
- if labelname in rname:
- if fence == 1:
- pointa = strtolstl(point)
- for poi in pointa:
- p1 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3),
- xyxy[3].cpu().item())
- p2 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3 * 2),
- xyxy[3].cpu().item())
- pt = [p1, p2]
- inflag = mat.Path(poi).contains_points(pt)
- if inflag.any():
- if persondet is None:
- self.flag = True
- # c = int(cls) # integer class
- # label = f'{self.names[c]} {conf:.2f}'
- label = None
- if c==0:
- dirp.setdefault(0,[])
- dirp[0].append(xyxy)
- elif c in [1,2]:
- dirp.setdefault(1,[])
- dirp[1].append(xyxy)
- dirf.setdefault(1,[])
- dirf[1].append(xyxy)
- elif c==3:
- dirp.setdefault(1,[])
- dirp[1].append(xyxy)
- elif c==4:
- dirf.setdefault(0,[])
- dirf[0].append(xyxy)
- #annotator.box_label(xyxy, None, color=(0, 0, 255))
- else:
- for person in persondet:
- personinflag = mat.Path(person).contains_points(pt)
- if personinflag.any():
- #annotator.box_label(xyxy, None, color=(0, 0, 255))
- #self.flag = True
- if c==0:
- dirp.setdefault(0,[])
- dirp[0].append(xyxy)
- elif c in [1,2]:
- dirp.setdefault(1,[])
- dirp[1].append(xyxy)
- dirf.setdefault(1,[])
- dirf[1].append(xyxy)
- elif c==3:
- dirp.setdefault(1,[])
- dirp[1].append(xyxy)
- elif c==4:
- dirf.setdefault(0,[])
- dirf[0].append(xyxy)
- else:
- if persondet is None:
- #self.flag = True
- # c = int(cls) # integer class
- # label = f'{self.names[c]} {conf:.2f}'
- #label = None
- #annotator.box_label(xyxy, None, color=(0, 0, 255))
- if c==0:
- dirp.setdefault(0,[])
- dirp[0].append(xyxy)
- elif c in [1,2]:
- dirp.setdefault(1,[])
- dirp[1].append(xyxy)
- dirf.setdefault(1,[])
- dirf[1].append(xyxy)
- elif c==3:
- dirp.setdefault(1,[])
- dirp[1].append(xyxy)
- elif c==4:
- dirf.setdefault(0,[])
- dirf[0].append(xyxy)
- else:
- for person in persondet:
- p1 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3),
- xyxy[3].cpu().item())
- p2 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3 * 2),
- xyxy[3].cpu().item())
- pt = [p1, p2]
- personinflag = mat.Path(person).contains_points(pt)
- if personinflag.any():
- #annotator.box_label(xyxy, None, color=(0, 0, 255))
- #self.flag = True
- if c==0:
- dirp.setdefault(0,[])
- dirp[0].append(xyxy)
- elif c in [1,2]:
- dirp.setdefault(1,[])
- dirp[1].append(xyxy)
- dirf.setdefault(1,[])
- dirf[1].append(xyxy)
- elif c==3:
- dirp.setdefault(1,[])
- dirp[1].append(xyxy)
- elif c==4:
- dirf.setdefault(0,[])
- dirf[0].append(xyxy)
- if len(dirp.keys()) == 2:
- claimg = []
- clapoint = []
- for person in dirp[0]:
- for other in dirp[1]:
- iou, newxyxy = compute_IOU(person, other)
- if iou>0.1:
- print(newxyxy)
- imgtmp = im0[int(newxyxy[1]):int(newxyxy[3]),int(newxyxy[0]):int(newxyxy[2])]
- imgtmp = imgtmp[...,::-1]
- imgtmp = Image.fromarray(imgtmp)
- imgten1 = test(imgtmp)
- claimg.append(imgten1)
- clapoint.append((newxyxy))
- result = clapre(self.classifier_model,claimg,clapoint)
- #imgten = imgten1[None]
- #imgten = imgten.to(0)
- #result = modelcla(imgten)
- #result = F.softmax(result, dim=1)
- #cla = result.argmax(1)
- if result is not None:
- self.flag = True
- for res in result:
- print(res)
- annotator.box_label(res, None, color=(0,0,255))
- if len(dirp.keys()) == 2:
- claimg = []
- clapoint = []
- for person in dirp[0]:
- for other in dirp[1]:
- iou, newxyxy = compute_IOU(person, other)
- if iou>0.1:
- print(newxyxy)
- imgtmp = im0[int(newxyxy[1]):int(newxyxy[3]),int(newxyxy[0]):int(newxyxy[2])]
- imgtmp = imgtmp[...,::-1]
- imgtmp = Image.fromarray(imgtmp)
- imgten1 = test(imgtmp)
- claimg.append(imgten1)
- clapoint.append((newxyxy))
- result = clapre(self.classifier_model,claimg,clapoint)
- #imgten = imgten1[None]
- #imgten = imgten.to(0)
- #result = modelcla(imgten)
- #result = F.softmax(result, dim=1)
- #cla = result.argmax(1)
- if result is not None:
- self.flag = True
- for res in result:
- print(res)
- annotator.box_label(res, None, color=(0,0,255))
- return self.flag
- class Danager:
- def __init__(self):
- self.flag = False
- def getflag(self, det, persondet,annotator, fence=0, point=None, names=None, rname=None,num=1):
- self.flag = False
- for *xyxy, conf, cls in reversed(det):
- c = int(cls)
- labelname = names[c]
- if labelname in rname:
- if fence == 1:
- pointa = strtolstl(point)
- for poi in pointa:
- p1 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3),
- xyxy[3].cpu().item())
- p2 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3 * 2),
- xyxy[3].cpu().item())
- pt = [p1, p2]
- inflag = mat.Path(poi).contains_points(pt)
- if inflag.any():
- if persondet is None:
- self.flag = True
- # c = int(cls) # integer class
- # label = f'{self.names[c]} {conf:.2f}'
- label = None
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- else:
- for person in persondet:
- personinflag = mat.Path(person).contains_points(pt)
- if personinflag.any():
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- self.flag = True
- else:
- if persondet is None:
- self.flag = True
- # c = int(cls) # integer class
- # label = f'{self.names[c]} {conf:.2f}'
- label = None
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- else:
- for person in persondet:
- p1 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3),
- xyxy[3].cpu().item())
- p2 = (int(xyxy[0].cpu().item() + (xyxy[2].cpu().item() - xyxy[0].cpu().item()) / 3 * 2),
- xyxy[3].cpu().item())
- pt = [p1, p2]
- personinflag = mat.Path(person).contains_points(pt)
- if personinflag.any():
- annotator.box_label(xyxy, None, color=(0, 0, 255))
- self.flag = True
- return self.flag
-
- class CarHelmetBelt:
- def __init__(self):
- self.flag = False
- def selectNoBeltPerson(self, person_objs, belt_objs):
- objs = []
- polygon_person = [Polygon(
- [(left, top), (right, top), (right, bottom), (left, bottom)]) for left, top, right, bottom, _, _ in person_objs]
- polygon_belt = [Polygon(
- [(left, top), (right, top), (right, bottom), (left, bottom)]) for left, top, right, bottom, _, _ in belt_objs]
- for person_obj, person in zip(person_objs, polygon_person):
- with_belt = False
- for belt in polygon_belt:
- if person.intersection(belt).area / belt.area > 0.5:
- with_belt = True
- break
- if not with_belt:
- objs.append(person_obj)
- return objs
- def selectWithPersonHead(self, person_objs, head_objs):
- objs = []
- polygon_person = [Polygon(
- [(left, top), (right, top), (right, top + (bottom - top)/2), (left, top + (bottom - top)/2)]) for left, top, right, bottom, _, _ in person_objs]
- polygon_head = [Polygon(
- [(left, top), (right, top), (right, bottom), (left, bottom)]) for left, top, right, bottom, _, _ in head_objs]
- for head_obj, head in zip(head_objs, polygon_head):
- with_person = False
- for person in polygon_person:
- if person.intersection(head).area / head.area > 0.5:
- with_person = True
- break
- if with_person:
- objs.append(head_obj)
- return objs
- def getflag(self, det, persondet,annotator, fence=0, point=None, names=None, rname=None,num=1):
- self.flag = False
- results = Results(annotator.result(),path=None,names=names,boxes=det)
- person_objs = []
- head_objs = []
- belt_objs = []
- polygon_area = strtolstl(point)[0]
- for result in results:
- boxes = result.boxes
- for box in boxes:
- left, top, right, bottom = box.xyxy.cpu().numpy().tolist()[0]
- polygon_box = Polygon([(left, top), (right, top), (right, bottom), (left, bottom)])
- intersection_areas = [polygon_area.intersection(polygon_box).area for polygon_area in self.polygon_areas ]
- if sum(intersection_areas) == 0:
- continue
- intersection_areas_ratio = sorted([intersection_area / polygon_box.area for intersection_area in intersection_areas])
- if intersection_areas_ratio[-1] < 0.9:
- continue
- conf = box.conf.cpu().numpy().tolist()[0]
- cls = box.cls.cpu().numpy().tolist()[0]
- if cls == 0:
- person_objs.append([left, top, right, bottom, conf, self.names[cls]])
- elif cls == 1:
- head_objs.append([left, top, right, bottom, conf, self.names[cls]])
- elif cls == 3:
- belt_objs.append([left, top, right, bottom, conf, self.names[cls]])
- illegal_objs = self.selectNoBeltPerson(person_objs, belt_objs) + self.selectWithPersonHead(person_objs, head_objs)
- if illegal_objs>0:
- for obj in illegal_objs:
- annotator.box_label(obj[:4], None, color=(0, 0, 255))
- self.flag = True
|