Coverage for src/main.py: 0%

1from fastapi import FastAPI, UploadFile, File 

2from fastapi.middleware.cors import CORSMiddleware 

3import numpy as np 

4import cv2 

5import os 

6 

7is_production = os.getenv("PY_ENV") == "production" 

8 

9app = FastAPI( 

10 docs_url=None if is_production else "/docs", 

11 redoc_url=None if is_production else "/redoc", 

12 openapi_url=None if is_production else "/openapi.json", 

13) 

14 

15 

16# Get allowed origins from environment variable 

17allowed_origins = os.environ.get("ALLOWED_ORIGINS", "http://localhost:3000").split(",") 

18 

19app.add_middleware( 

20 CORSMiddleware, 

21 allow_origins=allowed_origins, 

22 allow_credentials=True, 

23 allow_methods=["*"], 

24 allow_headers=["*"], 

25) 

26 

27 

28@app.post("/extract-palette") 

29async def extract_palette(image: UploadFile = File(...), colors: int = 5): 

30 """Extract a color palette from an uploaded image.""" 

31 contents = await image.read() 

32 nparr = np.frombuffer(contents, np.uint8) 

33 img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) 

34 

35 # Convert to RGB (from BGR) 

36 img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 

37 

38 # Reshape and get unique colors 

39 pixels = img.reshape(-1, 3) 

40 

41 # Use k-means to find dominant colors 

42 pixels = np.float32(pixels) 

43 criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 200, 0.1) 

44 _, labels, palette = cv2.kmeans(pixels, colors, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS) 

45 

46 # Convert to hex 

47 palette_hex = [] 

48 for color in palette: 

49 hex_color = "#{:02x}{:02x}{:02x}".format(int(color[0]), int(color[1]), int(color[2])) 

50 palette_hex.append(hex_color) 

51 

52 return {"palette": palette_hex} 

53 

54 

55@app.get("/health") 

56async def health_check(): 

57 return {"status": "healthy", "version": "1.0.0"} 

58 

59 

60@app.get("/test") 

61async def test_endpoint(): 

62 return {"message": "This is a test endpoint."}