[AI] Flask 웹 완성

갑자기 뚝딱
완성해버렸다

index.html 에 form tag 로 사진 등록하는 부분 만들어주었고

            <form method="post" action="/upload" enctype="multipart/form-data">
                <input type="file" name="file">
                <input type="submit" value="upload">
            </form>
            <p></p>
            {% if filename %}
            <div>
                <img src="{{ filename }}" alt = "{{ filename }}">
            </div>
            <div>
                <p> 품목 : {{ label01 }} <br>
                    품종 : {{ label02 }} <br>
                    등급 : {{ label03 }} <br> 
                    가로 : {{ width }} cm<br>
                    세로 : {{ height }} cm<br>
                    무게 : {{ weight }} g<br> 
                    </p>
            </div>
            {% else %}
            <div>    
                <h3>판정에 사용할 사진을 업로드해주세요.</h3>
            </div>
            {% endif %}

submit을 누르면 요 함수로 연결되도록 flask에 작성했다.

@app.route('/upload', methods = ['POST'])
def upload(): 
    file = request.files['file']
    filename = file.filename
    img_src = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    file.save(img_src)
    # solution 함수로 이미지 넘겨서 결과 받아오기
    label, prob, width, height, weight = solution(img_src)
    # 결과 포함해서 index.html로 넘겨주기
    label_encode = label.split('_')
    # return render_template('index.html', filename=img_src, label=label, probability=prob)
    return render_template('index.html', filename=img_src, 
    label01 = labels_kor[label_encode[0]],
    label02 = labels_kor[label_encode[1]],
    label03 = labels_kor[label_encode[2]],
    width = width,
    height = height,
    weight = weight)

solution() 함수에서 classification과 Linear Regression 모델을 모두 사용하여 결과를 반환한다.

classification 부분

  # 모델 불러오기
    model, size = models.initialize_model("resnet18", 54, use_pretrained=False)
    model.load_state_dict(torch.load("./resnet18_best.pt"))
    
    # 이미지 open, 전처리
    image = Image.open(filename).convert("RGB")
    image = test_transform(image)
    image = image.unsqueeze(0)
    
    # 모델에 넣어 결과 가져오기
    with torch.no_grad():
        model.eval()
        # Model outputs log probabilities
        output = model(image)
        output = torch.exp(output)
        _, index = torch.max(output, 1)
        label = classes[index]

여기에서 model.eval()을 안 해줘서 꽤 오랜 시간 삽질을 했다...
exp()는 꼭 해줄 필요는 없는데 그냥 숫자가 크게 보이니 좋아서 해줬다. 0.얼마들 보고 있으니 머리 아파서...

Linear regression 부분

width, height, weight = predictvalue.predict(index)

predictvalue.py의 predict 함수를 호출했다. 이 함수는 이렇게 생겼다.

def predict(label):
    data=pd.get_dummies(classes)
    # label=torch.Tensor([0])
    label=int(torch.Tensor.numpy(label)[0])
    for i in range(0,54):
        if i==classess[label]-1:
            x=data.values[i]
    x=np.array(x)
    x=torch.Tensor(x)

    model = nn.Linear(54,3)
    model.load_state_dict(loaded_model)
    model.eval()
    with torch.no_grad():
        outputs = model(x)
        width=outputs[0]
        height=outputs[1]
        weight=outputs[2]

    # print("STAT", width, height, weight)
    return width, height, weight

성우님이 거의 작성해주셔서 포맷만 맞춰서 리턴만 수정했다.