def _predict_batch(
model: Union[DetBenchTrain, DetBenchPredict],
batch: Sequence[torch.Tensor],
records: Sequence[BaseRecord],
detection_threshold: float = 0.5,
keep_images: bool = False,
device: Optional[torch.device] = None,
) -> List[Prediction]:
device = device or model_device(model)
imgs, img_info = batch
imgs = imgs.to(device)
img_info = {k: v.to(device) for k, v in img_info.items()}
bench = DetBenchPredict(unwrap_bench(model))
bench = bench.eval().to(device)
raw_preds = bench(x=imgs, img_info=img_info)
preds = convert_raw_predictions(
batch=batch,
raw_preds=raw_preds,
records=records,
detection_threshold=detection_threshold,
keep_images=keep_images,
)
return preds
def load_detect(self, checkpoint_path):
config = get_efficientdet_config('tf_efficientdet_d0')
config.image_size = [384, 384]
net = EfficientDet(config, pretrained_backbone=False)
net.reset_head(num_classes=1)
checkpoint = torch.load(checkpoint_path,
map_location=torch.device('cpu'))
net.load_state_dict(checkpoint['model_state_dict'])
del checkpoint
gc.collect()
self.net = DetBenchPredict(net)
self.net.eval()
self.net.cpu()
def predict(
model: Union[DetBenchTrain, DetBenchPredict],
batch: Sequence[torch.Tensor],
detection_threshold: float = 0.5,
device: Optional[torch.device] = None,
):
device = device or model_device(model)
batch = [o.to(device) for o in batch]
bench = DetBenchPredict(unwrap_bench(model), config=model.config)
bench = bench.eval().to(device)
raw_preds = bench(*batch)
return convert_raw_predictions(raw_preds,
detection_threshold=detection_threshold)
def Load_Model(checkpoint_path, which_model='effi5', img_size=1280):
if which_model == 'effi5':
config = get_efficientdet_config(
'tf_efficientdet_d5') # tf_effi5 model structure
elif which_model == 'effi4':
config = get_efficientdet_config(
'tf_efficientdet_d4') # tf_effi4 model structure
elif which_model == 'effi6':
config = get_efficientdet_config(
'tf_efficientdet_d6') # tf_effi4 model structure
config.image_size = (img_size, img_size)
config.num_classes = 32
config.norm_kwargs = dict(eps=.001, momentum=.01)
net = EfficientDet(config, pretrained_backbone=False)
net.class_net = HeadNet(config, num_outputs=config.num_classes)
ckp = torch.load(checkpoint_path)
net.load_state_dict(ckp['model_state_dict'])
del ckp
net = DetBenchPredict(net)
net.eval()
return net
def get_efficientdet(checkpoint_path):
config = get_efficientdet_config('efficientdet_d1')
net = EfficientDet(config, pretrained_backbone=False)
checkpoint = torch.load(checkpoint_path)
net.load_state_dict(checkpoint, strict=False)
net = DetBenchPredict(net)
return net
def load_net(checkpoint_path):
config = get_efficientdet_config('tf_efficientdet_d7')
net = EfficientDet(config, pretrained_backbone=False)
config.num_classes = 1
config.image_size=512
net.class_net = HeadNet(config, num_outputs=config.num_classes, norm_kwargs=dict(eps=.001, momentum=.01))
print(f"Load model {checkpoint_path}")
checkpoint = torch.load(checkpoint_path)
net.cuda()
new_state_dict = OrderedDict()
for k, v in checkpoint['model_state_dict'].items():
if "anchors" in k:
print("Ignore: ",k)
continue
name = re.sub("model.",'',k) if k.startswith('model') else k
new_state_dict[name] = v
# net.load_state_dict(checkpoint['model_state_dict'])
net.load_state_dict(new_state_dict)
del checkpoint
gc.collect()
net = DetBenchPredict(net, config)
net.eval()
return net.cuda()
def __init__(self, num_classes=11, checkpoint_path=None):
super(EfficientDet5AP, self).__init__()
config = get_efficientdet_config('tf_efficientdet_d5_ap')
config.image_size = [512, 512]
config.norm_kwargs = dict(eps=.001, momentum=.01)
config.soft_nms = True
config.label_smoothing = 0.1
# config.legacy_focal = True
net = EfficientDet(config, pretrained_backbone=False)
if checkpoint_path == None:
checkpoint = torch.load(
'./effdet_model/tf_efficientdet_d5_ap-3673ae5d.pth')
net.load_state_dict(checkpoint)
net.reset_head(num_classes=num_classes)
net.class_net = HeadNet(config, num_outputs=config.num_classes)
self.model = DetBenchTrain(net, config)
else:
checkpoint = torch.load(checkpoint_path)
checkpoint2 = {
'.'.join(k.split('.')[2:]): v
for k, v in checkpoint.items()
}
del checkpoint2['boxes']
net.reset_head(num_classes=num_classes)
net.class_net = HeadNet(config, num_outputs=config.num_classes)
net.load_state_dict(checkpoint2)
self.model = DetBenchPredict(net)
def predict(
model: Union[DetBenchTrain, DetBenchPredict],
batch: Sequence[torch.Tensor],
detection_threshold: float = 0.5,
device: Optional[torch.device] = None,
):
device = device or model_device(model)
imgs, img_info = batch
imgs = imgs.to(device)
img_info = {k: v.to(device) for k, v in img_info.items()}
bench = DetBenchPredict(unwrap_bench(model))
bench = bench.eval().to(device)
raw_preds = bench(x=imgs, img_info=img_info)
return convert_raw_predictions(raw_preds,
detection_threshold=detection_threshold)
def get_net():
config = get_efficientdet_config('efficientdet_d0')
config.num_classes = 29
config.image_size = [512, 512]
net = EfficientDet(config, pretrained_backbone=False)
checkpoint = torch.load(
'effdet_checkpoints/2021-02-28 19_55_41/last-checkpoint.pth')
net.load_state_dict(checkpoint['model_state_dict'])
return DetBenchPredict(net)
def __init__(self, model_weight, num_class):
super(EfficientDetPred, self).__init__()
config = get_efficientdet_config(f'tf_efficientdet_{MODEL_USE}')
config.num_classes = num_class
config.image_size = [TRAIN_SIZE, TRAIN_SIZE]
model = EfficientDet(config, pretrained_backbone=False)
model.class_net = HeadNet(config, num_outputs=config.num_classes)
new_keys = model.state_dict().keys()
values = torch.load(model_weight, map_location=lambda storage, loc: storage).values()
model.load_state_dict(OrderedDict(zip(new_keys, values)))
self.model = DetBenchPredict(model)
del new_keys, values
gc.collect()
def load_net(checkpoint_path):
config = get_efficientdet_config('tf_efficientdet_d5')
net = EfficientDet(config, pretrained_backbone=False)
config.num_classes = 1
config.image_size=512
net.class_net = HeadNet(config, num_outputs=config.num_classes, norm_kwargs=dict(eps=.001, momentum=.01))
checkpoint = torch.load(checkpoint_path)
net.load_state_dict(checkpoint['model_state_dict'])
del checkpoint
gc.collect()
net = DetBenchPredict(net, config)
net.eval();
return net.cuda()
def create_model_from_config(config, bench_name='', pretrained=False, checkpoint_path='', **kwargs):
model = EfficientDet(config, **kwargs)
# FIXME handle different head classes / anchors and re-init of necessary layers w/ pretrained load
if checkpoint_path:
load_checkpoint(model, checkpoint_path)
elif pretrained:
load_pretrained(model, config.url)
config = copy.deepcopy(config)
# override num classes
config.num_classes = 1
model.class_net = HeadNet(config, num_outputs=1, norm_kwargs=dict(eps=0.001, momentum=0.01))
# wrap model in task specific bench if set
if bench_name == 'train':
model = DetBenchTrain(model, config)
model.loss_fn = DetectionLoss(config)
elif bench_name == 'predict':
model = DetBenchPredict(model, config)
return model
def get_predict_model(config_name,
img_size,
model_ckpt,
useGN=False,
light=False):
config = get_efficientdet_config(config_name)
model = EfficientDet(config, pretrained_backbone=False)
config.num_classes = 1
config.image_size = img_size
model.class_net = HeadNet(
config,
num_outputs=config.num_classes,
norm_kwargs=dict(eps=0.001, momentum=0.01),
)
if useGN is True:
model = convert_layers(model,
nn.BatchNorm2d,
nn.GroupNorm,
True,
num_groups=2)
if light is True:
count = 0
state_dict = torch.load(model_ckpt)["state_dict"]
new_state_dict = OrderedDict()
for key, value in state_dict.items():
if key.startswith("model.model."):
new_key = reduce(lambda a, b: a + "." + b, key.split(".")[2:])
if new_key in model.state_dict():
new_state_dict[new_key] = value
count += 1
model.load_state_dict(new_state_dict)
else:
model.load_state_dict(torch.load(model_ckpt)["state_dict"])
print(f"loaded {count} keys")
return DetBenchPredict(model, config)
class Predictor(nn.Module):
def __init__(self,
detect_path='weights/best-checkpoint.bin',
model_arc='weights/model_eff_arc.json',
weights='weights/model_best_acc.h5'):
#clf
super().__init__()
with open(model_arc, 'r') as f:
self.model = tf.keras.models.model_from_json(f.read())
self.model.load_weights(weights)
self.mapper = [
'TH', 'ACB', 'Acecook', 'Addidas', 'Agribank', 'Bidv', 'Big C',
'Cai Lan', 'Chinsu', 'Colgate', 'FPT', 'Habeco', 'Hai Ha',
'Jollibee', 'KFC', 'Kinh Do', 'Lotte mart', 'Mbbank new',
'Mbbank old', 'Neptune', 'Nike', 'Pepsi', 'Petrolimex',
'Phuc Long', 'Samsung', 'SHB', 'Techcombank', 'The Coffe House',
'The gioi di dong', 'TPbank', 'Vietcombank', 'Vietinbank',
'Viettel', 'Vinamilk', 'Vinfast', 'Vinmart', 'Vifon', 'Vnpt',
'Vpbank'
]
#detect
self.transform = A.Compose([
A.Resize(height=512, width=512, p=1.0),
ToTensor(),
])
self.load_detect(detect_path)
def preprocess_detect(self, path_img):
image = cv2.imread(path_img)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
transformer = self.transform(image=image)
image = transformer['image']
return image.unsqueeze(0).cpu().float()
def preprocess_clf(self, batch):
tensor = []
for img in batch:
img = cv2.resize(img, (224, 224))
img = img / 255.0
tensor.append(img)
return np.array(tensor)
def make_predictions(self, images, score_threshold=0.21):
# images = torch.stack(images).cpu().float()
predictions = []
with torch.no_grad():
det = self.net(
images, {
'img_scale': torch.tensor(
[1.] * images.shape[0]).float().cpu(),
'img_size': torch.tensor([512, 512]).cpu()
})
for i in range(images.shape[0]):
boxes = det[i].detach().cpu().numpy()[:, :4]
scores = det[i].detach().cpu().numpy()[:, 4]
labels = det[i].detach().cpu().numpy()[:, 5]
indexes = np.where(scores > score_threshold)[0]
predictions.append({
'boxes': boxes[indexes],
'scores': scores[indexes],
'labels': labels[indexes]
})
return [predictions]
def load_detect(self, checkpoint_path):
config = get_efficientdet_config('tf_efficientdet_d5')
config.image_size = [512, 512]
net = EfficientDet(config, pretrained_backbone=False)
net.reset_head(num_classes=1)
checkpoint = torch.load(checkpoint_path,
map_location=torch.device('cpu'))
net.load_state_dict(checkpoint['model_state_dict'])
del checkpoint
gc.collect()
self.net = DetBenchPredict(net)
self.net.eval()
self.net.cpu()
def forward(self, path_test):
images = self.preprocess_detect(path_test)
s_t = time.time()
predictions = self.make_predictions(images)
print(time.time() - s_t)
keep_idx = torchvision.ops.nms(
torch.from_numpy(predictions[0][0]['boxes']),
torch.from_numpy(predictions[0][0]['scores']), 0.1)
boxes = []
scores = []
labels = []
for i in keep_idx:
boxes.append(predictions[0][0]['boxes'][i])
scores.append(predictions[0][0]['scores'][i])
labels.append(predictions[0][0]['labels'][i])
boxes = np.array(boxes).astype(np.float32).clip(min=0, max=511)
image_original = cv2.imread(path_test)
image_original = cv2.cvtColor(image_original, cv2.COLOR_BGR2RGB)
h, w = image_original.shape[0], image_original.shape[1]
batch_clf = []
boxes[:, 0] = boxes[:, 0] * (w / 512)
boxes[:, 1] = boxes[:, 1] * (h / 512)
boxes[:, 2] = boxes[:, 2] * (w / 512)
boxes[:, 3] = boxes[:, 3] * (h / 512)
for box in boxes:
batch_clf.append(image_original[int(box[1]):int(box[3]),
int(box[0]):int(box[2]), :])
tensor = self.preprocess_clf(batch_clf)
time_clf = time.time()
print(tensor.shape)
predictions = self.model.predict(tensor)
print(f"clf:{time.time()-time_clf}")
result = []
for box, score, prediction in zip(boxes, scores, predictions):
box[0] = int(box[0])
box[1] = int(box[1])
box[2] = int(box[2])
box[3] = int(box[3])
idx = int(np.argmax(prediction))
result.append({
'box': box,
'label': self.mapper[idx],
'score_detect': score,
'score_clf': prediction[idx]
})
return image_original, result
def validate(args):
# might as well try to validate something
args.pretrained = args.pretrained or not args.checkpoint
args.prefetcher = not args.no_prefetcher
args.redundant_bias = not args.no_redundant_bias
# create model
config = get_efficientdet_config(args.model)
config.redundant_bias = args.redundant_bias
config.num_classes = 3
model = EfficientDet(config)
if args.checkpoint:
state_dict = load_state_dict(args.checkpoint)
model.load_state_dict(state_dict, strict=True)
param_count = sum([m.numel() for m in model.parameters()])
print('Model %s created, param count: %d' % (args.model, param_count))
bench = DetBenchPredict(model, config)
bench = bench.cuda()
if has_amp:
print('Using AMP mixed precision.')
bench = amp.initialize(bench, opt_level='O1')
else:
print('AMP not installed, running network in FP32.')
if args.num_gpu > 1:
bench = torch.nn.DataParallel(bench,
device_ids=list(range(args.num_gpu)))
annotation_path = os.path.join(args.data, 'annotations',
f'test_annotations.json')
image_dir = os.path.join(args.data, 'test')
print(annotation_path)
print(image_dir)
dataset = CocoDetection(os.path.join(args.data, image_dir),
annotation_path)
print(dataset)
loader = create_loader(dataset,
input_size=config.image_size,
batch_size=args.batch_size,
use_prefetcher=args.prefetcher,
interpolation=args.interpolation,
fill_color=args.fill_color,
num_workers=args.workers,
pin_mem=args.pin_mem)
print(len(loader))
img_ids = []
results = []
model.eval()
batch_time = AverageMeter()
end = time.time()
with torch.no_grad():
for i, (input, target) in enumerate(loader):
output = bench(input, target['img_scale'], target['img_size'])
output = output.cpu()
sample_ids = target['img_id'].cpu()
for index, sample in enumerate(output):
image_id = int(sample_ids[index])
for det in sample:
score = float(det[4])
if score < .001: # stop when below this threshold, scores in descending order
break
coco_det = dict(image_id=image_id,
bbox=det[0:4].tolist(),
score=score,
category_id=int(det[5]))
img_ids.append(image_id)
results.append(coco_det)
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % args.log_freq == 0:
print(
'Test: [{0:>4d}/{1}] '
'Time: {batch_time.val:.3f}s ({batch_time.avg:.3f}s, {rate_avg:>7.2f}/s) '
.format(
i,
len(loader),
batch_time=batch_time,
rate_avg=input.size(0) / batch_time.avg,
))
json.dump(results, open(args.results, 'w'), indent=4)
return results
class Predictor(nn.Module):
def __init__(self, detect_path):
super().__init__()
#detect
self.transform = A.Compose([
A.Resize(height=384, width=384, p=1.0),
ToTensor(),
])
self.load_detect(detect_path)
def preprocess_detect(self, path_img):
image = cv2.imread(path_img)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
transformer = self.transform(image=image)
image = transformer['image']
return image.unsqueeze(0).cpu().float()
def make_predictions(self, images, score_threshold=0.25):
# images = torch.stack(images).cpu().float()
predictions = []
with torch.no_grad():
det = self.net(
images, {
'img_scale': torch.tensor(
[1.] * images.shape[0]).float().cpu(),
'img_size': torch.tensor([384, 384]).cpu()
})
for i in range(images.shape[0]):
boxes = det[i].detach().cpu().numpy()[:, :4]
scores = det[i].detach().cpu().numpy()[:, 4]
labels = det[i].detach().cpu().numpy()[:, 5]
indexes = np.where(scores > score_threshold)[0]
predictions.append({
'boxes': boxes[indexes],
'scores': scores[indexes],
'labels': labels[indexes]
})
return [predictions]
def load_detect(self, checkpoint_path):
config = get_efficientdet_config('tf_efficientdet_d0')
config.image_size = [384, 384]
net = EfficientDet(config, pretrained_backbone=False)
net.reset_head(num_classes=1)
checkpoint = torch.load(checkpoint_path,
map_location=torch.device('cpu'))
net.load_state_dict(checkpoint['model_state_dict'])
del checkpoint
gc.collect()
self.net = DetBenchPredict(net)
self.net.eval()
self.net.cpu()
def forward(self, path_test):
images = self.preprocess_detect(path_test)
s_t = time.time()
predictions = self.make_predictions(images)
print(time.time() - s_t)
batch_recognize = []
boxes = []
scores = []
labels = []
if (predictions[0][0]['boxes'] != []):
keep_idx = torchvision.ops.nms(
torch.from_numpy(predictions[0][0]['boxes']),
torch.from_numpy(predictions[0][0]['scores']), 0.1)
for i in keep_idx:
boxes.append(predictions[0][0]['boxes'][i])
scores.append(predictions[0][0]['scores'][i])
labels.append(predictions[0][0]['labels'][i])
boxes = np.array(boxes).astype(np.float32).clip(min=0, max=511)
image_original = cv2.imread(path_test)
image_original = cv2.cvtColor(image_original, cv2.COLOR_BGR2RGB)
h, w = image_original.shape[0], image_original.shape[1]
boxes[:, 0] = boxes[:, 0] * (w / 384)
boxes[:, 1] = boxes[:, 1] * (h / 384)
boxes[:, 2] = boxes[:, 2] * (w / 384)
boxes[:, 3] = boxes[:, 3] * (h / 384)
for box in boxes:
batch_recognize.append(
image_original[int(box[1]):int(box[3]),
int(box[0]):int(box[2]), :])
return batch_recognize, boxes