gnt_pseudo_label.py

import cv2

import sys, os
import json
from tqdm import tqdm

sys.path.insert(0, '.')
from configs import add_centernet_config
from detectron2.config import get_cfg
from inference.centernet import build_model
from detectron2.checkpoint import DetectionCheckpointer

def get_widerface_images():
    lines = open('datasets/wider_face_add_lm_10_10_add_mafa/ImageSets/Main/trainval_all.txt').readlines()
    root = 'datasets/wider_face_add_lm_10_10_add_mafa/JPEGImages/'
    images = [root + i.strip() + '.jpg' for i in lines]
    return images

def get_crowdhuman_images():
    json_file = 'datasets/crowd_human/Annotations/annotation_train.odgt'
    lines = open(json_file).readlines()
    root = 'datasets/crowd_human/JPEGImages/'
    images = [root + json.loads(line.strip('\n'))['ID']+ '.jpg' for line in lines]
    return images


if __name__ == "__main__":
    # cfg
    cfg = get_cfg()
    add_centernet_config(cfg)
    #cfg.merge_from_file("yamls/coco_det/centernet_r_50_C4_0.5x_coco_person.yaml")
    cfg.merge_from_file("yamls/person_face/face_res50.yaml")

    # model
    model = build_model(cfg)
    #DetectionCheckpointer(model).load("models/coco_det_crowd_R50_SGD.pth")
    DetectionCheckpointer(model).load("models/person_face_R50_face_adam.pth")
    model.eval()

    #txt
    #txt = open('datasets/wider_face_add_lm_10_10_add_mafa/widerface_pseudo_human.txt','w')
    txt = open('datasets/crowd_human/crowdhuman_pseudo_face.txt','w')
    class_name = 'face'

    # images
    images = get_crowdhuman_images()
    bs = 8
    for i in tqdm(range(0, len(images), bs)):
        images_rgb = [cv2.imread(j)[:,:,::-1] for j in images[i:i + bs]]
        img_names = [os.path.basename(j) for j in images[i:i + bs]]
        results = model.inference_on_images(images_rgb, K=100, max_size=640)
        for k,result in enumerate(results):
            cls = result['cls'].cpu().numpy()
            bbox = result['bbox'].cpu().numpy()
            scores = result['scores'].cpu().numpy()
            H,W,C = images_rgb[k].shape
            img = images_rgb[k][:,:,::-1]
            img_name = img_names[k]
            line = ';'.join([img_name,str(W),str(H)])
            for c,(x1,y1,x2,y2),s in zip(cls,bbox,scores):
                if c != 0.0 or s < 0.3:
                    continue
                x1 = str(max(0, int(x1)))
                y1 = str(max(0, int(y1)))
                x2 = str(min(W, int(x2)))
                y2 = str(min(H, int(y2)))
                s = str(round(float(s),3))
                line += ';'.join(['',class_name,x1,y1,x2,y2])
            txt.write(line+'\n')