adpro commited on
Commit
02cc722
·
verified ·
1 Parent(s): 7d25a11

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +35 -0
app.py ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI, UploadFile, File
2
+ import uvicorn
3
+ import cv2
4
+ import numpy as np
5
+ import torch
6
+ import torchvision.transforms as T
7
+ from PIL import Image
8
+ import io
9
+
10
+ app = FastAPI()
11
+
12
+ # Load AI model MiDaS
13
+ midas = torch.hub.load("intel-isl/MiDaS", "MiDaS_small")
14
+ midas.eval()
15
+ transform = T.Compose([T.Resize((256, 256)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
16
+
17
+ @app.post("/upload/")
18
+ async def upload_image(file: UploadFile = File(...)):
19
+ image_bytes = await file.read()
20
+ image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
21
+
22
+ # Convert to tensor & run AI model
23
+ img_tensor = transform(image).unsqueeze(0)
24
+ with torch.no_grad():
25
+ depth_map = midas(img_tensor).squeeze().cpu().numpy()
26
+
27
+ # Normalize depth map
28
+ depth_map = cv2.normalize(depth_map, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
29
+ depth_img = cv2.applyColorMap(depth_map, cv2.COLORMAP_JET)
30
+
31
+ _, buffer = cv2.imencode(".jpg", depth_img)
32
+ return {"depth_map": buffer.tobytes()}
33
+
34
+ if __name__ == "__main__":
35
+ uvicorn.run(app, host="0.0.0.0", port=7860)