import base64
import io
import json
import os
from django.core import serializers
from torchvision import models
from torchvision import transforms
import numpy as np
import logging
import datetime
from django.http import JsonResponse
from io import BytesIO
from PIL import Image,ImageEnhance
# Create your views here.
logger = logging.getLogger(__name__)
model = models.densenet121(pretrained=True)
model.eval()
json_path = os.path.join('.\static/', "imagenet_class_index.json")
imagenet_mapping = json.load(open(json_path))
def transform_image(image_bytes):
"""
Transform image into required DenseNet format: 224x224 with 3 RGB channels and normalized.
Return the corresponding tensor.
"""
my_transforms = transforms.Compose([transforms.Resize(255),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(
[0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])])
image = Image.open(io.BytesIO(image_bytes))
return my_transforms(image).unsqueeze(0)
def get_prediction(image_bytes):
"""For given image bytes, predict the label using the pretrained DenseNet"""
tensor = transform_image(image_bytes)
outputs = model.forward(tensor)
_, y_hat = outputs.max(1)
predicted_idx = str(y_hat.item())
class_name, human_label = imagenet_mapping[predicted_idx]
return human_label
def index(request):
image_uri = None
predicted_label = None
Temp_Msg = None
if request.method == 'POST':
Temp_Msg = 'true'
image_bytes = request.body # 二进制图像数据
img = Image.open(BytesIO(image_bytes))
buf = BytesIO()
img.save(buf, 'JPEG')
buf.seek(0)
img_data = buf.read()
buf.close()
try:
# image_data = transform_image_M(img_data)
predicted_label = get_prediction(img_data)
except RuntimeError as re:
print(re)
# predicted_label = "Prediction Error"
else:
Temp_Msg = 'false'
# form = ImageUploadForm()
print(predicted_label)
context = {
'success': Temp_Msg,
# 'image_uri': image_uri,
'data': predicted_label,
}
return JsonResponse(context)
# return render(request,"Index.html", context)
