File size: 2,953 Bytes
46dcd1f a940bc4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
# main.py
import json
from fastapi import FastAPI, File, UploadFile, Request
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.templating import Jinja2Templates
from PIL import Image
from io import BytesIO
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torchvision import models
class FruitRecognizer:
def __init__(self, model_path, num_classes):
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model = models.resnet18(pretrained=False)
self.model.fc = nn.Linear(self.model.fc.in_features, num_classes)
self.model.load_state_dict(torch.load(model_path, map_location=self.device))
self.model.to(self.device)
self.model.eval()
self.transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
def recognize_fruit_from_path(self, image_path, class_names):
img = Image.open(image_path).convert("RGB")
img = self.transform(img)
img = img.unsqueeze(0).to(self.device)
with torch.no_grad():
outputs = self.model(img)
_, predicted = torch.max(outputs.data, 1)
predicted_class = class_names[predicted.item()]
return predicted_class
def recognize_fruit(self, image, class_names):
img = self.transform(image)
img = img.unsqueeze(0).to(self.device)
with torch.no_grad():
outputs = self.model(img)
_, predicted = torch.max(outputs.data, 1)
predicted_class = class_names[predicted.item()]
return predicted_class
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
with open('metadata.json', 'r') as f:
metadata = json.load(f)
class_names = metadata['classes']
model_path = "models/fruit_recognition_model.pth"
recognizer = FruitRecognizer(model_path, len(class_names))
templates = Jinja2Templates(directory="templates")
@app.get("/", response_class=HTMLResponse)
async def root(request: Request):
return templates.TemplateResponse("index.html", {"request": request})
@app.post("/predict/")
async def predict_fruit(request: Request, file: UploadFile = File(...)):
try:
print("request")
img = Image.open(BytesIO(await file.read())).convert("RGB")
predicted_class = recognizer.recognize_fruit(img, class_names)
return JSONResponse({"predicted_class": predicted_class})
except Exception as e:
return JSONResponse({"error": str(e)}, status_code=400)
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="localhost", port=7860)
|