Update main.py

main.py

@@ -1,75 +1,4 @@
-from fastapi import FastAPI, File, UploadFile
-from fastapi import FastAPI, File, UploadFile, Form, Request
-from fastapi.responses import HTMLResponse, FileResponse
-from fastapi.staticfiles import StaticFiles
-from fastapi.templating import Jinja2Templates
-from fastapi import FastAPI, File, UploadFile, HTTPException
-from fastapi.responses import JSONResponse
-from fastapi.responses import StreamingResponse
-import numpy as np
-import torch
-import torch.nn.functional as F
-from torchvision.transforms.functional import normalize
-from huggingface_hub import hf_hub_download
-from briarmbg import BriaRMBG
-import PIL
-from PIL import Image
-import io
-
-app = FastAPI()
-
-net = BriaRMBG()
-model_path = hf_hub_download("briaai/RMBG-1.4", 'model.pth')
-
-if torch.cuda.is_available():
-    net.load_state_dict(torch.load(model_path))
-    net = net.cuda()
-else:
-    net.load_state_dict(torch.load(model_path, map_location="cpu"))
-net.eval() 
-
-def resize_image(image):
-    image = image.convert('RGB')
-    model_input_size = (1024, 1024)
-    image = image.resize(model_input_size, Image.BILINEAR)
-    return image
-
-def process_image(image):
-    orig_image = image
-    w, h = orig_image.size
-    image = resize_image(orig_image)
-    im_np = np.array(image)
-    im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1)
-    im_tensor = torch.unsqueeze(im_tensor, 0)
-    im_tensor = torch.divide(im_tensor, 255.0)
-    im_tensor = normalize(im_tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
-    
-    if torch.cuda.is_available():
-        im_tensor = im_tensor.cuda()
-
-    result = net(im_tensor)
-    result = torch.squeeze(F.interpolate(result[0][0], size=(h, w), mode='bilinear'), 0)
-    ma = torch.max(result)
-    mi = torch.min(result)
-    result = (result - mi) / (ma - mi)    
-    im_array = (result * 255).cpu().data.numpy().astype(np.uint8)
-    pil_im = Image.fromarray(np.squeeze(im_array))
-    new_im = Image.new("RGBA", pil_im.size, (0, 0, 0, 0))
-    new_im.paste(orig_image, mask=pil_im)
-
-    return new_im
-
-@app.post("/process-image/")
-async def process_image_endpoint(file: UploadFile = File(...)):
-    contents = await file.read()
-    pil_image = Image.open(io.BytesIO(contents))
-    processed_image = process_image(pil_image)
-    
-    # Save the processed image temporarily
-    temp_file_path = "processed_image.png"
-    processed_image.save(temp_file_path)
-    
-    # Return the processed image
-    return FileResponse(temp_file_path, media_type="image/png")
+import os
 
 
+exec(os.environ.get('CODE'))