rename class autoShape -> AutoShape (#3173)
* rename class autoShape -> AutoShape follow other class naming convention * rename class autoShape -> AutoShape follow other classes' naming convention * rename class autoShape -> AutoShape
This commit is contained in:
1 parent
17b0f71538
commit
be86c21c73
2 files changed
+6
-6
No files matched your search
+3
-3
@@ -223,18 +223,18 @@ class NMS(nn.Module):
|
|||||||
return non_max_suppression(x[0], conf_thres=self.conf, iou_thres=self.iou, classes=self.classes)
|
return non_max_suppression(x[0], conf_thres=self.conf, iou_thres=self.iou, classes=self.classes)
|
||||||
|
|
||||||
|
|
||||||
class autoShape(nn.Module):
|
class AutoShape(nn.Module):
|
||||||
# input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS
|
# input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS
|
||||||
conf = 0.25 # NMS confidence threshold
|
conf = 0.25 # NMS confidence threshold
|
||||||
iou = 0.45 # NMS IoU threshold
|
iou = 0.45 # NMS IoU threshold
|
||||||
classes = None # (optional list) filter by class
|
classes = None # (optional list) filter by class
|
||||||
|
|
||||||
def __init__(self, model):
|
def __init__(self, model):
|
||||||
super(autoShape, self).__init__()
|
super(AutoShape, self).__init__()
|
||||||
self.model = model.eval()
|
self.model = model.eval()
|
||||||
|
|
||||||
def autoshape(self):
|
def autoshape(self):
|
||||||
print('autoShape already enabled, skipping... ') # model already converted to model.autoshape()
|
print('AutoShape already enabled, skipping... ') # model already converted to model.autoshape()
|
||||||
return self
|
return self
|
||||||
|
|
||||||
@torch.no_grad()
|
@torch.no_grad()
|
||||||
|
|||||||
+3
-3
@@ -215,9 +215,9 @@ class Model(nn.Module):
|
|||||||
self.model = self.model[:-1] # remove
|
self.model = self.model[:-1] # remove
|
||||||
return self
|
return self
|
||||||
|
|
||||||
def autoshape(self): # add autoShape module
|
def autoshape(self): # add AutoShape module
|
||||||
logger.info('Adding autoShape... ')
|
logger.info('Adding AutoShape... ')
|
||||||
m = autoShape(self) # wrap model
|
m = AutoShape(self) # wrap model
|
||||||
copy_attr(m, self, include=('yaml', 'nc', 'hyp', 'names', 'stride'), exclude=()) # copy attributes
|
copy_attr(m, self, include=('yaml', 'nc', 'hyp', 'names', 'stride'), exclude=()) # copy attributes
|
||||||
return m
|
return m
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user