|
|
|
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) |
|
|
