Spaces:
Runtime error
Runtime error
File size: 3,498 Bytes
79b0a63 |
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 93 94 95 96 97 98 99 |
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, UploadFile, File, Request
from fastapi.responses import JSONResponse, HTMLResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from detection import initialize_directories, slice_geotiff, detect_weeds_in_slices, cleanup, create_zip # Import detection functions
from db_user_info import insert_user_info
from db_bucket import upload_file_to_bucket
from typing import List
from io import BytesIO
app = FastAPI()
# Mount static files
app.mount("/static", StaticFiles(directory="static"), name="static")
# WebSocket manager
class ConnectionManager:
def __init__(self):
self.active_connections: List[WebSocket] = []
async def connect(self, websocket: WebSocket):
await websocket.accept()
self.active_connections.append(websocket)
def disconnect(self, websocket: WebSocket):
self.active_connections.remove(websocket)
async def send_message(self, message: str):
for connection in self.active_connections:
await connection.send_text(message)
manager = ConnectionManager()
@app.get("/", response_class=HTMLResponse)
async def read_index():
with open("static/index.html") as f:
return HTMLResponse(content=f.read())
@app.get("/app", response_class=HTMLResponse)
async def read_app():
with open("static/app.html") as f:
return HTMLResponse(content=f.read())
@app.post("/register")
async def register_user(request: Request):
user_info = await request.json()
try:
insert_user_info(user_info)
return JSONResponse(content={"detail": "User registered successfully"}, status_code=200)
except Exception as e:
return JSONResponse(content={"detail": str(e)}, status_code=400)
@app.post("/upload_geotiff/")
async def upload_geotiff(file: UploadFile = File(...)):
# Initialize directories at the start
initialize_directories()
file_location = f"uploaded_geotiff.tif"
with open(file_location, "wb") as f:
f.write(file.file.read())
await manager.send_message("GeoTIFF file uploaded successfully. Slicing started.")
slices = await slice_geotiff(file_location, slice_size=3000)
await manager.send_message("Slicing complete. Starting weed detection.")
weed_bboxes = await detect_weeds_in_slices(slices)
await manager.send_message("Weed detection complete. Generating shapefile.")
# Create zip file
zip_file_path = await create_zip()
await manager.send_message("Shapefiles Generated. Zipping shapefile.")
# Upload the zip file to the bucket
response = upload_file_to_bucket(zip_file_path)
print(response)
await manager.send_message("Zip file uploaded to bucket storage.")
# Read zip file into buffer for download
zip_buffer = BytesIO()
with open(zip_file_path, 'rb') as f:
zip_buffer.write(f.read())
zip_buffer.seek(0)
# Cleanup files and directories
cleanup()
return StreamingResponse(zip_buffer, media_type="application/zip", headers={"Content-Disposition": "attachment; filename=weed_detections.zip"})
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
await manager.connect(websocket)
try:
while True:
data = await websocket.receive_text()
await manager.send_message(data)
except WebSocketDisconnect:
manager.disconnect(websocket)
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860) |