MonumentDetection/ContinualLearningFastAPI

1from fastapi import FastAPI, Request, Form, UploadFile2from fastapi.templating import Jinja2Templates3from fastapi.responses import HTMLResponse4from fastapi.staticfiles import StaticFiles5from pydantic import BaseModel6import io7import base648from matplotlib.backends.backend_agg import FigureCanvasAgg9from pathlib import Path10from monuments.faster_models.fasterrcnn import fasterrcnn_resnet50_fpn, filter_pred, classes, CLASSES1, CLASSES2, CLASSES3, CLASSES411import os12import shutil13from datetime import datetime14import torch15import torchvision.transforms as transforms16from PIL import Image17import numpy as np18import matplotlib.pyplot as plt19import matplotlib.patches as patches20 21 22app = FastAPI()23app.mount("/static", StaticFiles(directory="static"), name="static")24app.mount("/output", StaticFiles(directory="output"), name="output")25 26templates = Jinja2Templates(directory="templates")27imageBytes = None28model = None29class ImageForm(BaseModel):30    image: UploadFile31 32 33@app.get("/")34async def monuments(request: Request):35    return templates.TemplateResponse("welcome.html", {"request": request})36 37 38@app.get("/index/", response_class=HTMLResponse)39async def index(request: Request):40    return templates.TemplateResponse("index.html", {"request": request})41 42 43# @app.post("/upload")44# async def upload(request: Request, image: UploadFile = File(...)):45#     print('reached')46#     if image:47#         image_path = f"media/images/{image.filename}"48#         print(image_path)49#         with open(image_path, "wb") as image_file:50#             image_file.write(image.file.read())51 52#         return templates.TemplateResponse("upload.html", {"request": request, "form": image, "img_object": image_path})53    54#     return templates.TemplateResponse("image_form.html", {"request": request, "form": image})55@app.post("/upload/")56async def upload(request: Request,file: UploadFile):57    global imageBytes58    imageBytes = await file.read()59    image = Image.open(io.BytesIO(imageBytes))60    if image.mode != 'RGB':61        image = image.convert('RGB')62    image_bytes = io.BytesIO()63    image.save(image_bytes, format="JPEG") 64 65    contents = base64.b64encode(image_bytes.getvalue()).decode("utf-8")66    contents = contents.split('\n')[0]67 68    return templates.TemplateResponse("upload.html", {"request":request, "image_content": contents })69 70@app.post("/predict/")71async def predict(request: Request, model: str = Form(...), subset: str = Form(...)):72    global imageBytes73    try:74        image = Image.open(io.BytesIO(imageBytes))75        if image.mode != 'RGB':76            image = image.convert('RGB')77    except Exception as e:78        print(f"Error: {e}")79    if model:80        # Define transformations81        transform = transforms.Compose([82            transforms.ToTensor(),83        ])84        image = image.convert('RGB')85        img_tensor = transform(image)86        img_tensor = img_tensor.unsqueeze(0)87        88        base_dir = os.path.dirname(__file__)89        subsetList = ["subset1", "subset2", "subset3", "subsset4"]90        if subset in subsetList:91            subsetCategory = subsetList.index(subset)92            classesList = [CLASSES1,CLASSES2,CLASSES3,CLASSES4]93            CLASSES = classesList[subsetCategory]94            if model in ["joint_model", "meta_model", "MNAD_model"]:95                print(model)96                model_path = os.path.join(base_dir, 'Base Model', subset, f'{model}.pth')97        # if model == "joint_model":98        #     model_path = os.path.join(dir, 'Base Model', 'joint_model.pth')99        # elif model == "meta_model":100        #     model_path = os.path.join(dir, 'Base Model', 'meta_model.pth')101        # elif model == "MNAD_model":102        #     model_path = os.path.join(dir, 'Base Model', 'MNAD_model.pth')103        # else:104        #     # Handle the case when model is not any of the specified values105        #     raise ValueError(f"Unsupported model: {model}")106 107        model = fasterrcnn_resnet50_fpn(num_classes=21)108        model.load_state_dict(torch.load(model_path, map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu')))109 110        model.to('cuda' if torch.cuda.is_available() else 'cpu')111        model.eval()112 113        with torch.no_grad():114            predictions = model(img_tensor)115 116        outputs = filter_pred(predictions)117        boxes = outputs[0]['boxes'].cpu().numpy()118        labels = outputs[0]['labels'].cpu().numpy()119        scores = outputs[0]['scores'].cpu().numpy()120 121        original_np = np.array(image)122 123        # Original Image124        fig, axs = plt.subplots(figsize=(10, 5))125        axs.imshow(original_np)  # Assuming original images are in CHW format126        axs.axis('off')127        axs.set_title('Prediction')128        129        130        # Add predicted bounding boxes to the predicted image131        for j, box in enumerate(boxes):132            rect = patches.Rectangle(133                (box[0], box[1]), box[2] - box[0], box[3] - box[1], linewidth=2, edgecolor='r', facecolor='none'134            )135            axs.add_patch(rect)136            axs.text(137                box[0], box[1] - 5, f'{CLASSES[int(labels[j])]}' , color='r', fontsize=10,138                bbox=dict(facecolor='white', alpha=0.8, edgecolor='none', boxstyle='round,pad=0.2')139            )140        buffer = io.BytesIO()141        canvas = FigureCanvasAgg(plt.gcf())142        canvas.print_png(buffer)143        bytes_data = buffer.getvalue()144 145        image = Image.open(io.BytesIO(bytes_data))146        image_rgb = image.convert('RGB')147        output = io.BytesIO()148        image_rgb.save(output, format = "JPEG")149 150 151        contents = base64.b64encode(output.getvalue()).decode("utf-8")152        contents = contents.split('\n')[0]153 154        return templates.TemplateResponse("predict.html", {"request": request, "image_content": contents,"confidence_score": round(scores[0],3)})155    return {"detail": "Invalid model specified."}156