[ML][Enhancement] Image Classification

[yehyunsuh]

What

Image Classification

Why

For use in our project

How

• [x] type & shape classification
• [x] color classification
• [x] tablet/capsule classification

reference: #15 #25

[yehyunsuh]

type & shape classification

pill_type = ["알약_원형",'알약_장방형or타원형','알약_기타','알약_팔각형','알약_삼각형','알약_사각형','알약_오각형','알약_육각형','알약_마름모형','캡슐_장방형or타원형','캡슐_기타']

color classification

pill_type = sorted(['투명', '연두, 투명', '초록, 투명', '회색', '하양, 투명', '갈색or빨강', '노랑or주황', '주황, 투명', '하양, 노랑', '분홍, 투명', '갈색, 투명', '보라', '연두or초록', '보라, 투명', '청록, 투명', '하양, 갈색', '청록', '분홍', '하양, 파랑', '하양', '빨강, 투명', '남색', '파랑, 투명', '파랑', '검정', '자주', '노랑, 투명']) 

tablet/capsule classification

pill_type = ['알약','캡슐']

How to bring pretrained model

model = timm.create_model('resnet50', pretrained=True, num_classes=30)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 
model.to(device) 

PATH = ""
model.load_state_dict(torch.load(PATH, map_location=device))

If evaluation needed

with torch.no_grad():
    print("Calculating validation results...")
    model.eval()
    i, accuracy = 0, []
    for test in test_loader:
        inputs = test
        inputs = inputs.to(device)

        outs = model(inputs)
        preds = torch.argmax(outs, dim=-1)

        print(preds)   # result of the classification

[yehyunsuh]

Classification done. The train dataset will be changed to segmented data. Issue close.