def set_transform(self):
if self.augments:
color_jitter = OneOf(transforms=[
Identity(),
HSV(p=1,
hgain=self.aug_cfg['hsv_h'],
sgain=self.aug_cfg['hsv_s'],
vgain=self.aug_cfg['hsv_v']),
RandNoise()
])
mosaic = Mosaic(self.img_paths,
self.labels,
color_gitter=color_jitter,
target_size=self.img_size,
pad_val=self.aug_cfg['pad_val'],
rand_center=True)
mix_up = MixUp(self.img_paths,
self.labels,
color_gitter=color_jitter,
target_size=self.img_size,
pad_val=self.aug_cfg['pad_val'],
beta=self.aug_cfg['beta'])
perspective_transform = RandPerspective(
target_size=(self.img_size, self.img_size),
scale=self.aug_cfg['scale'],
degree=self.aug_cfg['degree'],
translate=self.aug_cfg['translate'],
shear=self.aug_cfg['shear'],
pad_val=self.aug_cfg['pad_val'])
basic_transform = Compose(transforms=[
color_jitter,
RandCutOut(),
ScalePadding(target_size=self.img_size,
padding_val=self.aug_cfg['pad_val']),
perspective_transform,
])
aug_mosaic = Mosaic(self.img_paths,
self.labels,
color_gitter=mix_up,
target_size=self.img_size,
pad_val=self.aug_cfg['pad_val'],
rand_center=True)
aug_mixup = MixUp(self.img_paths,
self.labels,
color_gitter=mosaic,
target_size=self.img_size,
pad_val=self.aug_cfg['pad_val'],
beta=self.aug_cfg['beta'])
self.transform = Compose(transforms=[
OneOf(transforms=[(0., basic_transform), (
1., mosaic), (0., mix_up), (0., aug_mosaic), (0.,
aug_mixup)]),
LRFlip()
])
else:
self.transform = ScalePadding(target_size=(self.img_size,
self.img_size),
padding_val=self.aug_cfg['pad_val'])
def write_coco_json():
from pycocotools.coco import COCO
img_root = "/home/wangchao/public_dataset/coco/images/val2017"
model = FCOS()
weights = torch.load("weights/auto_assign_1_10_best_map50.pth")['ema']
model.load_state_dict(weights)
model.cuda().eval()
basic_transform = ScalePadding(target_size=(768, 768),
padding_val=(103, 116, 123))
coco = COCO(
"/home/wangchao/public_dataset/coco/annotations/instances_val2017.json"
)
coco_predict_list = list()
for img_id in tqdm(coco.imgs.keys()):
file_name = coco.imgs[img_id]['file_name']
img_path = os.path.join(img_root, file_name)
img = cv.imread(img_path)
# ori_img = img.copy()
img, ratio, (left, top) = basic_transform.make_border(img)
h, w = img.shape[:2]
img_out = img[:, :, [2, 1, 0]].astype(np.float32) / 255.0
img_out = ((img_out - np.array(rgb_mean)) /
np.array(rgb_std)).transpose(2, 0, 1).astype(np.float32)
img_out = torch.from_numpy(
np.ascontiguousarray(img_out)).unsqueeze(0).float().cuda()
predicts = model(img_out)
for i in range(len(predicts)):
predicts[i][:, [0, 2]] = predicts[i][:, [0, 2]].clamp(min=0, max=w)
predicts[i][:, [1, 3]] = predicts[i][:, [1, 3]].clamp(min=0, max=h)
box = non_max_suppression(predicts,
conf_thresh=0.05,
iou_thresh=0.5,
max_det=300)[0]
if box is None:
continue
box[:, [0, 2]] = (box[:, [0, 2]] - left) / ratio[0]
box[:, [1, 3]] = (box[:, [1, 3]] - top) / ratio[1]
box = box.detach().cpu().numpy()
# ret_img = draw_box(ori_img, box[:, [4, 5, 0, 1, 2, 3]], colors=coco_colors)
# cv.imwrite(file_name, ret_img)
coco_box = box[:, :4]
coco_box[:, 2:] = coco_box[:, 2:] - coco_box[:, :2]
for p, b in zip(box.tolist(), coco_box.tolist()):
coco_predict_list.append({
'image_id': img_id,
'category_id': coco_ids[int(p[5])],
'bbox': [round(x, 3) for x in b],
'score': round(p[4], 5)
})
with open("predicts.json", 'w') as file:
json.dump(coco_predict_list, file)
def set_transform(self):
if self.augments:
self.transform = Compose(transforms=[
OneOf(transforms=[
(0.6, Compose(transforms=[
OneOf(transforms=[Identity(),
HSV(p=1,
hgain=self.aug_cfg['hsv_h'],
sgain=self.aug_cfg['hsv_s'],
vgain=self.aug_cfg['hsv_v']),
RandNoise()
]),
RandCutOut(),
ScalePadding(target_size=self.img_size, padding_val=self.aug_cfg['pad_val']),
RandPerspective(target_size=(self.img_size, self.img_size),
scale=self.aug_cfg['scale'],
degree=self.aug_cfg['degree'],
translate=self.aug_cfg['translate'],
shear=self.aug_cfg['shear'],
pad_val=self.aug_cfg['pad_val'])])),
(0.4, Mosaic(self.img_paths,
self.labels,
color_gitter=OneOf(transforms=[Identity(),
HSV(p=1,
hgain=self.aug_cfg['hsv_h'],
sgain=self.aug_cfg['hsv_s'],
vgain=self.aug_cfg['hsv_v']),
RandNoise()]),
target_size=self.img_size,
pad_val=self.aug_cfg['pad_val']))
]),
LRFlip()])
else:
self.transform = ScalePadding(target_size=(self.img_size, self.img_size),
padding_val=self.aug_cfg['pad_val'])
def write_coco_json():
from pycocotools.coco import COCO
device = torch.device("cuda:0")
img_root = "/home/ubuntu/wangchao/dataset/coco/val2017"
with open("config/centernet.yaml", 'r') as rf:
cfg = yaml.safe_load(rf)
model_cfg = cfg['model']
model = CenterNet(num_cls=model_cfg['num_cls'],
PIXEL_MEAN=model_cfg['PIXEL_MEAN'],
PIXEL_STD=model_cfg['PIXEL_STD'],
backbone=model_cfg['backbone'],
cfg=model_cfg)
weights = torch.load("weights/0_TTFNet_best_map.pth")['ema']
model.load_state_dict(weights)
model.to(device)
model.eval()
basic_transform = ScalePadding(target_size=(512, 512),
padding_val=(103, 116, 123))
coco = COCO(
"/home/ubuntu/wangchao/dataset/coco/annotations/instances_val2017.json"
)
coco_predict_list = list()
for img_id in tqdm(coco.imgs.keys()):
file_name = coco.imgs[img_id]['file_name']
img_path = os.path.join(img_root, file_name)
img = cv.imread(img_path)
# ori_img = img.copy()
img, ratio, (left, top) = basic_transform.make_border(img)
h, w = img.shape[:2]
img_out = img[:, :, [2, 1, 0]].astype(np.float32) / 255.0
img_out = ((img_out - np.array(rgb_mean)) /
np.array(rgb_std)).transpose(2, 0, 1).astype(np.float32)
img_out = torch.from_numpy(
np.ascontiguousarray(img_out)).unsqueeze(0).float().to(device)
predicts = model(
img_out
) #list(predict) predict.shape=[num_box,6] 6==>x1,y1,x2,y2,score,label
for i in range(len(predicts)):
predicts[i][:, [0, 2]] = predicts[i][:, [0, 2]].clamp(min=0, max=w)
predicts[i][:, [1, 3]] = predicts[i][:, [1, 3]].clamp(min=0, max=h)
box = predicts[0]
if box is None:
continue
box[:, [0, 2]] = (box[:, [0, 2]] - left) / ratio[0]
box[:, [1, 3]] = (box[:, [1, 3]] - top) / ratio[1]
box = box.detach().cpu().numpy()
# ret_img = draw_box(ori_img, box[:, [4, 5, 0, 1, 2, 3]], colors=coco_colors)
# cv.imwrite(file_name, ret_img)
coco_box = box[:, :4]
coco_box[:, 2:] = coco_box[:, 2:] - coco_box[:, :2]
for p, b in zip(box.tolist(), coco_box.tolist()):
coco_predict_list.append({
'image_id': img_id,
'category_id': coco_ids[int(p[5])],
'bbox': [round(x, 3) for x in b],
'score': round(p[4], 5)
})
with open("predicts.json", 'w') as file:
json.dump(coco_predict_list, file)