Django 基于densenet121的简单目标检测

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)